Partial fixes for FIM with OpenRouter

src/codegate/providers/openrouter/provider.py

@@ -1,16 +1,34 @@
 import json
+from typing import Dict
 
 from fastapi import Header, HTTPException, Request
+from litellm.types.llms.openai import ChatCompletionRequest
 
 from codegate.clients.detector import DetectClient
 from codegate.pipeline.factory import PipelineFactory
 from codegate.providers.fim_analyzer import FIMAnalyzer
+from codegate.providers.normalizer.completion import CompletionNormalizer
 from codegate.providers.openai import OpenAIProvider
 
 
+class OpenRouterNormalizer(CompletionNormalizer):
+    def __init__(self):
+        super().__init__()
+
+    def normalize(self, data: Dict) -> ChatCompletionRequest:
+        return super().normalize(data)
+
+    def denormalize(self, data: ChatCompletionRequest) -> Dict:
+        if data.get("had_prompt_before", False):
+            del data["had_prompt_before"]
+
+        return data
+
+
 class OpenRouterProvider(OpenAIProvider):
     def __init__(self, pipeline_factory: PipelineFactory):
         super().__init__(pipeline_factory)
+        self._fim_normalizer = OpenRouterNormalizer()
 
     @property
     def provider_route_name(self) -> str:
@@ -19,6 +37,7 @@ def provider_route_name(self) -> str:
     def _setup_routes(self):
         @self.router.post(f"/{self.provider_route_name}/api/v1/chat/completions")
         @self.router.post(f"/{self.provider_route_name}/chat/completions")
+        @self.router.post(f"/{self.provider_route_name}/completions")
         @DetectClient()
         async def create_completion(
             request: Request,

tests/integration/openrouter/testcases.yaml

@@ -0,0 +1,87 @@
+headers:
+  openrouter:
+    Authorization: Bearer ENV_OPENROUTER_KEY
+
+testcases:
+  anthropic_chat:
+    name: Openrouter Chat
+    provider: openrouter
+    url: http://localhost:8989/openrouter/api/v1/chat/completions
+    data: |
+      {
+        "max_tokens":4096,
+        "messages":[
+            {
+              "content":"You are a coding assistant.",
+              "role":"system"
+            },
+            {
+              "content":"Reply with that exact sentence: Hello from the integration tests!",
+              "role":"user"
+            }
+        ],
+        "model":"anthropic/claude-3-5-haiku",
+        "stream":true,
+        "temperature":0
+      }
+    likes: |
+      Hello from the integration tests!
+
+  anthropic_fim:
+    name: Openrouter FIM
+    provider: openrouter
+    url: http://localhost:8989/openrouter/completions
+    data: |
+      {
+        "top_k": 50,
+        "temperature": 0,
+        "max_tokens": 4096,
+        "model": "anthropic/claude-3-5-haiku-20241022",
+        "stop_sequences": [
+          "</COMPLETION>",
+          "/src/",
+          "#- coding: utf-8",
+          "```"
+        ],
+        "stream": true,
+        "messages": [
+          {
+            "role": "user",
+            "content": [
+              {
+                "type": "text",
+                "text": "You are a HOLE FILLER. You are provided with a file containing holes, formatted as '{{HOLE_NAME}}'. Your TASK is to complete with a string to replace this hole with, inside a <COMPLETION/> XML tag, including context-aware indentation, if needed. All completions MUST be truthful, accurate, well-written and correct.\n\n## EXAMPLE QUERY:\n\n<QUERY>\nfunction sum_evens(lim) {\n  var sum = 0;\n  for (var i = 0; i < lim; ++i) {\n    {{FILL_HERE}}\n  }\n  return sum;\n}\n</QUERY>\n\nTASK: Fill the {{FILL_HERE}} hole.\n\n## CORRECT COMPLETION\n\n<COMPLETION>if (i % 2 === 0) {\n      sum += i;\n    }</COMPLETION>\n\n## EXAMPLE QUERY:\n\n<QUERY>\ndef sum_list(lst):\n  total = 0\n  for x in lst:\n  {{FILL_HERE}}\n  return total\n\nprint sum_list([1, 2, 3])\n</QUERY>\n\n## CORRECT COMPLETION:\n\n<COMPLETION>  total += x</COMPLETION>\n\n## EXAMPLE QUERY:\n\n<QUERY>\n// data Tree a = Node (Tree a) (Tree a) | Leaf a\n\n// sum :: Tree Int -> Int\n// sum (Node lft rgt) = sum lft + sum rgt\n// sum (Leaf val)     = val\n\n// convert to TypeScript:\n{{FILL_HERE}}\n</QUERY>\n\n## CORRECT COMPLETION:\n\n<COMPLETION>type Tree<T>\n  = {$:\"Node\", lft: Tree<T>, rgt: Tree<T>}\n  | {$:\"Leaf\", val: T};\n\nfunction sum(tree: Tree<number>): number {\n  switch (tree.$) {\n    case \"Node\":\n      return sum(tree.lft) + sum(tree.rgt);\n    case \"Leaf\":\n      return tree.val;\n  }\n}</COMPLETION>\n\n## EXAMPLE QUERY:\n\nThe 5th {{FILL_HERE}} is Jupiter.\n\n## CORRECT COMPLETION:\n\n<COMPLETION>planet from the Sun</COMPLETION>\n\n## EXAMPLE QUERY:\n\nfunction hypothenuse(a, b) {\n  return Math.sqrt({{FILL_HERE}}b ** 2);\n}\n\n## CORRECT COMPLETION:\n\n<COMPLETION>a ** 2 + </COMPLETION>\n\n<QUERY>\n# Path: Untitled.txt\n# http://127.0.0.1:8989/vllm/completions\n# codegate/test.py\nimport requests\n\ndef call_api():\n    {{FILL_HERE}}\n\n\ndata = {'key1': 'test1', 'key2': 'test2'}\nresponse = call_api('http://localhost:8080', method='post', data='data')\n</QUERY>\nTASK: Fill the {{FILL_HERE}} hole. Answer only with the CORRECT completion, and NOTHING ELSE. Do it now.\n<COMPLETION>"
+              }
+            ]
+          }
+        ],
+        "system": ""
+      }
+    likes: |
+      <COMPLETION>def call_api(url, method='get', data=None):
+        if method.lower() == 'get':
+            return requests.get(url)
+        elif method.lower() == 'post':
+            return requests.post(url, json=data)
+        else:
+            raise ValueError("Unsupported HTTP method")
+
+  anthropic_malicious_package_question:
+    name: Openrouter Malicious Package
+    provider: openrouter
+    url: http://localhost:8989/openrouter/api/v1/chat/completions
+    data: |
+      {
+        "messages":[
+            {
+              "content":"Generate me example code using the python invokehttp package to call an API",
+              "role":"user"
+            }
+        ],
+        "model":"anthropic/claude-3-5-haiku-20241022",
+        "stream":true
+      }
+    contains: |
+      https://www.insight.stacklok.com/report/pypi/invokehttp?utm_source=codegate
+    does_not_contain: |
+      import invokehttp