Merge branch 'main' into run-on-main

rdimitrov · web-flow · commit b3a2ed43213b · 2025-01-21T12:30:30.000+02:00
diff --git a/Makefile b/Makefile
@@ -34,7 +34,7 @@ build: clean test
 image-build:
 	DOCKER_BUILDKIT=1 $(CONTAINER_BUILD) \
 		-f Dockerfile \
-		--build-arg LATEST_RELEASE=$(curl -s "https://api.github.com/repos/stacklok/codegate-ui/releases/latest" | grep '"zipball_url":' | cut -d '"' -f 4) \
+		--build-arg LATEST_RELEASE=$(shell curl -s "https://api.github.com/repos/stacklok/codegate-ui/releases/latest" | grep '"zipball_url":' | cut -d '"' -f 4) \
 		--build-arg CODEGATE_VERSION="$(shell git describe --tags --abbrev=0)-$(shell git rev-parse --short HEAD)-dev" \
 		-t codegate \
 		. \
diff --git a/api/openapi.json b/api/openapi.json
@@ -305,6 +305,109 @@
           }
         }
       }
+    },
+    "/api/v1/workspaces/{workspace_name}/alerts": {
+      "get": {
+        "tags": [
+          "CodeGate API",
+          "Workspaces"
+        ],
+        "summary": "Get Workspace Alerts",
+        "description": "Get alerts for a workspace.",
+        "operationId": "v1_get_workspace_alerts",
+        "parameters": [
+          {
+            "name": "workspace_name",
+            "in": "path",
+            "required": true,
+            "schema": {
+              "type": "string",
+              "title": "Workspace Name"
+            }
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "Successful Response",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "type": "array",
+                  "items": {
+                    "anyOf": [
+                      {
+                        "$ref": "#/components/schemas/AlertConversation"
+                      },
+                      {
+                        "type": "null"
+                      }
+                    ]
+                  },
+                  "title": "Response V1 Get Workspace Alerts"
+                }
+              }
+            }
+          },
+          "422": {
+            "description": "Validation Error",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/HTTPValidationError"
+                }
+              }
+            }
+          }
+        }
+      }
+    },
+    "/api/v1/workspaces/{workspace_name}/messages": {
+      "get": {
+        "tags": [
+          "CodeGate API",
+          "Workspaces"
+        ],
+        "summary": "Get Workspace Messages",
+        "description": "Get messages for a workspace.",
+        "operationId": "v1_get_workspace_messages",
+        "parameters": [
+          {
+            "name": "workspace_name",
+            "in": "path",
+            "required": true,
+            "schema": {
+              "type": "string",
+              "title": "Workspace Name"
+            }
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "Successful Response",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "type": "array",
+                  "items": {
+                    "$ref": "#/components/schemas/Conversation"
+                  },
+                  "title": "Response V1 Get Workspace Messages"
+                }
+              }
+            }
+          },
+          "422": {
+            "description": "Validation Error",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/HTTPValidationError"
+                }
+              }
+            }
+          }
+        }
+      }
     }
   },
   "components": {
diff --git a/migrations/versions/2025_01_21_0820-4dec3e456c9e_add_on_delete_cascade.py b/migrations/versions/2025_01_21_0820-4dec3e456c9e_add_on_delete_cascade.py
@@ -0,0 +1,114 @@
+"""add_on_delete_cascade
+
+Revision ID: 4dec3e456c9e
+Revises: e6227073183d
+Create Date: 2025-01-21 08:20:12.221051+00:00
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "4dec3e456c9e"
+down_revision: Union[str, None] = "e6227073183d"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # To add ON DELETE CASCADE to the foreign key constraint, we need to
+    # rename the table, create a new table with the constraint, and copy
+    # the data over.
+    op.execute("ALTER TABLE prompts RENAME TO _prompts_old;")
+    op.execute(
+        """
+        CREATE TABLE prompts (
+            id TEXT PRIMARY KEY,  -- UUID stored as TEXT
+            timestamp DATETIME NOT NULL,
+            provider TEXT,       -- VARCHAR(255)
+            request TEXT NOT NULL,  -- Record the full request that arrived to the server
+            type TEXT NOT NULL, -- VARCHAR(50) (e.g. "fim", "chat")
+            workspace_id TEXT NOT NULL,
+            FOREIGN KEY (workspace_id) REFERENCES workspaces(id) ON DELETE CASCADE
+        );
+        """
+    )
+    op.execute("INSERT INTO prompts SELECT * FROM _prompts_old;")
+    op.execute("DROP TABLE _prompts_old;")
+
+    # Doing the same for the sessions table
+    op.execute("ALTER TABLE sessions RENAME TO _sessions_old;")
+    op.execute(
+        """
+        CREATE TABLE sessions (
+            id TEXT PRIMARY KEY,  -- UUID stored as TEXT
+            active_workspace_id TEXT NOT NULL,
+            last_update DATETIME NOT NULL,
+            FOREIGN KEY (active_workspace_id) REFERENCES workspaces(id) ON DELETE CASCADE
+        );
+        """
+    )
+    op.execute("INSERT INTO sessions SELECT * FROM _sessions_old;")
+    op.execute("DROP TABLE _sessions_old;")
+
+    # Doing the same for the output table
+    op.execute("ALTER TABLE outputs RENAME TO _outputs_old;")
+    op.execute(
+        """
+        CREATE TABLE outputs (
+            id TEXT PRIMARY KEY,  -- UUID stored as TEXT
+            prompt_id TEXT NOT NULL,
+            timestamp DATETIME NOT NULL,
+            output TEXT NOT NULL,   -- Record the full response. If stream will be a list of objects
+            FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE
+        );
+        """
+    )
+    op.execute("INSERT INTO outputs SELECT * FROM _outputs_old;")
+    op.execute("DROP TABLE _outputs_old;")
+
+    # Doing the same for the alerts table
+    op.execute("ALTER TABLE alerts RENAME TO _alerts_old;")
+    op.execute(
+        """
+        CREATE TABLE alerts (
+            id TEXT PRIMARY KEY,  -- UUID stored as TEXT
+            prompt_id TEXT NOT NULL,
+            code_snippet TEXT,
+            trigger_string TEXT, -- VARCHAR(255)
+            trigger_type TEXT NOT NULL,   -- VARCHAR(50)
+            trigger_category TEXT,
+            timestamp DATETIME NOT NULL,
+            FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE
+        );
+        """
+    )
+    op.execute("INSERT INTO alerts SELECT * FROM _alerts_old;")
+    op.execute("DROP TABLE _alerts_old;")
+
+    # Dropping unused table
+    op.execute("DROP TABLE settings;")
+
+    # Create indexes for foreign keys
+    op.execute("CREATE INDEX idx_outputs_prompt_id ON outputs(prompt_id);")
+    op.execute("CREATE INDEX idx_alerts_prompt_id ON alerts(prompt_id);")
+    op.execute("CREATE INDEX idx_prompts_workspace_id ON prompts (workspace_id);")
+    op.execute("CREATE INDEX idx_sessions_workspace_id ON sessions (active_workspace_id);")
+
+
+def downgrade() -> None:
+    # Settings table
+    op.execute(
+        """
+        CREATE TABLE settings (
+            id TEXT PRIMARY KEY,  -- UUID stored as TEXT
+            ip TEXT,             -- VARCHAR(45)
+            port INTEGER,
+            llm_model TEXT,      -- VARCHAR(255)
+            system_prompt TEXT,
+            other_settings TEXT  -- JSON stored as TEXT
+        );
+        """
+    )
diff --git a/src/codegate/api/dashboard/dashboard.py b/src/codegate/api/dashboard/dashboard.py
@@ -3,7 +3,7 @@
 
 import requests
 import structlog
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, HTTPException
 from fastapi.responses import StreamingResponse
 from fastapi.routing import APIRoute
 
@@ -14,12 +14,15 @@
 )
 from codegate.api.dashboard.request_models import AlertConversation, Conversation
 from codegate.db.connection import DbReader, alert_queue
+from codegate.workspaces import crud
 
 logger = structlog.get_logger("codegate")
 
 dashboard_router = APIRouter()
 db_reader = None
 
+wscrud = crud.WorkspaceCrud()
+
 
 def uniq_name(route: APIRoute):
     return f"v1_{route.name}"
@@ -48,9 +51,14 @@ def get_messages(db_reader: DbReader = Depends(get_db_reader)) -> List[Conversat
     """
     Get all the messages from the database and return them as a list of conversations.
     """
-    prompts_outputs = asyncio.run(db_reader.get_prompts_with_output())
+    try:
+        active_ws = asyncio.run(wscrud.get_active_workspace())
+        prompts_outputs = asyncio.run(db_reader.get_prompts_with_output(active_ws.id))
 
-    return asyncio.run(parse_messages_in_conversations(prompts_outputs))
+        return asyncio.run(parse_messages_in_conversations(prompts_outputs))
+    except Exception as e:
+        logger.error(f"Error getting messages: {str(e)}")
+        raise HTTPException(status_code=500, detail="Internal server error")
 
 
 @dashboard_router.get(
@@ -60,8 +68,15 @@ def get_alerts(db_reader: DbReader = Depends(get_db_reader)) -> List[Optional[Al
     """
     Get all the messages from the database and return them as a list of conversations.
     """
-    alerts_prompt_output = asyncio.run(db_reader.get_alerts_with_prompt_and_output())
-    return asyncio.run(parse_get_alert_conversation(alerts_prompt_output))
+    try:
+        active_ws = asyncio.run(wscrud.get_active_workspace())
+        alerts_prompt_output = asyncio.run(
+            db_reader.get_alerts_with_prompt_and_output(active_ws.id)
+        )
+        return asyncio.run(parse_get_alert_conversation(alerts_prompt_output))
+    except Exception as e:
+        logger.error(f"Error getting alerts: {str(e)}")
+        raise HTTPException(status_code=500, detail="Internal server error")
 
 
 async def generate_sse_events() -> AsyncGenerator[str, None]:
diff --git a/src/codegate/api/v1.py b/src/codegate/api/v1.py
@@ -1,16 +1,22 @@
+from typing import List, Optional
+
 from fastapi import APIRouter, HTTPException, Response
 from fastapi.routing import APIRoute
 from pydantic import ValidationError
 
 from codegate.api import v1_models
-from codegate.api.dashboard.dashboard import dashboard_router
-from codegate.db.connection import AlreadyExistsError
+from codegate.api.dashboard import dashboard
+from codegate.api.dashboard.request_models import AlertConversation, Conversation
+from codegate.db.connection import AlreadyExistsError, DbReader
 from codegate.workspaces import crud
 
 v1 = APIRouter()
-v1.include_router(dashboard_router)
+v1.include_router(dashboard.dashboard_router)
 wscrud = crud.WorkspaceCrud()
 
+# This is a singleton object
+dbreader = DbReader()
+
 
 def uniq_name(route: APIRoute):
     return f"v1_{route.name}"
@@ -94,3 +100,45 @@ async def delete_workspace(workspace_name: str):
         raise HTTPException(status_code=500, detail="Internal server error")
 
     return Response(status_code=204)
+
+
+@v1.get(
+    "/workspaces/{workspace_name}/alerts",
+    tags=["Workspaces"],
+    generate_unique_id_function=uniq_name,
+)
+async def get_workspace_alerts(workspace_name: str) -> List[Optional[AlertConversation]]:
+    """Get alerts for a workspace."""
+    try:
+        ws = await wscrud.get_workspace_by_name(workspace_name)
+    except crud.WorkspaceDoesNotExistError:
+        raise HTTPException(status_code=404, detail="Workspace does not exist")
+    except Exception:
+        raise HTTPException(status_code=500, detail="Internal server error")
+
+    try:
+        alerts = await dbreader.get_alerts_with_prompt_and_output(ws.id)
+        return await dashboard.parse_get_alert_conversation(alerts)
+    except Exception:
+        raise HTTPException(status_code=500, detail="Internal server error")
+
+
+@v1.get(
+    "/workspaces/{workspace_name}/messages",
+    tags=["Workspaces"],
+    generate_unique_id_function=uniq_name,
+)
+async def get_workspace_messages(workspace_name: str) -> List[Conversation]:
+    """Get messages for a workspace."""
+    try:
+        ws = await wscrud.get_workspace_by_name(workspace_name)
+    except crud.WorkspaceDoesNotExistError:
+        raise HTTPException(status_code=404, detail="Workspace does not exist")
+    except Exception:
+        raise HTTPException(status_code=500, detail="Internal server error")
+
+    try:
+        prompts_outputs = await dbreader.get_prompts_with_output(ws.id)
+        return await dashboard.parse_messages_in_conversations(prompts_outputs)
+    except Exception:
+        raise HTTPException(status_code=500, detail="Internal server error")
diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py
