forked from stacklok/codegate
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathserver.py
More file actions
104 lines (88 loc) · 3.12 KB
/
server.py
File metadata and controls
104 lines (88 loc) · 3.12 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import traceback
import structlog
from fastapi import APIRouter, FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from starlette.middleware.errors import ServerErrorMiddleware
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
from codegate.providers.ollama.provider import OllamaProvider
from codegate.providers.openai.provider import OpenAIProvider
from codegate.providers.registry import ProviderRegistry
from codegate.providers.vllm.provider import VLLMProvider
logger = structlog.get_logger("codegate")
async def custom_error_handler(request, exc: Exception):
"""This is a Middleware to handle exceptions and log them."""
# Capture the stack trace
extracted_traceback = traceback.extract_tb(exc.__traceback__)
# Log only the last 3 items of the stack trace. 3 is an arbitrary number.
logger.error(traceback.print_list(extracted_traceback[-3:]))
return JSONResponse({"error": str(exc)}, status_code=500)
def init_app(pipeline_factory: PipelineFactory) -> FastAPI:
"""Create the FastAPI application."""
app = FastAPI(
title="CodeGate",
description=__description__,
version=__version__,
)
@app.middleware("http")
async def log_user_agent(request: Request, call_next):
user_agent = request.headers.get("user-agent")
client_host = request.client.host if request.client else "unknown"
logger.debug(f"User-Agent header received: {user_agent} from {client_host}")
response = await call_next(request)
return response
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Apply error handling middleware
app.add_middleware(ServerErrorMiddleware, handler=custom_error_handler)
# Create provider registry
registry = ProviderRegistry(app)
# Register all known providers
registry.add_provider(
"openai",
OpenAIProvider(pipeline_factory),
)
registry.add_provider(
"anthropic",
AnthropicProvider(
pipeline_factory,
),
)
registry.add_provider(
"llamacpp",
LlamaCppProvider(
pipeline_factory,
),
)
registry.add_provider(
"vllm",
VLLMProvider(
pipeline_factory,
),
)
registry.add_provider(
"ollama",
OllamaProvider(
pipeline_factory,
),
)
# Create and add system routes
system_router = APIRouter(tags=["System"])
@system_router.get("/health")
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