diff --git a/src/codegate/providers/openrouter/provider.py b/src/codegate/providers/openrouter/provider.py
index c3124282..13aad02a 100644
--- a/src/codegate/providers/openrouter/provider.py
+++ b/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,
diff --git a/tests/integration/openrouter/testcases.yaml b/tests/integration/openrouter/testcases.yaml
new file mode 100644
index 00000000..1ced50b7
--- /dev/null
+++ b/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": [
+ "",
+ "/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 XML tag, including context-aware indentation, if needed. All completions MUST be truthful, accurate, well-written and correct.\n\n## EXAMPLE QUERY:\n\n\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\n\nTASK: Fill the {{FILL_HERE}} hole.\n\n## CORRECT COMPLETION\n\nif (i % 2 === 0) {\n sum += i;\n }\n\n## EXAMPLE QUERY:\n\n\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\n\n## CORRECT COMPLETION:\n\n total += x\n\n## EXAMPLE QUERY:\n\n\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\n\n## CORRECT COMPLETION:\n\ntype Tree\n = {$:\"Node\", lft: Tree, rgt: Tree}\n | {$:\"Leaf\", val: T};\n\nfunction sum(tree: Tree): number {\n switch (tree.$) {\n case \"Node\":\n return sum(tree.lft) + sum(tree.rgt);\n case \"Leaf\":\n return tree.val;\n }\n}\n\n## EXAMPLE QUERY:\n\nThe 5th {{FILL_HERE}} is Jupiter.\n\n## CORRECT COMPLETION:\n\nplanet from the Sun\n\n## EXAMPLE QUERY:\n\nfunction hypothenuse(a, b) {\n return Math.sqrt({{FILL_HERE}}b ** 2);\n}\n\n## CORRECT COMPLETION:\n\na ** 2 + \n\n\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\nTASK: Fill the {{FILL_HERE}} hole. Answer only with the CORRECT completion, and NOTHING ELSE. Do it now.\n"
+ }
+ ]
+ }
+ ],
+ "system": ""
+ }
+ likes: |
+ 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