Skip to content

Commit d01f27c

Browse files
Fix FIM for OpenRouter (#1097)
* Fix FIM for OpenRouter FIM was not working with OpenRouter in Continue. The reason was that we get FIM requests in `/completions`. LiteLLM when using `acompletion` is forcing `/chat/completions` and the return format `{..., "choices":[{"delta":{"content":"some text"}}]}`. However, Continue was expecting the format: `{..., "choices":[{"text":"some text"}]}` becuase of the endpoint it called. With this PR we force the return format to be the latter using `atext_completion` from LiteLLM * Force denormalization in OpenRouter FIM to have a prompt key
1 parent 868c687 commit d01f27c

File tree

2 files changed

+39
-6
lines changed

2 files changed

+39
-6
lines changed

src/codegate/providers/openai/provider.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,9 @@ class OpenAIProvider(BaseProvider):
1818
def __init__(
1919
self,
2020
pipeline_factory: PipelineFactory,
21+
# Enable receiving other completion handlers from childs, i.e. OpenRouter and LM Studio
22+
completion_handler: LiteLLmShim = LiteLLmShim(stream_generator=sse_stream_generator),
2123
):
22-
completion_handler = LiteLLmShim(stream_generator=sse_stream_generator)
2324
super().__init__(
2425
OpenAIInputNormalizer(),
2526
OpenAIOutputNormalizer(),

src/codegate/providers/openrouter/provider.py

+37-5
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,14 @@
22
from typing import Dict
33

44
from fastapi import Header, HTTPException, Request
5+
from litellm import atext_completion
56
from litellm.types.llms.openai import ChatCompletionRequest
67

78
from codegate.clients.clients import ClientType
89
from codegate.clients.detector import DetectClient
910
from codegate.pipeline.factory import PipelineFactory
1011
from codegate.providers.fim_analyzer import FIMAnalyzer
12+
from codegate.providers.litellmshim import LiteLLmShim, sse_stream_generator
1113
from codegate.providers.normalizer.completion import CompletionNormalizer
1214
from codegate.providers.openai import OpenAIProvider
1315

@@ -20,15 +22,45 @@ def normalize(self, data: Dict) -> ChatCompletionRequest:
2022
return super().normalize(data)
2123

2224
def denormalize(self, data: ChatCompletionRequest) -> Dict:
23-
if data.get("had_prompt_before", False):
24-
del data["had_prompt_before"]
25-
26-
return data
25+
"""
26+
Denormalize a FIM OpenRouter request. Force it to be an accepted atext_completion format.
27+
"""
28+
denormalized_data = super().denormalize(data)
29+
# We are forcing atext_completion which expects to have a "prompt" key in the data
30+
# Forcing it in case is not present
31+
if "prompt" in data:
32+
return denormalized_data
33+
custom_prompt = ""
34+
for msg_dict in denormalized_data.get("messages", []):
35+
content_obj = msg_dict.get("content")
36+
if not content_obj:
37+
continue
38+
if isinstance(content_obj, list):
39+
for content_dict in content_obj:
40+
custom_prompt += (
41+
content_dict.get("text", "") if isinstance(content_dict, dict) else ""
42+
)
43+
elif isinstance(content_obj, str):
44+
custom_prompt += content_obj
45+
46+
# Erase the original "messages" key. Replace it by "prompt"
47+
del denormalized_data["messages"]
48+
denormalized_data["prompt"] = custom_prompt
49+
50+
return denormalized_data
2751

2852

2953
class OpenRouterProvider(OpenAIProvider):
3054
def __init__(self, pipeline_factory: PipelineFactory):
31-
super().__init__(pipeline_factory)
55+
super().__init__(
56+
pipeline_factory,
57+
# We get FIM requests in /completions. LiteLLM is forcing /chat/completions
58+
# which returns "choices":[{"delta":{"content":"some text"}}]
59+
# instead of "choices":[{"text":"some text"}] expected by the client (Continue)
60+
completion_handler=LiteLLmShim(
61+
stream_generator=sse_stream_generator, fim_completion_func=atext_completion
62+
),
63+
)
3264
self._fim_normalizer = OpenRouterNormalizer()
3365

3466
@property

0 commit comments

Comments
 (0)