openai fim fixed

Author: aponcedeleonch

src/codegate/providers/litellmshim/litellmshim.py

@@ -3,11 +3,7 @@
 import litellm
 import structlog
 from fastapi.responses import JSONResponse, StreamingResponse
-from litellm import (
-    ChatCompletionRequest,
-    ModelResponse,
-    acompletion,
-)
+from litellm import ChatCompletionRequest, ModelResponse, acompletion, atext_completion
 
 from codegate.clients.clients import ClientType
 from codegate.providers.base import BaseCompletionHandler, StreamGenerator
@@ -52,6 +48,11 @@ async def execute_completion(
         request["api_key"] = api_key
         request["base_url"] = base_url
         if is_fim_request:
+            # We need to force atext_completion if there is "prompt" in the request.
+            # The default function acompletion can only handle "messages" in the request.
+            if "prompt" in request:
+                logger.debug("Forcing atext_completion in FIM")
+                return await atext_completion(**request)
             return await self._fim_completion_func(**request)
         return await self._completion_func(**request)
 

src/codegate/providers/normalizer/completion.py

@@ -19,13 +19,9 @@ def normalize(self, data: Dict) -> ChatCompletionRequest:
         if "prompt" in data:
             data["messages"] = [{"content": data.pop("prompt"), "role": "user"}]
             # We can add as many parameters as we like to data. ChatCompletionRequest is not strict.
+            data["had_prompt_before"] = True
 
-            # NOTE: Not adding the flag. This will also skip the denormalize step.
-            # LiteLLM seems to not support anymore having "prompt" as key and only "message" is
-            # supported
-            # data["had_prompt_before"] = True
-
-        # Litelllm says the we need to have max a list of length 4 in stop.
+        # Litelllm says the we need to have max a list of length 4 in stop. Forcing it.
         stop_list = data.get("stop", [])
         trimmed_stop_list = stop_list[:4]
         data["stop"] = trimmed_stop_list