Skip to content

Commit

Permalink
Merge branch 'main' into cline
Browse files Browse the repository at this point in the history
  • Loading branch information
yrobla authored Jan 21, 2025
2 parents 9cff3fb + e49c72f commit 2ffd1dc
Show file tree
Hide file tree
Showing 12 changed files with 167 additions and 139 deletions.
9 changes: 7 additions & 2 deletions api/openapi.json
Original file line number Diff line number Diff line change
Expand Up @@ -285,8 +285,13 @@
}
],
"responses": {
"204": {
"description": "Successful Response"
"200": {
"description": "Successful Response",
"content": {
"application/json": {
"schema": {}
}
}
},
"422": {
"description": "Validation Error",
Expand Down
30 changes: 30 additions & 0 deletions migrations/versions/8e4b4b8d1a88_add_soft_delete.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
"""add soft delete
Revision ID: 8e4b4b8d1a88
Revises: 5c2f3eee5f90
Create Date: 2025-01-20 14:08:40.851647
"""

from typing import Sequence, Union

from alembic import op

# revision identifiers, used by Alembic.
revision: str = "8e4b4b8d1a88"
down_revision: Union[str, None] = "5c2f3eee5f90"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
op.execute(
"""
ALTER TABLE workspaces
ADD COLUMN deleted_at DATETIME DEFAULT NULL;
"""
)


def downgrade() -> None:
op.execute("ALTER TABLE workspaces DROP COLUMN deleted_at;")
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
"""merging system prompt and soft-deletes
Revision ID: e6227073183d
Revises: 8e4b4b8d1a88, a692c8b52308
Create Date: 2025-01-20 16:08:40.645298
"""

from typing import Sequence, Union

# revision identifiers, used by Alembic.
revision: str = "e6227073183d"
down_revision: Union[str, None] = ("8e4b4b8d1a88", "a692c8b52308")
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
pass


def downgrade() -> None:
pass
146 changes: 25 additions & 121 deletions poetry.lock

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,9 @@ tree-sitter-java = "==0.23.5"
tree-sitter-javascript = "==0.23.1"
tree-sitter-python = "==0.23.6"
tree-sitter-rust = "==0.23.2"
sqlite-vec-sl-tmp = "==0.0.4"
alembic = "==1.14.1"
pygments = "==2.19.1"
sqlite-vec = "==0.1.6"

[tool.poetry.group.dev.dependencies]
pytest = "==8.3.4"
Expand Down
4 changes: 2 additions & 2 deletions scripts/import_packages.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import sqlite3

import numpy as np
import sqlite_vec_sl_tmp
import sqlite_vec

from codegate.inference.inference_engine import LlamaCppInferenceEngine
from codegate.utils.utils import generate_vector_string
Expand All @@ -27,7 +27,7 @@ def __init__(self, jsonl_dir="data", vec_db_path="./sqlite_data/vectordb.db"):
def _get_connection(self):
conn = sqlite3.connect(self.vec_db_path)
conn.enable_load_extension(True)
sqlite_vec_sl_tmp.load(conn)
sqlite_vec.load(conn)
conn.enable_load_extension(False)
return conn

Expand Down
18 changes: 13 additions & 5 deletions src/codegate/api/v1.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,11 +45,11 @@ async def activate_workspace(request: v1_models.ActivateWorkspaceRequest, status
try:
await wscrud.activate_workspace(request.name)
except crud.WorkspaceAlreadyActiveError:
return HTTPException(status_code=409, detail="Workspace already active")
raise HTTPException(status_code=409, detail="Workspace already active")
except crud.WorkspaceDoesNotExistError:
return HTTPException(status_code=404, detail="Workspace does not exist")
raise HTTPException(status_code=404, detail="Workspace does not exist")
except Exception:
return HTTPException(status_code=500, detail="Internal server error")
raise HTTPException(status_code=500, detail="Internal server error")

return Response(status_code=204)

Expand Down Expand Up @@ -79,8 +79,16 @@ async def create_workspace(request: v1_models.CreateWorkspaceRequest) -> v1_mode
"/workspaces/{workspace_name}",
tags=["Workspaces"],
generate_unique_id_function=uniq_name,
status_code=204,
)
async def delete_workspace(workspace_name: str):
"""Delete a workspace by name."""
raise NotImplementedError
try:
_ = await wscrud.soft_delete_workspace(workspace_name)
except crud.WorkspaceDoesNotExistError:
raise HTTPException(status_code=404, detail="Workspace does not exist")
except crud.WorkspaceCrudError as e:
raise HTTPException(status_code=400, detail=str(e))
except Exception:
raise HTTPException(status_code=500, detail="Internal server error")

return Response(status_code=204)
17 changes: 16 additions & 1 deletion src/codegate/db/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -304,6 +304,20 @@ async def update_session(self, session: Session) -> Optional[Session]:
active_session = await self._execute_update_pydantic_model(session, sql, should_raise=True)
return active_session

async def soft_delete_workspace(self, workspace: Workspace) -> Optional[Workspace]:
sql = text(
"""
UPDATE workspaces
SET deleted_at = CURRENT_TIMESTAMP
WHERE id = :id
RETURNING *
"""
)
deleted_workspace = await self._execute_update_pydantic_model(
workspace, sql, should_raise=True
)
return deleted_workspace


class DbReader(DbCodeGate):

Expand Down Expand Up @@ -401,6 +415,7 @@ async def get_workspaces(self) -> List[WorkspaceActive]:
w.id, w.name, s.active_workspace_id
FROM workspaces w
LEFT JOIN sessions s ON w.id = s.active_workspace_id
WHERE w.deleted_at IS NULL
"""
)
workspaces = await self._execute_select_pydantic_model(WorkspaceActive, sql)
Expand All @@ -412,7 +427,7 @@ async def get_workspace_by_name(self, name: str) -> Optional[Workspace]:
SELECT
id, name, system_prompt
FROM workspaces
WHERE name = :name
WHERE name = :name AND deleted_at IS NULL
"""
)
conditions = {"name": name}
Expand Down
22 changes: 22 additions & 0 deletions src/codegate/pipeline/cli/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,7 @@ def subcommands(self) -> Dict[str, Callable[[List[str]], Awaitable[str]]]:
"list": self._list_workspaces,
"add": self._add_workspace,
"activate": self._activate_workspace,
"remove": self._remove_workspace,
}

async def _list_workspaces(self, flags: Dict[str, str], args: List[str]) -> str:
Expand Down Expand Up @@ -211,6 +212,27 @@ async def _activate_workspace(self, flags: Dict[str, str], args: List[str]) -> s
return "An error occurred while activating the workspace"
return f"Workspace **{workspace_name}** has been activated"

async def _remove_workspace(self, flags: Dict[str, str], args: List[str]) -> str:
"""
Remove a workspace
"""
if args is None or len(args) == 0:
return "Please provide a name. Use `codegate workspace remove workspace_name`"

workspace_name = args[0]
if not workspace_name:
return "Please provide a name. Use `codegate workspace remove workspace_name`"

try:
await self.workspace_crud.soft_delete_workspace(workspace_name)
except crud.WorkspaceDoesNotExistError:
return f"Workspace **{workspace_name}** does not exist"
except crud.WorkspaceCrudError as e:
return str(e)
except Exception:
return "An error occurred while removing the workspace"
return f"Workspace **{workspace_name}** has been removed"

@property
def help(self) -> str:
return (
Expand Down
4 changes: 2 additions & 2 deletions src/codegate/storage/storage_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from typing import List, Optional

import numpy as np
import sqlite_vec_sl_tmp
import sqlite_vec
import structlog

from codegate.config import Config
Expand Down Expand Up @@ -71,7 +71,7 @@ def _get_connection(self):
try:
conn = sqlite3.connect(self.db_path)
conn.enable_load_extension(True)
sqlite_vec_sl_tmp.load(conn)
sqlite_vec.load(conn)
conn.enable_load_extension(False)
return conn
except Exception as e:
Expand Down
29 changes: 25 additions & 4 deletions src/codegate/workspaces/crud.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,10 +68,6 @@ async def _is_workspace_active(
async def activate_workspace(self, workspace_name: str):
"""
Activate a workspace
Will return:
- True if the workspace was activated
- False if the workspace is already active or does not exist
"""
is_active, session, workspace = await self._is_workspace_active(workspace_name)
if is_active:
Expand Down Expand Up @@ -100,6 +96,31 @@ async def update_workspace_system_prompt(
updated_workspace = await db_recorder.update_workspace(workspace_update)
return updated_workspace

async def soft_delete_workspace(self, workspace_name: str):
"""
Soft delete a workspace
"""
if workspace_name == "":
raise WorkspaceCrudError("Workspace name cannot be empty.")
if workspace_name == "default":
raise WorkspaceCrudError("Cannot delete default workspace.")

selected_workspace = await self._db_reader.get_workspace_by_name(workspace_name)
if not selected_workspace:
raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.")

# Check if workspace is active, if it is, make the default workspace active
active_workspace = await self._db_reader.get_active_workspace()
if active_workspace and active_workspace.id == selected_workspace.id:
raise WorkspaceCrudError("Cannot delete active workspace.")

db_recorder = DbRecorder()
try:
_ = await db_recorder.soft_delete_workspace(selected_workspace)
except Exception:
raise WorkspaceCrudError(f"Error deleting workspace {workspace_name}")
return

async def get_workspace_by_name(self, workspace_name: str) -> Workspace:
workspace = await self._db_reader.get_workspace_by_name(workspace_name)
if not workspace:
Expand Down
2 changes: 1 addition & 1 deletion tests/vectordb/test_sqlitevec.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@

@pytest.fixture(scope="module")
def mock_sqlite_vec():
with patch("sqlite_vec_sl_tmp.load") as mock_load:
with patch("sqlite_vec.load") as mock_load:
# Mock the vector similarity extension loading
def setup_vector_similarity(conn):
cursor = conn.cursor()
Expand Down

0 comments on commit 2ffd1dc

Please sign in to comment.