Move instantiating the pipelines into a pipeline factory

This would allow us to reuse the pipelines in the copilot proxy server
without having to redefine them over and over.

Author: jhrozek

src/codegate/cli.py

@@ -15,6 +15,8 @@
 from codegate.codegate_logging import LogFormat, LogLevel, setup_logging
 from codegate.config import Config, ConfigurationError
 from codegate.db.connection import init_db_sync
+from codegate.pipeline.factory import PipelineFactory
+from codegate.pipeline.secrets.manager import SecretsManager
 from codegate.providers.copilot.provider import CopilotProvider
 from codegate.server import init_app
 from codegate.storage.utils import restore_storage_backup
@@ -307,7 +309,11 @@ def serve(
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
 
-        app = init_app()
+        # Initialize secrets manager and pipeline factory
+        secrets_manager = SecretsManager()
+        pipeline_factory = PipelineFactory(secrets_manager)
+
+        app = init_app(pipeline_factory)
 
         # Run the server
         try:

src/codegate/pipeline/factory.py

@@ -0,0 +1,56 @@
+from typing import List
+
+from codegate.config import Config
+from codegate.pipeline.base import PipelineStep, SequentialPipelineProcessor
+from codegate.pipeline.codegate_context_retriever.codegate import CodegateContextRetriever
+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,
+    SecretRedactionNotifier,
+    SecretUnredactionStep,
+)
+from codegate.pipeline.system_prompt.codegate import SystemPrompt
+from codegate.pipeline.version.version import CodegateVersion
+
+
+class PipelineFactory:
+    def __init__(self, secrets_manager: SecretsManager):
+        self.secrets_manager = secrets_manager
+
+    def create_input_pipeline(self) -> SequentialPipelineProcessor:
+        input_steps: List[PipelineStep] = [
+            # make sure that this step is always first in the pipeline
+            # the other steps might send the request to a LLM for it to be analyzed
+            # and without obfuscating the secrets, we'd leak the secrets during those
+            # later steps
+            CodegateSecrets(),
+            CodegateVersion(),
+            CodeSnippetExtractor(),
+            SystemPrompt(Config.get_config().prompts.default_chat),
+            CodegateContextRetriever(),
+        ]
+        return SequentialPipelineProcessor(input_steps, self.secrets_manager)
+
+    def create_fim_pipeline(self) -> SequentialPipelineProcessor:
+        fim_steps: List[PipelineStep] = [
+            CodegateSecrets(),
+        ]
+        return SequentialPipelineProcessor(fim_steps, self.secrets_manager)
+
+    def create_output_pipeline(self) -> OutputPipelineProcessor:
+        output_steps: List[OutputPipelineStep] = [
+            SecretRedactionNotifier(),
+            SecretUnredactionStep(),
+            CodeCommentStep(),
+        ]
+        return OutputPipelineProcessor(output_steps)
+
+    def create_fim_output_pipeline(self) -> OutputPipelineProcessor:
+        fim_output_steps: List[OutputPipelineStep] = [
+            # temporarily disabled
+            # SecretUnredactionStep(),
+        ]
+        return OutputPipelineProcessor(fim_output_steps)

src/codegate/server.py

@@ -1,25 +1,10 @@
-from typing import List
-
 from fastapi import APIRouter, FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
 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.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,
-    SecretRedactionNotifier,
-    SecretUnredactionStep,
-)
+from codegate.pipeline.factory import PipelineFactory
 from codegate.pipeline.secrets.signatures import CodegateSignatures
-from codegate.pipeline.system_prompt.codegate import SystemPrompt
-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
@@ -28,7 +13,7 @@
 from codegate.providers.vllm.provider import VLLMProvider
 
 
-def init_app() -> FastAPI:
+def init_app(pipeline_factory: PipelineFactory) -> FastAPI:
     """Create the FastAPI application."""
     app = FastAPI(
         title="CodeGate",
@@ -43,44 +28,6 @@ def init_app() -> FastAPI:
         allow_headers=["*"],
     )
 
-    # 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()
-
-    # Define input pipeline steps
-    input_steps: List[PipelineStep] = [
-        CodegateSecrets(),
-        CodegateVersion(),
-        CodeSnippetExtractor(),
-        SystemPrompt(Config.get_config().prompts.default_chat),
-        CodegateContextRetriever(),
-    ]
-
-    # Define FIM pipeline steps
-    fim_steps: List[PipelineStep] = [
-        CodegateSecrets(),
-    ]
-
-    # Initialize input pipeline processors
-    input_pipeline_processor = SequentialPipelineProcessor(input_steps, secrets_manager)
-    fim_pipeline_processor = SequentialPipelineProcessor(fim_steps, secrets_manager)
-
-    # Define output pipeline steps
-    output_steps: List[OutputPipelineStep] = [
-        SecretRedactionNotifier(),
-        SecretUnredactionStep(),
-        CodeCommentStep(),
-    ]
-    fim_output_steps: List[OutputPipelineStep] = [
-        # temporarily disabled
-        # SecretUnredactionStep(),
-    ]
-
-    output_pipeline_processor = OutputPipelineProcessor(output_steps)
-    fim_output_pipeline_processor = OutputPipelineProcessor(fim_output_steps)
-
     # Create provider registry
     registry = ProviderRegistry(app)
 
@@ -91,60 +38,57 @@ def init_app() -> FastAPI:
     registry.add_provider(
         "openai",
         OpenAIProvider(
-            pipeline_processor=input_pipeline_processor,
-            fim_pipeline_processor=fim_pipeline_processor,
-            output_pipeline_processor=output_pipeline_processor,
-            fim_output_pipeline_processor=fim_output_pipeline_processor,
+            pipeline_processor=pipeline_factory.create_input_pipeline(),
+            fim_pipeline_processor=pipeline_factory.create_fim_pipeline(),
+            output_pipeline_processor=pipeline_factory.create_output_pipeline(),
+            fim_output_pipeline_processor=pipeline_factory.create_fim_output_pipeline(),
         ),
     )
     registry.add_provider(
         "anthropic",
         AnthropicProvider(
-            pipeline_processor=input_pipeline_processor,
-            fim_pipeline_processor=fim_pipeline_processor,
-            output_pipeline_processor=output_pipeline_processor,
-            fim_output_pipeline_processor=fim_output_pipeline_processor,
+            pipeline_processor=pipeline_factory.create_input_pipeline(),
+            fim_pipeline_processor=pipeline_factory.create_fim_pipeline(),
+            output_pipeline_processor=pipeline_factory.create_output_pipeline(),
+            fim_output_pipeline_processor=pipeline_factory.create_fim_output_pipeline(),
         ),
     )
     registry.add_provider(
         "llamacpp",
         LlamaCppProvider(
-            pipeline_processor=input_pipeline_processor,
-            fim_pipeline_processor=fim_pipeline_processor,
-            output_pipeline_processor=output_pipeline_processor,
-            fim_output_pipeline_processor=fim_output_pipeline_processor,
+            pipeline_processor=pipeline_factory.create_input_pipeline(),
+            fim_pipeline_processor=pipeline_factory.create_fim_pipeline(),
+            output_pipeline_processor=pipeline_factory.create_output_pipeline(),
+            fim_output_pipeline_processor=pipeline_factory.create_fim_output_pipeline(),
         ),
     )
     registry.add_provider(
         "vllm",
         VLLMProvider(
-            pipeline_processor=input_pipeline_processor,
-            fim_pipeline_processor=fim_pipeline_processor,
-            output_pipeline_processor=output_pipeline_processor,
-            fim_output_pipeline_processor=fim_output_pipeline_processor,
+            pipeline_processor=pipeline_factory.create_input_pipeline(),
+            fim_pipeline_processor=pipeline_factory.create_fim_pipeline(),
+            output_pipeline_processor=pipeline_factory.create_output_pipeline(),
+            fim_output_pipeline_processor=pipeline_factory.create_fim_output_pipeline(),
         ),
     )
     registry.add_provider(
         "ollama",
         OllamaProvider(
-            pipeline_processor=input_pipeline_processor,
-            fim_pipeline_processor=fim_pipeline_processor,
-            output_pipeline_processor=output_pipeline_processor,
-            fim_output_pipeline_processor=fim_output_pipeline_processor,
+            pipeline_processor=pipeline_factory.create_input_pipeline(),
+            fim_pipeline_processor=pipeline_factory.create_fim_pipeline(),
+            output_pipeline_processor=pipeline_factory.create_output_pipeline(),
+            fim_output_pipeline_processor=pipeline_factory.create_fim_output_pipeline(),
         ),
     )
 
     # Create and add system routes
-    system_router = APIRouter(tags=["System"])  # Tags group endpoints in the docs
+    system_router = APIRouter(tags=["System"])
 
     @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