From cdbac96910919e3e03dbf60c8843013e9773b5fa Mon Sep 17 00:00:00 2001 From: Alex McGovern <58784948+alex-mcgovern@users.noreply.github.com> Date: Mon, 20 Jan 2025 07:31:00 +0000 Subject: [PATCH] Add workspaces to OpenAPI spec (#634) * Add workspaces to OpenAPI spec * feat(server): move dashboard endpoints under api v1 router * fix: broken test * chore: tidy up dashboard api naming * fix: ruff format pass --- api/openapi.json | 364 +++++++++++++++++- pyproject.toml | 2 +- src/codegate/{ => api}/dashboard/dashboard.py | 62 ++- .../{ => api}/dashboard/post_processing.py | 2 +- .../{ => api}/dashboard/request_models.py | 0 src/codegate/api/v1.py | 3 + src/codegate/db/connection.py | 2 +- src/codegate/server.py | 15 +- tests/dashboard/test_post_processing.py | 8 +- tests/test_server.py | 6 +- 10 files changed, 408 insertions(+), 56 deletions(-) rename src/codegate/{ => api}/dashboard/dashboard.py (71%) 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 61049270..aadbee66 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": "v1_get_messages", "responses": { "200": { "description": "Successful Response", @@ -23,7 +44,7 @@ "$ref": "#/components/schemas/Conversation" }, "type": "array", - "title": "Response Get Messages Dashboard Messages Get" + "title": "Response V1 Get Messages" } } } @@ -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": "v1_get_alerts", "responses": { "200": { "description": "Successful Response", @@ -56,7 +78,7 @@ ] }, "type": "array", - "title": "Response Get Alerts Dashboard Alerts Get" + "title": "Response V1 Get Alerts" } } } @@ -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": "v1_stream_sse", "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": "v1_version_check", "responses": { "200": { "description": "Successful Response", @@ -102,10 +126,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": { @@ -287,6 +518,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": { @@ -310,6 +599,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 26564796..8c330b1e 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/api/dashboard/dashboard.py similarity index 71% rename from src/codegate/dashboard/dashboard.py rename to src/codegate/api/dashboard/dashboard.py index ee71424f..f59b31a5 100644 --- a/src/codegate/dashboard/dashboard.py +++ b/src/codegate/api/dashboard/dashboard.py @@ -1,43 +1,49 @@ import asyncio -import json from typing import AsyncGenerator, List, Optional import requests import structlog -from fastapi import APIRouter, Depends, FastAPI +from fastapi import APIRouter, Depends from fastapi.responses import StreamingResponse +from fastapi.routing import APIRoute 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") -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: 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") + +@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. @@ -47,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") +@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. @@ -65,21 +73,26 @@ 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() # 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 @@ -95,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)}") @@ -103,20 +116,5 @@ def version_check(): "current_version": __version__, "latest_version": "unknown", "is_latest": None, - "error": "An unexpected error occurred" + "error": "An unexpected error occurred", } - - -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/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/db/connection.py b/src/codegate/db/connection.py index b83ceb7c..111c582f 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -54,7 +54,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..dc9a8b0f 100644 --- a/src/codegate/server.py +++ b/src/codegate/server.py @@ -1,4 +1,6 @@ +import json import traceback +from unittest.mock import Mock import structlog from fastapi import APIRouter, FastAPI, Request @@ -8,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 @@ -96,9 +97,19 @@ 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"]) return app + + +def generate_openapi(): + app = init_app(Mock(spec=PipelineFactory)) + + # 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 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..0eecd92c 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.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