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
124 lines (109 loc) · 4.32 KB
/
server.py
File metadata and controls
124 lines (109 loc) · 4.32 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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
from typing import List
from fastapi import APIRouter, FastAPI
from codegate import __description__, __version__
from codegate.config import Config
from codegate.dashboard.dashboard import dashboard_router
from codegate.pipeline.base import PipelineStep, SequentialPipelineProcessor
from codegate.pipeline.codegate_context_retriever.codegate import CodegateContextRetriever
from codegate.pipeline.codegate_system_prompt.codegate import CodegateSystemPrompt
from codegate.pipeline.extract_snippets.extract_snippets import CodeSnippetExtractor
from codegate.pipeline.extract_snippets.output import CodeCommentStep
from codegate.pipeline.output import OutputPipelineProcessor, OutputPipelineStep
from codegate.pipeline.secrets.manager import SecretsManager
from codegate.pipeline.secrets.secrets import CodegateSecrets, SecretUnredactionStep
from codegate.pipeline.secrets.signatures import CodegateSignatures
from codegate.pipeline.version.version import CodegateVersion
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
def init_app() -> FastAPI:
"""Create the FastAPI application."""
app = FastAPI(
title="CodeGate",
description=__description__,
version=__version__,
)
# Initialize secrets manager
# TODO: we need to clean up the secrets manager
# after the conversation is concluded
# this was done in the pipeline step but I just removed it for now
secrets_manager = SecretsManager()
steps: List[PipelineStep] = [
CodegateVersion(),
CodeSnippetExtractor(),
CodegateSystemPrompt(Config.get_config().prompts.codegate_chat),
CodegateContextRetriever(Config.get_config().prompts.codegate_chat),
CodegateSecrets(),
]
# Leaving the pipeline empty for now
fim_steps: List[PipelineStep] = []
pipeline = SequentialPipelineProcessor(steps)
fim_pipeline = SequentialPipelineProcessor(fim_steps)
output_steps: List[OutputPipelineStep] = [
SecretUnredactionStep(),
CodeCommentStep(),
]
output_pipeline = OutputPipelineProcessor(output_steps)
# Create provider registry
registry = ProviderRegistry(app)
# Initialize SignaturesFinder
CodegateSignatures.initialize("signatures.yaml")
# Register all known providers
registry.add_provider(
"openai",
OpenAIProvider(
secrets_manager=secrets_manager,
pipeline_processor=pipeline,
fim_pipeline_processor=fim_pipeline,
output_pipeline_processor=output_pipeline,
),
)
registry.add_provider(
"anthropic",
AnthropicProvider(
secrets_manager=secrets_manager,
pipeline_processor=pipeline,
fim_pipeline_processor=fim_pipeline,
output_pipeline_processor=output_pipeline,
),
)
registry.add_provider(
"llamacpp",
LlamaCppProvider(
secrets_manager=secrets_manager,
pipeline_processor=pipeline,
fim_pipeline_processor=fim_pipeline,
output_pipeline_processor=output_pipeline,
),
)
registry.add_provider(
"vllm",
VLLMProvider(
secrets_manager=secrets_manager,
pipeline_processor=pipeline,
fim_pipeline_processor=fim_pipeline,
output_pipeline_processor=output_pipeline,
),
)
registry.add_provider(
"ollama",
OllamaProvider(
secrets_manager=secrets_manager,
pipeline_processor=pipeline,
fim_pipeline_processor=fim_pipeline,
output_pipeline_processor=output_pipeline,
),
)
# Create and add system routes
system_router = APIRouter(tags=["System"]) # Tags group endpoints in the docs
@system_router.get("/health")
async def health_check():
return {"status": "healthy"}
# Include the router in the app - this exposes the health check endpoint
app.include_router(system_router)
# Include the routes for the dashboard
app.include_router(dashboard_router)
return app