diff --git a/src/codegate/providers/ollama/completion_handler.py b/src/codegate/providers/ollama/completion_handler.py index cfd04acd..9e602dc9 100644 --- a/src/codegate/providers/ollama/completion_handler.py +++ b/src/codegate/providers/ollama/completion_handler.py @@ -89,9 +89,21 @@ async def execute_completion( """Stream response directly from Ollama API.""" self.base_tool = base_tool if is_fim_request: - prompt = request["messages"][0].get("content", "") + prompt = "" + for i in reversed(range(len(request["messages"]))): + if request["messages"][i]["role"] == "user": + prompt = request["messages"][i]["content"] # type: ignore + break + if not prompt: + raise ValueError("No user message found in FIM request") + response = await self.client.generate( - model=request["model"], prompt=prompt, stream=stream, options=request["options"] # type: ignore + model=request["model"], + prompt=prompt, + raw=request.get("raw", False), + suffix=request.get("suffix", ""), + stream=stream, + options=request["options"], # type: ignore ) else: response = await self.client.chat( diff --git a/tests/integration/integration_tests.py b/tests/integration/integration_tests.py index ebed95ed..c9930353 100644 --- a/tests/integration/integration_tests.py +++ b/tests/integration/integration_tests.py @@ -139,6 +139,7 @@ async def run_test(self, test: dict, test_headers: dict) -> bool: try: parsed_response = self.parse_response_message(response, streaming=streaming) + logger.debug(f"Response message: {parsed_response}") # Load appropriate checks for this test checks = CheckLoader.load(test) diff --git a/tests/integration/testcases.yaml b/tests/integration/testcases.yaml index 9190dc9c..9109088c 100644 --- a/tests/integration/testcases.yaml +++ b/tests/integration/testcases.yaml @@ -297,48 +297,32 @@ testcases: url: http://127.0.0.1:8989/ollama/api/generate data: | { - "model": "qwen2.5-coder:0.5b", - "max_tokens": 4096, - "temperature": 0, "stream": true, - "stop": [ - "", - "", - "", - "", - "", - "", - "/src/", - "#- coding: utf-8", - "```", - "" - ], + "model": "qwen2.5-coder:0.5b", + "raw": true, + "options": { + "temperature": 0.01, + "num_predict": 4096, + "stop": [ + "<|endoftext|>", + "<|fim_prefix|>", + "<|fim_middle|>", + "<|fim_suffix|>", + "<|fim_pad|>", + "<|repo_name|>", + "<|file_sep|>", + "<|im_start|>", + "<|im_end|>", + "/src/", + "#- coding: utf-8", + "```" + ], + "num_ctx": 8096 + }, "prompt":"<|fim_prefix|>\n# codegate/test.py\nimport invokehttp\nimport requests\n\nkey = \"mysecret-key\"\n\ndef call_api():\n <|fim_suffix|>\n\n\ndata = {'key1': 'test1', 'key2': 'test2'}\nresponse = call_api('http://localhost:8080', method='post', data='data')\n<|fim_middle|>" } likes: | ```python - import invokehttp - import requests - - key = "mysecret-key" - - def call_api(url, method='get', data=None): - headers = { - 'Authorization': f'Bearer {key}' - } - - if method == 'get': - response = requests.get(url, headers=headers) - elif method == 'post': - response = requests.post(url, headers=headers, json=data) - else: - raise ValueError("Unsupported HTTP method") - - return response - - data = {'key1': 'test1', 'key2': 'test2'} - response = call_api('http://localhost:8080', method='post', data=data) - - print(response.status_code) - print(response.json()) + if __name__ == '__main__': + invokehttp.run(call_api) ``` diff --git a/tests/providers/ollama/test_ollama_completion_handler.py b/tests/providers/ollama/test_ollama_completion_handler.py index cc32e915..df0eb149 100644 --- a/tests/providers/ollama/test_ollama_completion_handler.py +++ b/tests/providers/ollama/test_ollama_completion_handler.py @@ -40,6 +40,8 @@ async def test_execute_completion_is_fim_request(handler, chat_request): prompt="FIM prompt", stream=False, options=chat_request["options"], + suffix="", + raw=False, )