Fixed unit tests

Author: aponcedeleonch

scripts/import_packages.py

@@ -3,7 +3,6 @@
 import os
 import shutil
 
-
 import weaviate
 from weaviate.classes.config import DataType, Property
 from weaviate.embedded import EmbeddedOptions

src/codegate/llm_utils/llmclient.py

@@ -132,7 +132,7 @@ async def _complete_litellm(
                     temperature=request["temperature"],
                     base_url=base_url,
                     response_format=request["response_format"],
-            )
+                )
             content = response["choices"][0]["message"]["content"]
 
             # Clean up code blocks if present

src/codegate/providers/ollama/adapter.py

@@ -19,6 +19,7 @@ def normalize(self, data: Dict) -> ChatCompletionRequest:
         """
         # Make a copy of the data to avoid modifying the original and normalize the message content
         normalized_data = self._normalize_content_messages(data)
+        normalized_data["model"] = data.get("model", "").strip()
         normalized_data["options"] = data.get("options", {})
         # In Ollama force the stream to be True. Continue is not setting this parameter and
         # most of our functionality is for streaming completions.

src/codegate/providers/ollama/completion_handler.py

@@ -45,13 +45,16 @@ async def execute_completion(
     ) -> Union[ChatResponse, GenerateResponse]:
         """Stream response directly from Ollama API."""
         if is_fim_request:
-            prompt = request["messages"][0].content
-            response = self.client.generate(model=request["model"], prompt=prompt, stream=stream)
+            prompt = request["messages"][0]["content"]
+            response = self.client.generate(
+                model=request["model"], prompt=prompt, stream=stream, options=request["options"]
+            )
         else:
             response = self.client.chat(
                 model=request["model"],
                 messages=request["messages"],
                 stream=stream,
+                options=request["options"],
             )
         return response
 

src/codegate/storage/storage_engine.py

@@ -67,7 +67,6 @@ def __del__(self):
         except Exception as e:
             logger.error(f"Failed to close client: {str(e)}")
 
-
     def get_client(self, data_path):
         try:
             # Configure Weaviate logging

tests/providers/ollama/test_ollama_adapter.py

@@ -18,21 +18,6 @@ def test_normalize_ollama_input():
     normalized = normalizer.normalize(data)
     assert normalized["model"] == "codellama:7b-instruct"  # Space removed
 
-    # Test base URL handling
-    data = {"model": "llama2", "base_url": "http://localhost:11434"}
-    normalized = normalizer.normalize(data)
-    assert normalized["base_url"] == "http://localhost:11434/api"
-
-    # Test base URL already has /api
-    data = {"model": "llama2", "base_url": "http://localhost:11434/api"}
-    normalized = normalizer.normalize(data)
-    assert normalized["base_url"] == "http://localhost:11434/api"
-
-    # Test base URL with trailing slash
-    data = {"model": "llama2", "base_url": "http://localhost:11434/"}
-    normalized = normalizer.normalize(data)
-    assert normalized["base_url"] == "http://localhost:11434/api"
-
 
 def test_normalize_native_ollama_input():
     """Test input normalization for native Ollama API requests."""
@@ -58,7 +43,6 @@ def test_normalize_native_ollama_input():
         "base_url": "http://localhost:11434",
     }
     normalized = normalizer.normalize(data)
-    assert normalized["base_url"] == "http://localhost:11434/api"
 
 
 def test_normalize_ollama_message_format():
@@ -128,11 +112,6 @@ def test_normalize_ollama_output():
     """Test output normalization for Ollama."""
     normalizer = OllamaOutputNormalizer()
 
-    # Test streaming response passthrough
-    response = {"message": {"role": "assistant", "content": "test"}}
-    normalized = normalizer.normalize_streaming(response)
-    assert normalized == response
-
     # Test regular response passthrough
     response = {"message": {"role": "assistant", "content": "test"}}
     normalized = normalizer.normalize(response)
@@ -147,8 +126,3 @@ def test_normalize_ollama_output():
     response = {"message": {"role": "assistant", "content": "test"}}
     denormalized = normalizer.denormalize(response)
     assert denormalized == response
-
-    # Test streaming denormalize passthrough
-    response = {"message": {"role": "assistant", "content": "test"}}
-    denormalized = normalizer.denormalize_streaming(response)
-    assert denormalized == response

tests/providers/ollama/test_ollama_completion_handler.py

@@ -0,0 +1,54 @@
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from litellm import ChatCompletionRequest
+from ollama import ChatResponse, GenerateResponse, Message
+
+from codegate.providers.ollama.completion_handler import OllamaShim
+
+
+@pytest.fixture
+def mock_client():
+    client = MagicMock()
+    client.generate = AsyncMock(return_value=GenerateResponse(response="FIM response"))
+    client.chat = AsyncMock(
+        return_value=ChatResponse(message=Message(content="Chat response", role="assistant"))
+    )
+    return client
+
+
+@pytest.fixture
+def handler(mock_client):
+    ollama_shim = OllamaShim()
+    ollama_shim.client = mock_client
+    return ollama_shim
+
+
+@pytest.fixture
+def chat_request():
+    return ChatCompletionRequest(
+        model="test-model", messages=[{"role": "user", "content": "Hello"}], options={}
+    )
+
+
+@pytest.mark.asyncio
+async def test_execute_completion_is_fim_request(handler, chat_request):
+    chat_request["messages"][0]["content"] = "FIM prompt"
+    await handler.execute_completion(chat_request, api_key=None, stream=False, is_fim_request=True)
+    handler.client.generate.assert_called_once_with(
+        model=chat_request["model"],
+        prompt="FIM prompt",
+        stream=False,
+        options=chat_request["options"],
+    )
+
+
+@pytest.mark.asyncio
+async def test_execute_completion_not_is_fim_request(handler, chat_request):
+    await handler.execute_completion(chat_request, api_key=None, stream=False, is_fim_request=False)
+    handler.client.chat.assert_called_once_with(
+        model=chat_request["model"],
+        messages=chat_request["messages"],
+        stream=False,
+        options=chat_request["options"],
+    )