Add a FIM pipeline to providers

aponcedeleonch · aponcedeleonch · commit c36a53e78ea9 · 2024-11-28T10:05:46.000+02:00
Related: #87, #43 The PR adds a FIM pipeline independent from chat completion pipeline. It could still be faulty since we need: - Message normalizer. We now expect all messages to have the key `messages`. However, there are incoming messages with `prompt`. - Secreets detector. There's the skeleton of a class called SecretAnalyzer that is meant to analyze the messages and return a warning if it detected a secret.
diff --git a/src/codegate/pipeline/fim/secret_analyzer.py b/src/codegate/pipeline/fim/secret_analyzer.py
@@ -0,0 +1,48 @@
+from litellm import ChatCompletionRequest
+
+from codegate.pipeline.base import PipelineContext, PipelineResult, PipelineStep, PipelineResponse
+
+
+class SecretAnalyzer(PipelineStep):
+    """Pipeline step that handles version information requests."""
+
+    message_blocked = """
+        ⚠️ CodeGate Security Warning! Analysis Report ⚠️
+        Potential leak of sensitive credentials blocked
+
+        Recommendations:
+        - Use environment variables for secrets
+    """
+
+    @property
+    def name(self) -> str:
+        """
+        Returns the name of this pipeline step.
+
+        Returns:
+            str: The identifier 'fim-secret-analyzer'
+        """
+        return "fim-secret-analyzer"
+
+    async def process(
+            self,
+            request: ChatCompletionRequest,
+            context: PipelineContext
+    ) -> PipelineResult:
+        # We should call here Secrets Blocking module to see if the request messages contain secrets
+        # messages_contain_secrets = [analyze_msg_secrets(msg) for msg in request.messages]
+        # message_with_secrets = any(messages_contain_secretes)
+
+        # For the moment to test shortcutting just treat all messages as if they contain secrets
+        message_with_secrets = True
+        if message_with_secrets:
+            return PipelineResult(
+                response=PipelineResponse(
+                    step_name=self.name,
+                    content=self.message_blocked,
+                    model=request["model"],
+                ),
+            )
+
+        # No messages with secrets, execute the rest of the pipeline
+        return PipelineResult(request=request)
diff --git a/src/codegate/providers/anthropic/provider.py b/src/codegate/providers/anthropic/provider.py
@@ -1,4 +1,5 @@
 import json
+from typing import Optional
 
 from fastapi import Header, HTTPException, Request
 
diff --git a/src/codegate/providers/litellmshim/litellmshim.py b/src/codegate/providers/litellmshim/litellmshim.py
@@ -43,3 +43,11 @@ def create_streaming_response(self, stream: AsyncIterator[Any]) -> StreamingResp
             },
             status_code=200,
         )
+
+    def is_fim_request(self, data: Dict) -> bool:
+        """
+        Determine from the raw incoming data if it's a FIM request.
+        This is needed here since completion_handler is used by provider and provider
+        doesn't know about the adapter.
+        """
+        return self._adapter.is_fim_request(data)
diff --git a/src/codegate/providers/llamacpp/completion_handler.py b/src/codegate/providers/llamacpp/completion_handler.py
@@ -65,3 +65,11 @@ def create_streaming_response(self, stream: Iterator[Any]) -> StreamingResponse:
             },
             status_code=200,
         )
+
+    def is_fim_request(self, data: Dict) -> bool:
+        """
+        Determine from the raw incoming data if it's a FIM request.
+        This is needed here since completion_handler is used by provider and provider
+        doesn't know about the adapter.
+        """
+        return self._adapter.is_fim_request(data)
diff --git a/src/codegate/providers/llamacpp/provider.py b/src/codegate/providers/llamacpp/provider.py
@@ -1,7 +1,9 @@
 import json
+from typing import Optional
 
 from fastapi import Request
 
+from codegate.pipeline.base import SequentialPipelineProcessor
 from codegate.providers.base import BaseProvider
 from codegate.providers.llamacpp.completion_handler import LlamaCppCompletionHandler
 from codegate.providers.llamacpp.normalizer import LLamaCppInputNormalizer, LLamaCppOutputNormalizer
diff --git a/src/codegate/providers/openai/provider.py b/src/codegate/providers/openai/provider.py
@@ -1,7 +1,9 @@
 import json
+from typing import Optional
 
 from fastapi import Header, HTTPException, Request
 
+from codegate.pipeline.base import SequentialPipelineProcessor
 from codegate.providers.base import BaseProvider
 from codegate.providers.litellmshim import LiteLLmShim, sse_stream_generator
 from codegate.providers.openai.adapter import OpenAIInputNormalizer, OpenAIOutputNormalizer
@@ -42,3 +44,18 @@ async def create_completion(
 
             stream = await self.complete(data, api_key)
             return self._completion_handler.create_streaming_response(stream)
+
+        @self.router.post(f"/{self.provider_route_name}/completions")
+        async def create_fim(
+            request: Request,
+            authorization: str = Header(..., description="Bearer token"),
+        ):
+            if not authorization.startswith("Bearer "):
+                raise HTTPException(status_code=401, detail="Invalid authorization header")
+
+            api_key = authorization.split(" ")[1]
+            body = await request.body()
+            data = json.loads(body)
+
+            stream = await self.complete(data, api_key)
+            return self._completion_handler.create_streaming_response(stream)
diff --git a/src/codegate/server.py b/src/codegate/server.py
@@ -5,6 +5,7 @@
 from codegate import __description__, __version__
 from codegate.pipeline.base import PipelineStep, SequentialPipelineProcessor
 from codegate.pipeline.version.version import CodegateVersion
+from codegate.pipeline.fim.secret_analyzer import SecretAnalyzer
 from codegate.providers.anthropic.provider import AnthropicProvider
 from codegate.providers.llamacpp.provider import LlamaCppProvider
 from codegate.providers.openai.provider import OpenAIProvider
@@ -21,15 +22,20 @@ def init_app() -> FastAPI:
     steps: List[PipelineStep] = [
         CodegateVersion(),
     ]
+    fim_steps: List[PipelineStep] = [
+        SecretAnalyzer(),
+    ]
 
     pipeline = SequentialPipelineProcessor(steps)
+    fim_pipeline = SequentialPipelineProcessor(fim_steps)
+
     # Create provider registry
     registry = ProviderRegistry(app)
 
     # Register all known providers
-    registry.add_provider("openai", OpenAIProvider(pipeline_processor=pipeline))
-    registry.add_provider("anthropic", AnthropicProvider(pipeline_processor=pipeline))
-    registry.add_provider("llamacpp", LlamaCppProvider(pipeline_processor=pipeline))
+    registry.add_provider("openai", OpenAIProvider(pipeline_processor=pipeline, fim_pipeline_processor=fim_pipeline))
+    registry.add_provider("anthropic", AnthropicProvider(pipeline_processor=pipeline, fim_pipeline_processor=fim_pipeline))
+    registry.add_provider("llamacpp", LlamaCppProvider(pipeline_processor=pipeline, fim_pipeline_processor=fim_pipeline))
 
     # Create and add system routes
     system_router = APIRouter(tags=["System"])  # Tags group endpoints in the docs

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`import json`
	`2`	`+from typing import Optional`
`2`	`3`
`3`	`4`	`from fastapi import Header, HTTPException, Request`
`4`	`5`