From df79c2e2011317ebab37beee36eead922d4a8a8d Mon Sep 17 00:00:00 2001 From: alex-mcgovern Date: Fri, 17 Jan 2025 11:05:38 +0000 Subject: [PATCH 1/5] Add workspaces to OpenAPI spec --- api/openapi.json | 316 ++++++++++++++++++++++++++++ pyproject.toml | 2 +- src/codegate/dashboard/dashboard.py | 16 -- src/codegate/db/connection.py | 2 +- src/codegate/server.py | 16 ++ 5 files changed, 334 insertions(+), 18 deletions(-) diff --git a/api/openapi.json b/api/openapi.json index 6b495bb3..db9678aa 100644 --- a/api/openapi.json +++ b/api/openapi.json @@ -83,10 +83,217 @@ } } } + }, + "/api/v1/workspaces": { + "get": { + "tags": [ + "CodeGate API", + "Workspaces" + ], + "summary": "List Workspaces", + "description": "List all workspaces.", + "operationId": "v1_list_workspaces", + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ListWorkspacesResponse" + } + } + } + } + } + }, + "post": { + "tags": [ + "CodeGate API", + "Workspaces" + ], + "summary": "Create Workspace", + "description": "Create a new workspace.", + "operationId": "v1_create_workspace", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CreateWorkspaceRequest" + } + } + }, + "required": true + }, + "responses": { + "201": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": {} + } + } + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + } + }, + "/api/v1/workspaces/active": { + "get": { + "tags": [ + "CodeGate API", + "Workspaces" + ], + "summary": "List Active Workspaces", + "description": "List all active workspaces.\n\nIn it's current form, this function will only return one workspace. That is,\nthe globally active workspace.", + "operationId": "v1_list_active_workspaces", + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ListActiveWorkspacesResponse" + } + } + } + } + } + }, + "post": { + "tags": [ + "CodeGate API", + "Workspaces" + ], + "summary": "Activate Workspace", + "description": "Activate a workspace by name.", + "operationId": "v1_activate_workspace", + "parameters": [ + { + "name": "status_code", + "in": "query", + "required": false, + "schema": { + "default": 204, + "title": "Status Code" + } + } + ], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ActivateWorkspaceRequest" + } + } + } + }, + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": {} + } + } + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + } + }, + "/api/v1/workspaces/{workspace_name}": { + "delete": { + "tags": [ + "CodeGate API", + "Workspaces" + ], + "summary": "Delete Workspace", + "description": "Delete a workspace by name.", + "operationId": "v1_delete_workspace", + "parameters": [ + { + "name": "workspace_name", + "in": "path", + "required": true, + "schema": { + "type": "string", + "title": "Workspace Name" + } + } + ], + "responses": { + "204": { + "description": "Successful Response" + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + } } }, "components": { "schemas": { + "ActivateWorkspaceRequest": { + "properties": { + "name": { + "type": "string", + "title": "Name" + } + }, + "type": "object", + "required": [ + "name" + ], + "title": "ActivateWorkspaceRequest" + }, + "ActiveWorkspace": { + "properties": { + "name": { + "type": "string", + "title": "Name" + }, + "is_active": { + "type": "boolean", + "title": "Is Active" + }, + "last_updated": { + "title": "Last Updated" + } + }, + "type": "object", + "required": [ + "name", + "is_active", + "last_updated" + ], + "title": "ActiveWorkspace" + }, "AlertConversation": { "properties": { "conversation": { @@ -268,6 +475,64 @@ "title": "Conversation", "description": "Represents a conversation." }, + "CreateWorkspaceRequest": { + "properties": { + "name": { + "type": "string", + "title": "Name" + } + }, + "type": "object", + "required": [ + "name" + ], + "title": "CreateWorkspaceRequest" + }, + "HTTPValidationError": { + "properties": { + "detail": { + "items": { + "$ref": "#/components/schemas/ValidationError" + }, + "type": "array", + "title": "Detail" + } + }, + "type": "object", + "title": "HTTPValidationError" + }, + "ListActiveWorkspacesResponse": { + "properties": { + "workspaces": { + "items": { + "$ref": "#/components/schemas/ActiveWorkspace" + }, + "type": "array", + "title": "Workspaces" + } + }, + "type": "object", + "required": [ + "workspaces" + ], + "title": "ListActiveWorkspacesResponse" + }, + "ListWorkspacesResponse": { + "properties": { + "workspaces": { + "items": { + "$ref": "#/components/schemas/Workspace" + }, + "type": "array", + "title": "Workspaces" + } + }, + "type": "object", + "required": [ + "workspaces" + ], + "title": "ListWorkspacesResponse" + }, "QuestionAnswer": { "properties": { "question": { @@ -291,6 +556,57 @@ ], "title": "QuestionAnswer", "description": "Represents a question and answer pair." + }, + "ValidationError": { + "properties": { + "loc": { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + }, + "type": "array", + "title": "Location" + }, + "msg": { + "type": "string", + "title": "Message" + }, + "type": { + "type": "string", + "title": "Error Type" + } + }, + "type": "object", + "required": [ + "loc", + "msg", + "type" + ], + "title": "ValidationError" + }, + "Workspace": { + "properties": { + "name": { + "type": "string", + "title": "Name" + }, + "is_active": { + "type": "boolean", + "title": "Is Active" + } + }, + "type": "object", + "required": [ + "name", + "is_active" + ], + "title": "Workspace" } } } diff --git a/pyproject.toml b/pyproject.toml index f90a6488..c38ca31b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,7 +50,7 @@ build-backend = "poetry.core.masonry.api" [tool.poetry.scripts] codegate = "codegate.cli:main" -generate-openapi = "src.codegate.dashboard.dashboard:generate_openapi" +generate-openapi = "src.codegate.server:generate_openapi" [tool.black] line-length = 100 diff --git a/src/codegate/dashboard/dashboard.py b/src/codegate/dashboard/dashboard.py index 35d1efd5..c9844cc2 100644 --- a/src/codegate/dashboard/dashboard.py +++ b/src/codegate/dashboard/dashboard.py @@ -1,5 +1,4 @@ import asyncio -import json from typing import AsyncGenerator, List, Optional import structlog @@ -60,18 +59,3 @@ async def stream_sse(): Send alerts event """ return StreamingResponse(generate_sse_events(), media_type="text/event-stream") - - -def generate_openapi(): - # Create a temporary FastAPI app instance - app = FastAPI() - - # Include your defined router - app.include_router(dashboard_router) - - # Generate OpenAPI JSON - openapi_schema = app.openapi() - - # Convert the schema to JSON string for easier handling or storage - openapi_json = json.dumps(openapi_schema, indent=2) - print(openapi_json) diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index 006616ba..304dbef0 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -51,7 +51,7 @@ def __init__(self, sqlite_path: Optional[str] = None): ) self._db_path = Path(sqlite_path).absolute() self._db_path.parent.mkdir(parents=True, exist_ok=True) - logger.debug(f"Connecting to DB from path: {self._db_path}") + # logger.debug(f"Connecting to DB from path: {self._db_path}") engine_dict = { "url": f"sqlite+aiosqlite:///{self._db_path}", "echo": False, # Set to False in production diff --git a/src/codegate/server.py b/src/codegate/server.py index d1da668e..99642b05 100644 --- a/src/codegate/server.py +++ b/src/codegate/server.py @@ -1,3 +1,4 @@ +import json import traceback import structlog @@ -102,3 +103,18 @@ async def health_check(): app.include_router(v1, prefix="/api/v1", tags=["CodeGate API"]) return app + + +def generate_openapi(): + # Create a temporary FastAPI app instance + app = FastAPI() + + app.include_router(dashboard_router) + app.include_router(v1, prefix="/api/v1", tags=["CodeGate API"]) + + # Generate OpenAPI JSON + openapi_schema = app.openapi() + + # Convert the schema to JSON string for easier handling or storage + openapi_json = json.dumps(openapi_schema, indent=2) + print(openapi_json) \ No newline at end of file From 23ec130d1ebddcf5309d5b59144b32b4f364e8b3 Mon Sep 17 00:00:00 2001 From: alex-mcgovern Date: Fri, 17 Jan 2025 15:36:30 +0000 Subject: [PATCH 2/5] feat(server): move dashboard endpoints under api v1 router --- api/openapi.json | 48 ++++++++++++++----- src/codegate/{ => api}/dashboard/dashboard.py | 8 ++-- .../{ => api}/dashboard/post_processing.py | 2 +- .../{ => api}/dashboard/request_models.py | 0 src/codegate/api/v1.py | 3 ++ src/codegate/server.py | 9 +--- tests/dashboard/test_post_processing.py | 8 ++-- tests/test_server.py | 6 +-- 8 files changed, 53 insertions(+), 31 deletions(-) rename src/codegate/{ => api}/dashboard/dashboard.py (95%) rename src/codegate/{ => api}/dashboard/post_processing.py (99%) rename src/codegate/{ => api}/dashboard/request_models.py (100%) diff --git a/api/openapi.json b/api/openapi.json index 38345cdf..080cb00f 100644 --- a/api/openapi.json +++ b/api/openapi.json @@ -1,18 +1,39 @@ { "openapi": "3.1.0", "info": { - "title": "FastAPI", - "version": "0.1.0" + "title": "CodeGate", + "description": "Generative AI CodeGen security gateway", + "version": "0.1.7" }, "paths": { - "/dashboard/messages": { + "/health": { "get": { "tags": [ + "System" + ], + "summary": "Health Check", + "operationId": "health_check_health_get", + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": {} + } + } + } + } + } + }, + "/api/v1/dashboard/messages": { + "get": { + "tags": [ + "CodeGate API", "Dashboard" ], "summary": "Get Messages", "description": "Get all the messages from the database and return them as a list of conversations.", - "operationId": "get_messages_dashboard_messages_get", + "operationId": "get_messages_api_v1_dashboard_messages_get", "responses": { "200": { "description": "Successful Response", @@ -23,7 +44,7 @@ "$ref": "#/components/schemas/Conversation" }, "type": "array", - "title": "Response Get Messages Dashboard Messages Get" + "title": "Response Get Messages Api V1 Dashboard Messages Get" } } } @@ -31,14 +52,15 @@ } } }, - "/dashboard/alerts": { + "/api/v1/dashboard/alerts": { "get": { "tags": [ + "CodeGate API", "Dashboard" ], "summary": "Get Alerts", "description": "Get all the messages from the database and return them as a list of conversations.", - "operationId": "get_alerts_dashboard_alerts_get", + "operationId": "get_alerts_api_v1_dashboard_alerts_get", "responses": { "200": { "description": "Successful Response", @@ -56,7 +78,7 @@ ] }, "type": "array", - "title": "Response Get Alerts Dashboard Alerts Get" + "title": "Response Get Alerts Api V1 Dashboard Alerts Get" } } } @@ -64,14 +86,15 @@ } } }, - "/dashboard/alerts_notification": { + "/api/v1/dashboard/alerts_notification": { "get": { "tags": [ + "CodeGate API", "Dashboard" ], "summary": "Stream Sse", "description": "Send alerts event", - "operationId": "stream_sse_dashboard_alerts_notification_get", + "operationId": "stream_sse_api_v1_dashboard_alerts_notification_get", "responses": { "200": { "description": "Successful Response", @@ -84,13 +107,14 @@ } } }, - "/dashboard/version": { + "/api/v1/dashboard/version": { "get": { "tags": [ + "CodeGate API", "Dashboard" ], "summary": "Version Check", - "operationId": "version_check_dashboard_version_get", + "operationId": "version_check_api_v1_dashboard_version_get", "responses": { "200": { "description": "Successful Response", diff --git a/src/codegate/dashboard/dashboard.py b/src/codegate/api/dashboard/dashboard.py similarity index 95% rename from src/codegate/dashboard/dashboard.py rename to src/codegate/api/dashboard/dashboard.py index 2406f5e0..125f2ae0 100644 --- a/src/codegate/dashboard/dashboard.py +++ b/src/codegate/api/dashboard/dashboard.py @@ -5,13 +5,13 @@ import structlog from fastapi import APIRouter, Depends, FastAPI from fastapi.responses import StreamingResponse - from codegate import __version__ -from codegate.dashboard.post_processing import ( + +from codegate.api.dashboard.post_processing import ( parse_get_alert_conversation, parse_messages_in_conversations, ) -from codegate.dashboard.request_models import AlertConversation, Conversation +from codegate.api.dashboard.request_models import AlertConversation, Conversation from codegate.db.connection import DbReader, alert_queue logger = structlog.get_logger("codegate") @@ -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, diff --git a/src/codegate/dashboard/post_processing.py b/src/codegate/api/dashboard/post_processing.py similarity index 99% rename from src/codegate/dashboard/post_processing.py rename to src/codegate/api/dashboard/post_processing.py index 2ffa841e..1e4135d2 100644 --- a/src/codegate/dashboard/post_processing.py +++ b/src/codegate/api/dashboard/post_processing.py @@ -6,7 +6,7 @@ import structlog -from codegate.dashboard.request_models import ( +from codegate.api.dashboard.request_models import ( AlertConversation, ChatMessage, Conversation, diff --git a/src/codegate/dashboard/request_models.py b/src/codegate/api/dashboard/request_models.py similarity index 100% rename from src/codegate/dashboard/request_models.py rename to src/codegate/api/dashboard/request_models.py diff --git a/src/codegate/api/v1.py b/src/codegate/api/v1.py index 0f8dbcd3..85892c43 100644 --- a/src/codegate/api/v1.py +++ b/src/codegate/api/v1.py @@ -6,8 +6,11 @@ from codegate.api import v1_models from codegate.db.connection import AlreadyExistsError from codegate.workspaces.crud import WorkspaceCrud +from codegate.api.dashboard.dashboard import dashboard_router v1 = APIRouter() +v1.include_router(dashboard_router) + wscrud = WorkspaceCrud() diff --git a/src/codegate/server.py b/src/codegate/server.py index 99642b05..dc9a8b0f 100644 --- a/src/codegate/server.py +++ b/src/codegate/server.py @@ -1,5 +1,6 @@ import json import traceback +from unittest.mock import Mock import structlog from fastapi import APIRouter, FastAPI, Request @@ -9,7 +10,6 @@ from codegate import __description__, __version__ from codegate.api.v1 import v1 -from codegate.dashboard.dashboard import dashboard_router from codegate.pipeline.factory import PipelineFactory from codegate.providers.anthropic.provider import AnthropicProvider from codegate.providers.llamacpp.provider import LlamaCppProvider @@ -97,7 +97,6 @@ async def health_check(): return {"status": "healthy"} app.include_router(system_router) - app.include_router(dashboard_router) # CodeGate API app.include_router(v1, prefix="/api/v1", tags=["CodeGate API"]) @@ -106,11 +105,7 @@ async def health_check(): def generate_openapi(): - # Create a temporary FastAPI app instance - app = FastAPI() - - app.include_router(dashboard_router) - app.include_router(v1, prefix="/api/v1", tags=["CodeGate API"]) + app = init_app(Mock(spec=PipelineFactory)) # Generate OpenAPI JSON openapi_schema = app.openapi() diff --git a/tests/dashboard/test_post_processing.py b/tests/dashboard/test_post_processing.py index aa35cff2..d6359efb 100644 --- a/tests/dashboard/test_post_processing.py +++ b/tests/dashboard/test_post_processing.py @@ -4,14 +4,14 @@ import pytest -from codegate.dashboard.post_processing import ( +from codegate.api.dashboard.post_processing import ( _get_question_answer, _group_partial_messages, _is_system_prompt, parse_output, parse_request, ) -from codegate.dashboard.request_models import ( +from codegate.api.dashboard.request_models import ( PartialQuestions, ) from codegate.db.models import GetPromptWithOutputsRow @@ -162,10 +162,10 @@ async def test_parse_output(output_dict, expected_str): ) async def test_get_question_answer(request_msg_list, output_msg_str, row): with patch( - "codegate.dashboard.post_processing.parse_request", new_callable=AsyncMock + "codegate.api.dashboard.post_processing.parse_request", new_callable=AsyncMock ) as mock_parse_request: with patch( - "codegate.dashboard.post_processing.parse_output", new_callable=AsyncMock + "codegate.api.dashboard.post_processing.parse_output", new_callable=AsyncMock ) as mock_parse_output: # Set return values for the mocks mock_parse_request.return_value = request_msg_list diff --git a/tests/test_server.py b/tests/test_server.py index ad3b3541..6ae0538a 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -81,10 +81,10 @@ def test_health_check(test_client: TestClient) -> None: assert response.status_code == 200 assert response.json() == {"status": "healthy"} -@patch("codegate.dashboard.dashboard.fetch_latest_version", return_value="foo") +@patch("codegate.api.dashboard.fetch_latest_version", return_value="foo") def test_version_endpoint(mock_fetch_latest_version, test_client: TestClient) -> None: """Test the version endpoint.""" - response = test_client.get("/dashboard/version") + response = test_client.get("/api/v1/dashboard/version") assert response.status_code == 200 response_data = response.json() @@ -139,7 +139,7 @@ def test_dashboard_routes(mock_pipeline_factory) -> None: routes = [route.path for route in app.routes] # Verify dashboard endpoints are included - dashboard_routes = [route for route in routes if route.startswith("/dashboard")] + dashboard_routes = [route for route in routes if route.startswith("/api/v1/dashboard")] assert len(dashboard_routes) > 0 From 7acb78f30a1555784bee6061664960950e06cb7f Mon Sep 17 00:00:00 2001 From: alex-mcgovern Date: Fri, 17 Jan 2025 15:39:01 +0000 Subject: [PATCH 3/5] fix: broken test --- tests/test_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_server.py b/tests/test_server.py index 6ae0538a..0eecd92c 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -81,7 +81,7 @@ def test_health_check(test_client: TestClient) -> None: assert response.status_code == 200 assert response.json() == {"status": "healthy"} -@patch("codegate.api.dashboard.fetch_latest_version", return_value="foo") +@patch("codegate.api.dashboard.dashboard.fetch_latest_version", return_value="foo") def test_version_endpoint(mock_fetch_latest_version, test_client: TestClient) -> None: """Test the version endpoint.""" response = test_client.get("/api/v1/dashboard/version") From 9420fe3f23e56e9ee26ef3ee2af26dc9b9433918 Mon Sep 17 00:00:00 2001 From: alex-mcgovern Date: Fri, 17 Jan 2025 15:48:42 +0000 Subject: [PATCH 4/5] chore: tidy up dashboard api naming --- api/openapi.json | 12 ++++++------ src/codegate/api/dashboard/dashboard.py | 14 +++++++++----- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/api/openapi.json b/api/openapi.json index 080cb00f..aadbee66 100644 --- a/api/openapi.json +++ b/api/openapi.json @@ -33,7 +33,7 @@ ], "summary": "Get Messages", "description": "Get all the messages from the database and return them as a list of conversations.", - "operationId": "get_messages_api_v1_dashboard_messages_get", + "operationId": "v1_get_messages", "responses": { "200": { "description": "Successful Response", @@ -44,7 +44,7 @@ "$ref": "#/components/schemas/Conversation" }, "type": "array", - "title": "Response Get Messages Api V1 Dashboard Messages Get" + "title": "Response V1 Get Messages" } } } @@ -60,7 +60,7 @@ ], "summary": "Get Alerts", "description": "Get all the messages from the database and return them as a list of conversations.", - "operationId": "get_alerts_api_v1_dashboard_alerts_get", + "operationId": "v1_get_alerts", "responses": { "200": { "description": "Successful Response", @@ -78,7 +78,7 @@ ] }, "type": "array", - "title": "Response Get Alerts Api V1 Dashboard Alerts Get" + "title": "Response V1 Get Alerts" } } } @@ -94,7 +94,7 @@ ], "summary": "Stream Sse", "description": "Send alerts event", - "operationId": "stream_sse_api_v1_dashboard_alerts_notification_get", + "operationId": "v1_stream_sse", "responses": { "200": { "description": "Successful Response", @@ -114,7 +114,7 @@ "Dashboard" ], "summary": "Version Check", - "operationId": "version_check_api_v1_dashboard_version_get", + "operationId": "v1_version_check", "responses": { "200": { "description": "Successful Response", diff --git a/src/codegate/api/dashboard/dashboard.py b/src/codegate/api/dashboard/dashboard.py index 125f2ae0..9faea546 100644 --- a/src/codegate/api/dashboard/dashboard.py +++ b/src/codegate/api/dashboard/dashboard.py @@ -1,6 +1,7 @@ import asyncio from typing import AsyncGenerator, List, Optional +from fastapi.routing import APIRoute import requests import structlog from fastapi import APIRouter, Depends, FastAPI @@ -16,9 +17,12 @@ logger = structlog.get_logger("codegate") -dashboard_router = APIRouter(tags=["Dashboard"]) +dashboard_router = APIRouter() db_reader = None +def uniq_name(route: APIRoute): + return f"v1_{route.name}" + def get_db_reader(): global db_reader if db_reader is None: @@ -36,7 +40,7 @@ def fetch_latest_version() -> str: data = response.json() return data.get("tag_name", "unknown") -@dashboard_router.get("/dashboard/messages") +@dashboard_router.get("/dashboard/messages", tags=["Dashboard"], generate_unique_id_function=uniq_name) def get_messages(db_reader: DbReader = Depends(get_db_reader)) -> List[Conversation]: """ Get all the messages from the database and return them as a list of conversations. @@ -46,7 +50,7 @@ def get_messages(db_reader: DbReader = Depends(get_db_reader)) -> List[Conversat return asyncio.run(parse_messages_in_conversations(prompts_outputs)) -@dashboard_router.get("/dashboard/alerts") +@dashboard_router.get("/dashboard/alerts", tags=["Dashboard"], generate_unique_id_function=uniq_name) def get_alerts(db_reader: DbReader = Depends(get_db_reader)) -> List[Optional[AlertConversation]]: """ Get all the messages from the database and return them as a list of conversations. @@ -64,14 +68,14 @@ async def generate_sse_events() -> AsyncGenerator[str, None]: yield f"data: {message}\n\n" -@dashboard_router.get("/dashboard/alerts_notification") +@dashboard_router.get("/dashboard/alerts_notification", tags=["Dashboard"], generate_unique_id_function=uniq_name) async def stream_sse(): """ Send alerts event """ return StreamingResponse(generate_sse_events(), media_type="text/event-stream") -@dashboard_router.get("/dashboard/version") +@dashboard_router.get("/dashboard/version", tags=["Dashboard"], generate_unique_id_function=uniq_name) def version_check(): try: latest_version = fetch_latest_version() From 520d44773b710ace8ef9bd80c27bf4910e13f242 Mon Sep 17 00:00:00 2001 From: alex-mcgovern Date: Fri, 17 Jan 2025 15:52:55 +0000 Subject: [PATCH 5/5] fix: ruff format pass --- src/codegate/api/dashboard/dashboard.py | 44 +++++++++++++++---------- 1 file changed, 27 insertions(+), 17 deletions(-) diff --git a/src/codegate/api/dashboard/dashboard.py b/src/codegate/api/dashboard/dashboard.py index 9faea546..f59b31a5 100644 --- a/src/codegate/api/dashboard/dashboard.py +++ b/src/codegate/api/dashboard/dashboard.py @@ -1,13 +1,13 @@ import asyncio from typing import AsyncGenerator, List, Optional -from fastapi.routing import APIRoute import requests import structlog -from fastapi import APIRouter, Depends, FastAPI +from fastapi import APIRouter, Depends from fastapi.responses import StreamingResponse -from codegate import __version__ +from fastapi.routing import APIRoute +from codegate import __version__ from codegate.api.dashboard.post_processing import ( parse_get_alert_conversation, parse_messages_in_conversations, @@ -20,27 +20,30 @@ dashboard_router = APIRouter() db_reader = None + def uniq_name(route: APIRoute): return f"v1_{route.name}" + def get_db_reader(): global db_reader if db_reader is None: db_reader = DbReader() return db_reader + def fetch_latest_version() -> str: url = "https://api.github.com/repos/stacklok/codegate/releases/latest" - headers = { - "Accept": "application/vnd.github+json", - "X-GitHub-Api-Version": "2022-11-28" - } + headers = {"Accept": "application/vnd.github+json", "X-GitHub-Api-Version": "2022-11-28"} response = requests.get(url, headers=headers, timeout=5) response.raise_for_status() data = response.json() return data.get("tag_name", "unknown") -@dashboard_router.get("/dashboard/messages", tags=["Dashboard"], generate_unique_id_function=uniq_name) + +@dashboard_router.get( + "/dashboard/messages", tags=["Dashboard"], generate_unique_id_function=uniq_name +) def get_messages(db_reader: DbReader = Depends(get_db_reader)) -> List[Conversation]: """ Get all the messages from the database and return them as a list of conversations. @@ -50,7 +53,9 @@ def get_messages(db_reader: DbReader = Depends(get_db_reader)) -> List[Conversat return asyncio.run(parse_messages_in_conversations(prompts_outputs)) -@dashboard_router.get("/dashboard/alerts", tags=["Dashboard"], generate_unique_id_function=uniq_name) +@dashboard_router.get( + "/dashboard/alerts", tags=["Dashboard"], generate_unique_id_function=uniq_name +) def get_alerts(db_reader: DbReader = Depends(get_db_reader)) -> List[Optional[AlertConversation]]: """ Get all the messages from the database and return them as a list of conversations. @@ -68,24 +73,29 @@ async def generate_sse_events() -> AsyncGenerator[str, None]: yield f"data: {message}\n\n" -@dashboard_router.get("/dashboard/alerts_notification", tags=["Dashboard"], generate_unique_id_function=uniq_name) +@dashboard_router.get( + "/dashboard/alerts_notification", tags=["Dashboard"], generate_unique_id_function=uniq_name +) async def stream_sse(): """ Send alerts event """ return StreamingResponse(generate_sse_events(), media_type="text/event-stream") -@dashboard_router.get("/dashboard/version", tags=["Dashboard"], generate_unique_id_function=uniq_name) + +@dashboard_router.get( + "/dashboard/version", tags=["Dashboard"], generate_unique_id_function=uniq_name +) def version_check(): try: latest_version = fetch_latest_version() # normalize the versions as github will return them with a 'v' prefix - current_version = __version__.lstrip('v') - latest_version_stripped = latest_version.lstrip('v') + current_version = __version__.lstrip("v") + 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, @@ -98,7 +108,7 @@ def version_check(): "current_version": __version__, "latest_version": "unknown", "is_latest": None, - "error": "An error occurred while fetching the latest version" + "error": "An error occurred while fetching the latest version", } except Exception as e: logger.error(f"Unexpected error: {str(e)}") @@ -106,5 +116,5 @@ def version_check(): "current_version": __version__, "latest_version": "unknown", "is_latest": None, - "error": "An unexpected error occurred" - } \ No newline at end of file + "error": "An unexpected error occurred", + }