diff --git a/src/codegate/api/v1.py b/src/codegate/api/v1.py index ebd9be79..f9d4b00b 100644 --- a/src/codegate/api/v1.py +++ b/src/codegate/api/v1.py @@ -397,6 +397,33 @@ async def get_workspace_alerts(workspace_name: str) -> List[Optional[v1_models.A raise HTTPException(status_code=500, detail="Internal server error") +@v1.get( + "/workspaces/{workspace_name}/alerts-summary", + tags=["Workspaces"], + generate_unique_id_function=uniq_name, +) +async def get_workspace_alerts_summary(workspace_name: str) -> v1_models.AlertSummary: + """Get alert summary 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: + logger.exception("Error while getting workspace") + raise HTTPException(status_code=500, detail="Internal server error") + + try: + summary = await dbreader.get_alerts_summary_by_workspace(ws.id) + return v1_models.AlertSummary( + malicious_packages=summary["codegate_context_retriever_count"], + pii=summary["codegate_pii_count"], + secrets=summary["codegate_secrets_count"], + ) + except Exception: + logger.exception("Error while getting alerts summary") + raise HTTPException(status_code=500, detail="Internal server error") + + @v1.get( "/workspaces/{workspace_name}/messages", tags=["Workspaces"], diff --git a/src/codegate/api/v1_models.py b/src/codegate/api/v1_models.py index c608484c..51f65ea9 100644 --- a/src/codegate/api/v1_models.py +++ b/src/codegate/api/v1_models.py @@ -190,6 +190,16 @@ def from_db_model(db_model: db_models.Alert) -> "Alert": timestamp: datetime.datetime +class AlertSummary(pydantic.BaseModel): + """ + Represents a set of summary alerts + """ + + malicious_packages: int + pii: int + secrets: int + + class PartialQuestionAnswer(pydantic.BaseModel): """ Represents a partial conversation. diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index 803943b3..38bf6010 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -746,6 +746,38 @@ async def get_alerts_by_workspace( ) return prompts + async def get_alerts_summary_by_workspace(self, workspace_id: str) -> dict: + """Get aggregated alert summary counts for a given workspace_id.""" + sql = text( + """ + SELECT + COUNT(*) AS total_alerts, + SUM(CASE WHEN a.trigger_type = 'codegate-secrets' THEN 1 ELSE 0 END) + AS codegate_secrets_count, + SUM(CASE WHEN a.trigger_type = 'codegate-context-retriever' THEN 1 ELSE 0 END) + AS codegate_context_retriever_count, + SUM(CASE WHEN a.trigger_type = 'codegate-pii' THEN 1 ELSE 0 END) + AS codegate_pii_count + FROM alerts a + INNER JOIN prompts p ON p.id = a.prompt_id + WHERE p.workspace_id = :workspace_id + """ + ) + conditions = {"workspace_id": workspace_id} + + async with self._async_db_engine.begin() as conn: + result = await conn.execute(sql, conditions) + row = result.fetchone() + + # Return a dictionary with counts (handling None values safely) + return { + "codegate_secrets_count": row.codegate_secrets_count or 0 if row else 0, + "codegate_context_retriever_count": ( + row.codegate_context_retriever_count or 0 if row else 0 + ), + "codegate_pii_count": row.codegate_pii_count or 0 if row else 0, + } + async def get_workspaces(self) -> List[WorkspaceWithSessionInfo]: sql = text( """