loader Public

Watch 0 Fork 0 Star 0
Python · 11701 bytes Raw Blame History
1 """Focused tests for Ollama text tool parsing."""
2
3 from __future__ import annotations
4
5 import json
6 import sys
7 import types
8
9 import pytest
10
11 try:
12 import httpx as _httpx # noqa: F401
13 except ModuleNotFoundError:
14 class _AsyncClientStub:
15 def __init__(self, *args, **kwargs) -> None:
16 pass
17
18 async def aclose(self) -> None:
19 return None
20
21 sys.modules["httpx"] = types.SimpleNamespace(AsyncClient=_AsyncClientStub)
22
23 from loader.llm.base import Message, Role, StreamChunk, ToolCall
24 from loader.llm.ollama import OllamaBackend
25
26
27 class FakeResponse:
28 """Small response stub for Ollama complete() tests."""
29
30 def __init__(self, payload: dict, status_code: int = 200) -> None:
31 self._payload = payload
32 self.status_code = status_code
33 self.content = json.dumps(payload).encode()
34
35 def json(self) -> dict:
36 return self._payload
37
38 def raise_for_status(self) -> None:
39 if self.status_code >= 400:
40 raise AssertionError(f"unexpected status {self.status_code}")
41
42
43 class FakeClient:
44 """Small async client stub for Ollama complete() tests."""
45
46 def __init__(self, responses: list[FakeResponse]) -> None:
47 self.responses = list(responses)
48
49 async def post(self, url: str, json: dict) -> FakeResponse: # noqa: A002
50 assert self.responses, f"unexpected Ollama POST to {url}"
51 return self.responses.pop(0)
52
53 async def aclose(self) -> None:
54 return None
55
56
57 class FakeStreamResponse:
58 """Small streaming response stub for _stream_response tests."""
59
60 def __init__(self, payloads: list[dict]) -> None:
61 self._payloads = payloads
62
63 async def aiter_lines(self):
64 for payload in self._payloads:
65 yield json.dumps(payload)
66
67
68 def test_ollama_format_messages_adds_anchor_for_empty_tool_call_content() -> None:
69 backend = OllamaBackend()
70
71 formatted = backend._format_messages(
72 [
73 Message(
74 role=Role.ASSISTANT,
75 content="",
76 tool_calls=[
77 ToolCall(
78 id="write-1",
79 name="write",
80 arguments={"file_path": "index.html", "content": "..."},
81 )
82 ],
83 )
84 ]
85 )
86
87 assert formatted[0]["content"] == "Calling tools: write."
88 assert formatted[0]["tool_calls"] == [
89 {
90 "function": {
91 "name": "write",
92 "arguments": {"file_path": "index.html", "content": "..."},
93 }
94 }
95 ]
96
97
98 @pytest.mark.asyncio
99 async def test_ollama_complete_uses_shared_parser_with_allowed_tool_names() -> None:
100 backend = OllamaBackend()
101
102 async def fake_describe_model() -> None:
103 return None
104
105 backend.describe_model = fake_describe_model # type: ignore[method-assign]
106 backend._client = FakeClient(
107 [
108 FakeResponse(
109 {
110 "message": {
111 "content": (
112 '[calls askuserquestion tool with: '
113 'question="Which path should we take?"]'
114 )
115 },
116 "prompt_eval_count": 4,
117 "eval_count": 2,
118 }
119 )
120 ]
121 )
122
123 response = await backend.complete(
124 messages=[],
125 tools=[{"name": "AskUserQuestion"}, {"name": "TodoWrite"}],
126 )
127
128 assert response.content == ""
129 assert response.tool_calls[0].name == "AskUserQuestion"
130 assert response.tool_calls[0].arguments == {
131 "question": "Which path should we take?"
132 }
133 await backend.close()
134
135
136 @pytest.mark.asyncio
137 async def test_ollama_complete_canonicalizes_native_tool_aliases() -> None:
138 backend = OllamaBackend()
139
140 async def fake_describe_model() -> None:
141 return None
142
143 backend.describe_model = fake_describe_model # type: ignore[method-assign]
144 backend._client = FakeClient(
145 [
146 FakeResponse(
147 {
148 "message": {
149 "content": "",
150 "tool_calls": [
151 {
152 "id": "call_read",
153 "function": {
154 "name": "read_file",
155 "arguments": {"file_path": "/tmp/test.txt"},
156 },
157 }
158 ],
159 },
160 "prompt_eval_count": 4,
161 "eval_count": 2,
162 }
163 )
164 ]
165 )
166
167 response = await backend.complete(
168 messages=[],
169 tools=[{"name": "read"}, {"name": "write"}, {"name": "patch"}],
170 )
171
172 assert response.tool_calls[0].name == "read"
173 assert response.tool_calls[0].arguments == {"file_path": "/tmp/test.txt"}
174 await backend.close()
175
176
177 @pytest.mark.asyncio
178 async def test_ollama_stream_response_uses_shared_parser_for_text_tool_calls() -> None:
179 backend = OllamaBackend()
180
181 chunks = [
182 chunk
183 async for chunk in backend._stream_response(
184 FakeStreamResponse(
185 [
186 {
187 "message": {
188 "content": (
189 '[calls askuserquestion tool with: '
190 'question="Which path should we take?"]'
191 )
192 },
193 "done": False,
194 },
195 {
196 "message": {"content": ""},
197 "done": True,
198 "prompt_eval_count": 4,
199 "eval_count": 2,
200 },
201 ]
202 ),
203 tools=[{"name": "AskUserQuestion"}, {"name": "TodoWrite"}],
204 )
205 ]
206
207 final_chunk = chunks[-1]
208 assert isinstance(final_chunk, StreamChunk)
209 assert final_chunk.tool_calls[0].name == "AskUserQuestion"
210 assert final_chunk.tool_calls[0].arguments == {
211 "question": "Which path should we take?"
212 }
213 await backend.close()
214
215
216 @pytest.mark.asyncio
217 async def test_ollama_stream_response_canonicalizes_native_tool_aliases() -> None:
218 backend = OllamaBackend()
219
220 chunks = [
221 chunk
222 async for chunk in backend._stream_response(
223 FakeStreamResponse(
224 [
225 {
226 "message": {
227 "content": "",
228 "tool_calls": [
229 {
230 "id": "call_read",
231 "function": {
232 "name": "read_file",
233 "arguments": {"file_path": "/tmp/test.txt"},
234 },
235 }
236 ],
237 },
238 "done": True,
239 "prompt_eval_count": 4,
240 "eval_count": 2,
241 }
242 ]
243 ),
244 tools=[{"name": "read"}, {"name": "write"}, {"name": "patch"}],
245 )
246 ]
247
248 final_chunk = chunks[-1]
249 assert final_chunk.tool_calls[0].name == "read"
250 assert final_chunk.tool_calls[0].arguments == {"file_path": "/tmp/test.txt"}
251 await backend.close()
252
253
254 @pytest.mark.asyncio
255 async def test_ollama_stream_response_parses_fenced_read_command() -> None:
256 backend = OllamaBackend()
257
258 chunks = [
259 chunk
260 async for chunk in backend._stream_response(
261 FakeStreamResponse(
262 [
263 {
264 "message": {
265 "content": (
266 "I need to inspect the file first.\n"
267 "```bash\nread /tmp/test.txt\n```"
268 )
269 },
270 "done": False,
271 },
272 {
273 "message": {"content": ""},
274 "done": True,
275 "prompt_eval_count": 4,
276 "eval_count": 2,
277 },
278 ]
279 ),
280 tools=[{"name": "read"}, {"name": "glob"}, {"name": "bash"}],
281 )
282 ]
283
284 final_chunk = chunks[-1]
285 assert final_chunk.tool_calls[0].name == "read"
286 assert final_chunk.tool_calls[0].arguments == {"file_path": "/tmp/test.txt"}
287 await backend.close()
288
289
290 @pytest.mark.asyncio
291 async def test_ollama_stream_response_defers_raw_json_detection_to_final_parse() -> None:
292 backend = OllamaBackend()
293 raw_json = (
294 '{"name": "AskUserQuestion", '
295 '"arguments": {"question": "Which path should we take?"}}'
296 )
297
298 chunks = [
299 chunk
300 async for chunk in backend._stream_response(
301 FakeStreamResponse(
302 [
303 {
304 "message": {"content": raw_json[:30]},
305 "done": False,
306 },
307 {
308 "message": {"content": raw_json[30:]},
309 "done": False,
310 },
311 {
312 "message": {"content": ""},
313 "done": True,
314 "prompt_eval_count": 4,
315 "eval_count": 2,
316 },
317 ]
318 ),
319 tools=[{"name": "AskUserQuestion"}, {"name": "TodoWrite"}],
320 )
321 ]
322
323 assert [chunk.pending_tool_call for chunk in chunks[:-1]] == [None, None]
324 assert [chunk.content for chunk in chunks[:-1]] == [raw_json[:30], raw_json[30:]]
325 final_chunk = chunks[-1]
326 assert final_chunk.full_content == ""
327 assert final_chunk.tool_calls[0].name == "AskUserQuestion"
328 assert final_chunk.tool_calls[0].arguments == {
329 "question": "Which path should we take?"
330 }
331 await backend.close()
332
333
334 @pytest.mark.asyncio
335 async def test_ollama_stream_response_filters_tool_call_tags_and_parses_at_end() -> None:
336 backend = OllamaBackend()
337
338 chunks = [
339 chunk
340 async for chunk in backend._stream_response(
341 FakeStreamResponse(
342 [
343 {
344 "message": {
345 "content": 'Before <tool_call>{"name": "AskUserQuestion", '
346 '"arguments": {"question": "Which path should we take?"}}'
347 },
348 "done": False,
349 },
350 {
351 "message": {"content": "</tool_call> After"},
352 "done": False,
353 },
354 {
355 "message": {"content": ""},
356 "done": True,
357 "prompt_eval_count": 4,
358 "eval_count": 2,
359 },
360 ]
361 ),
362 tools=[{"name": "AskUserQuestion"}, {"name": "TodoWrite"}],
363 )
364 ]
365
366 assert [chunk.pending_tool_call for chunk in chunks[:-1]] == [None, None]
367 assert [chunk.content for chunk in chunks[:-1]] == ["Before ", " After"]
368 final_chunk = chunks[-1]
369 assert final_chunk.full_content == "Before After"
370 assert final_chunk.tool_calls[0].name == "AskUserQuestion"
371 assert final_chunk.tool_calls[0].arguments == {
372 "question": "Which path should we take?"
373 }
374 await backend.close()