Merge branch 'main' of github.com:stacklok/codegate into workspaces-to-open-api-spec

Author: alex-mcgovern

src/codegate/api/v1.py

@@ -1,8 +1,10 @@
 from fastapi import APIRouter, Response
 from fastapi.exceptions import HTTPException
 from fastapi.routing import APIRoute
+from pydantic import ValidationError
 
 from codegate.api import v1_models
+from codegate.db.connection import AlreadyExistsError
 from codegate.workspaces.crud import WorkspaceCrud
 
 v1 = APIRouter()
@@ -52,13 +54,19 @@ async def activate_workspace(request: v1_models.ActivateWorkspaceRequest, status
 async def create_workspace(request: v1_models.CreateWorkspaceRequest):
     """Create a new workspace."""
     # Input validation is done in the model
-    created = await wscrud.add_workspace(request.name)
-
-    # TODO: refactor to use a more specific exception
-    if not created:
-        raise HTTPException(status_code=400, detail="Failed to create workspace")
-
-    return v1_models.Workspace(name=request.name)
+    try:
+        created = await wscrud.add_workspace(request.name)
+    except AlreadyExistsError:
+        raise HTTPException(status_code=409, detail="Workspace already exists")
+    except ValidationError:
+        raise HTTPException(status_code=400,
+                            detail=("Invalid workspace name. "
+                                    "Please use only alphanumeric characters and dashes"))
+    except Exception:
+        raise HTTPException(status_code=500, detail="Internal server error")
+
+    if created:
+        return v1_models.Workspace(name=created.name)
 
 
 @v1.delete(

src/codegate/dashboard/dashboard.py

@@ -5,8 +5,8 @@
 import structlog
 from fastapi import APIRouter, Depends, FastAPI
 from fastapi.responses import StreamingResponse
-from codegate import __version__
 
+from codegate import __version__
 from codegate.dashboard.post_processing import (
     parse_get_alert_conversation,
     parse_messages_in_conversations,
@@ -81,7 +81,7 @@ def version_check():
         latest_version_stripped = latest_version.lstrip('v')
 
         is_latest: bool = latest_version_stripped == current_version
-        
+
         return {
             "current_version": current_version,
             "latest_version": latest_version_stripped,

src/codegate/db/connection.py

@@ -7,9 +7,9 @@
 import structlog
 from alembic import command as alembic_command
 from alembic.config import Config as AlembicConfig
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel
 from sqlalchemy import CursorResult, TextClause, text
-from sqlalchemy.exc import OperationalError
+from sqlalchemy.exc import IntegrityError, OperationalError
 from sqlalchemy.ext.asyncio import create_async_engine
 
 from codegate.db.fim_cache import FimCache
@@ -30,6 +30,8 @@
 alert_queue = asyncio.Queue()
 fim_cache = FimCache()
 
+class AlreadyExistsError(Exception):
+    pass
 
 class DbCodeGate:
     _instance = None
@@ -70,11 +72,11 @@ def __init__(self, sqlite_path: Optional[str] = None):
         super().__init__(sqlite_path)
 
     async def _execute_update_pydantic_model(
-        self, model: BaseModel, sql_command: TextClause
+        self, model: BaseModel, sql_command: TextClause, should_raise: bool = False
     ) -> Optional[BaseModel]:
         """Execute an update or insert command for a Pydantic model."""
-        async with self._async_db_engine.begin() as conn:
-            try:
+        try:
+            async with self._async_db_engine.begin() as conn:
                 result = await conn.execute(sql_command, model.model_dump())
                 row = result.first()
                 if row is None:
@@ -83,9 +85,11 @@ async def _execute_update_pydantic_model(
                 # Get the class of the Pydantic object to create a new object
                 model_class = model.__class__
                 return model_class(**row._asdict())
-            except Exception as e:
-                logger.error(f"Failed to update model: {model}.", error=str(e))
-                return None
+        except Exception as e:
+            logger.error(f"Failed to update model: {model}.", error=str(e))
+            if should_raise:
+                raise e
+            return None
 
     async def record_request(self, prompt_params: Optional[Prompt] = None) -> Optional[Prompt]:
         if prompt_params is None:
@@ -243,11 +247,14 @@ async def record_context(self, context: Optional[PipelineContext]) -> None:
             logger.error(f"Failed to record context: {context}.", error=str(e))
 
     async def add_workspace(self, workspace_name: str) -> Optional[Workspace]:
-        try:
-            workspace = Workspace(id=str(uuid.uuid4()), name=workspace_name)
-        except ValidationError as e:
-            logger.error(f"Failed to create workspace with name: {workspace_name}: {str(e)}")
-            return None
+        """Add a new workspace to the DB.
+
+        This handles validation and insertion of a new workspace.
+
+        It may raise a ValidationError if the workspace name is invalid.
+        or a AlreadyExistsError if the workspace already exists.
+        """
+        workspace = Workspace(id=str(uuid.uuid4()), name=workspace_name)
 
         sql = text(
             """
@@ -256,12 +263,13 @@ async def add_workspace(self, workspace_name: str) -> Optional[Workspace]:
             RETURNING *
             """
         )
-        try:
-            added_workspace = await self._execute_update_pydantic_model(workspace, sql)
-        except Exception as e:
-            logger.error(f"Failed to add workspace: {workspace_name}.", error=str(e))
-            return None
 
+        try:
+            added_workspace = await self._execute_update_pydantic_model(
+                workspace, sql, should_raise=True)
+        except IntegrityError as e:
+            logger.debug(f"Exception type: {type(e)}")
+            raise AlreadyExistsError(f"Workspace {workspace_name} already exists.")
         return added_workspace
 
     async def update_session(self, session: Session) -> Optional[Session]:

src/codegate/pipeline/cli/cli.py

@@ -32,10 +32,11 @@ async def codegate_cli(command):
     available_commands = {
         "version": Version().exec,
         "workspace": Workspace().exec,
-        "-h": lambda _: HELP_TEXT,
     }
     out_func = available_commands.get(command[0])
     if out_func is None:
+        if command[0] == "-h":
+            return HELP_TEXT
         return NOT_FOUND_TEXT
 
     return await out_func(command[1:])

src/codegate/pipeline/cli/commands.py

@@ -1,7 +1,10 @@
 from abc import ABC, abstractmethod
 from typing import List
 
+from pydantic import ValidationError
+
 from codegate import __version__
+from codegate.db.connection import AlreadyExistsError
 from codegate.workspaces.crud import WorkspaceCrud
 
 
@@ -63,37 +66,39 @@ async def _add_workspace(self, args: List[str]) -> str:
         Add a workspace
         """
         if args is None or len(args) == 0:
-            return "Please provide a name. Use `codegate-workspace add your_workspace_name`"
+            return "Please provide a name. Use `codegate workspace add your_workspace_name`"
 
         new_workspace_name = args[0]
         if not new_workspace_name:
-            return "Please provide a name. Use `codegate-workspace add your_workspace_name`"
+            return "Please provide a name. Use `codegate workspace add your_workspace_name`"
+
+        try:
+            _ = await self.workspace_crud.add_workspace(new_workspace_name)
+        except ValidationError:
+            return "Invalid workspace name: It should be alphanumeric and dashes"
+        except AlreadyExistsError:
+            return f"Workspace **{new_workspace_name}** already exists"
+        except Exception:
+            return "An error occurred while adding the workspace"
 
-        workspace_created = await self.workspace_crud.add_workspace(new_workspace_name)
-        if not workspace_created:
-            return (
-                "Something went wrong. Workspace could not be added.\n"
-                "1. Check if the name is alphanumeric and only contains dashes, and underscores.\n"
-                "2. Check if the workspace already exists."
-            )
         return f"Workspace **{new_workspace_name}** has been added"
 
     async def _activate_workspace(self, args: List[str]) -> str:
         """
         Activate a workspace
         """
         if args is None or len(args) == 0:
-            return "Please provide a name. Use `codegate-workspace activate workspace_name`"
+            return "Please provide a name. Use `codegate workspace activate workspace_name`"
 
         workspace_name = args[0]
         if not workspace_name:
-            return "Please provide a name. Use `codegate-workspace activate workspace_name`"
+            return "Please provide a name. Use `codegate workspace activate workspace_name`"
 
         was_activated = await self.workspace_crud.activate_workspace(workspace_name)
         if not was_activated:
             return (
                 f"Workspace **{workspace_name}** does not exist or was already active. "
-                f"Use `codegate-workspace add {workspace_name}` to add it"
+                f"Use `codegate workspace add {workspace_name}` to add it"
             )
         return f"Workspace **{workspace_name}** has been activated"
 

src/codegate/providers/ollama/adapter.py

@@ -25,10 +25,8 @@ def normalize(self, data: Dict) -> ChatCompletionRequest:
                 {"content": normalized_data.pop("prompt"), "role": "user"}
             ]
 
-        # In Ollama force the stream to be True. Continue is not setting this parameter and
-        # most of our functionality is for streaming completions.
-        normalized_data["stream"] = True
-
+        # if we have the stream flag in data we set it, otherwise defaults to true
+        normalized_data["stream"] = data.get("stream", True)
         return ChatCompletionRequest(**normalized_data)
 
     def denormalize(self, data: ChatCompletionRequest) -> Dict:

src/codegate/providers/ollama/completion_handler.py

@@ -60,7 +60,7 @@ def _create_streaming_response(self, stream: AsyncIterator[ChatResponse]) -> Str
         """
         return StreamingResponse(
             ollama_stream_generator(stream),
-            media_type="application/x-ndjson",
+            media_type="application/x-ndjson; charset=utf-8",
             headers={
                 "Cache-Control": "no-cache",
                 "Connection": "keep-alive",
@@ -70,4 +70,4 @@ def _create_streaming_response(self, stream: AsyncIterator[ChatResponse]) -> Str
     def _create_json_response(
         self, response: Union[GenerateResponse, ChatResponse]
     ) -> JSONResponse:
-        return JSONResponse(content=response.model_dump_json(), status_code=200)
+        return JSONResponse(status_code=200, content=response.model_dump())

src/codegate/providers/ollama/provider.py

@@ -65,7 +65,7 @@ async def show_model(request: Request):
                 response = await client.post(
                     f"{self.base_url}/api/show",
                     content=body,
-                    headers={"Content-Type": "application/json"},
+                    headers={"Content-Type": "application/json; charset=utf-8"},
                 )
                 return response.json()
 

src/codegate/providers/registry.py

@@ -16,7 +16,7 @@ def add_provider(self, name: str, provider: BaseProvider):
         to the FastAPI app.
         """
         self.providers[name] = provider
-        self.app.include_router(provider.get_routes())
+        self.app.include_router(provider.get_routes(), include_in_schema=False)
 
     def get_provider(self, name: str) -> Optional[BaseProvider]:
         """

src/codegate/workspaces/crud.py

@@ -1,16 +1,19 @@
 import datetime
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 from codegate.db.connection import DbReader, DbRecorder
-from codegate.db.models import Session, Workspace
+from codegate.db.models import ActiveWorkspace, Session, Workspace, WorkspaceActive
 
 
+class WorkspaceCrudError(Exception):
+    pass
+
 class WorkspaceCrud:
 
     def __init__(self):
         self._db_reader = DbReader()
 
-    async def add_workspace(self, new_workspace_name: str) -> bool:
+    async def add_workspace(self, new_workspace_name: str) -> Workspace:
         """
         Add a workspace
 
@@ -19,14 +22,20 @@ async def add_workspace(self, new_workspace_name: str) -> bool:
         """
         db_recorder = DbRecorder()
         workspace_created = await db_recorder.add_workspace(new_workspace_name)
-        return bool(workspace_created)
+        return workspace_created
 
-    async def get_workspaces(self):
+    async def get_workspaces(self)-> List[WorkspaceActive]:
         """
         Get all workspaces
         """
         return await self._db_reader.get_workspaces()
 
+    async def get_active_workspace(self) -> Optional[ActiveWorkspace]:
+        """
+        Get the active workspace
+        """
+        return await self._db_reader.get_active_workspace()
+
     async def _is_workspace_active_or_not_exist(
         self, workspace_name: str
     ) -> Tuple[bool, Optional[Session], Optional[Workspace]]: