diff --git a/src/codegate/api/dashboard/dashboard.py b/src/codegate/api/dashboard/dashboard.py index f59b31a5..fa558a91 100644 --- 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 index 1c112a78..5a31066e 100644 --- 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}" @@ -92,3 +98,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 index 6e014916..103a8eec 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -365,7 +365,7 @@ async def _exec_select_conditions_to_pydantic( raise e return None - async def get_prompts_with_output(self) -> List[GetPromptWithOutputsRow]: + async def get_prompts_with_output(self, workpace_id: str) -> List[GetPromptWithOutputsRow]: sql = text( """ SELECT @@ -375,13 +375,19 @@ async def get_prompts_with_output(self) -> List[GetPromptWithOutputsRow]: o.timestamp as output_timestamp FROM prompts p LEFT JOIN outputs o ON p.id = o.prompt_id + WHERE p.workspace_id = :workspace_id ORDER BY o.timestamp DESC """ ) - prompts = await self._execute_select_pydantic_model(GetPromptWithOutputsRow, sql) + conditions = {"workspace_id": workpace_id} + prompts = await self._exec_select_conditions_to_pydantic( + GetPromptWithOutputsRow, sql, conditions, should_raise=True + ) return prompts - async def get_alerts_with_prompt_and_output(self) -> List[GetAlertsWithPromptAndOutputRow]: + async def get_alerts_with_prompt_and_output( + self, workspace_id: str + ) -> List[GetAlertsWithPromptAndOutputRow]: sql = text( """ SELECT @@ -402,10 +408,14 @@ async def get_alerts_with_prompt_and_output(self) -> List[GetAlertsWithPromptAnd FROM alerts a LEFT JOIN prompts p ON p.id = a.prompt_id LEFT JOIN outputs o ON p.id = o.prompt_id + WHERE p.workspace_id = :workspace_id ORDER BY a.timestamp DESC """ ) - prompts = await self._execute_select_pydantic_model(GetAlertsWithPromptAndOutputRow, sql) + conditions = {"workspace_id": workspace_id} + prompts = await self._exec_select_conditions_to_pydantic( + GetAlertsWithPromptAndOutputRow, sql, conditions, should_raise=True + ) return prompts async def get_workspaces(self) -> List[WorkspaceActive]: