wehub-resource-sync
66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from opensquilla.provider.ollama import OllamaProvider
|
|
from opensquilla.provider.types import ChatConfig, DoneEvent, Message
|
|
|
|
|
|
def test_ollama_provider_writes_llm_trace(monkeypatch: Any, tmp_path: Any) -> None:
|
|
trace_path = tmp_path / "ollama-llm-calls.jsonl"
|
|
monkeypatch.setenv("OPENSQUILLA_LLM_TRACE_RECORDER", "full")
|
|
monkeypatch.setenv("OPENSQUILLA_LLM_TRACE_PATH", str(trace_path))
|
|
body = b"".join(
|
|
[
|
|
json.dumps({"message": {"content": "ok"}, "done": False}).encode() + b"\n",
|
|
json.dumps(
|
|
{"message": {}, "done": True, "prompt_eval_count": 5, "eval_count": 2}
|
|
).encode()
|
|
+ b"\n",
|
|
]
|
|
)
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
headers={"content-type": "application/x-ndjson"},
|
|
content=body,
|
|
)
|
|
|
|
transport = httpx.MockTransport(handler)
|
|
real_async_client = httpx.AsyncClient
|
|
|
|
def patched_async_client(*args: Any, **kwargs: Any) -> httpx.AsyncClient:
|
|
kwargs["transport"] = transport
|
|
return real_async_client(*args, **kwargs)
|
|
|
|
monkeypatch.setattr("opensquilla.provider.ollama.httpx.AsyncClient", patched_async_client)
|
|
provider = OllamaProvider(model="llama3")
|
|
|
|
async def _run() -> list[Any]:
|
|
return [
|
|
event
|
|
async for event in provider.chat(
|
|
[Message(role="user", content="hi")],
|
|
config=ChatConfig(max_tokens=12),
|
|
)
|
|
]
|
|
|
|
events = asyncio.run(_run())
|
|
|
|
assert any(isinstance(event, DoneEvent) for event in events)
|
|
rows = [json.loads(line) for line in trace_path.read_text(encoding="utf-8").splitlines()]
|
|
assert [row["event"] for row in rows] == [
|
|
"llm.request",
|
|
"llm.response_chunk",
|
|
"llm.response_chunk",
|
|
"llm.response",
|
|
]
|
|
assert rows[0]["provider"] == "ollama"
|
|
assert rows[-1]["assistant_text"] == "ok"
|
|
assert rows[-1]["usage"]["input_tokens"] == 5
|