Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions tests/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
"""TinyAgent test suite."""
1 change: 1 addition & 0 deletions tests/fixtures/ollama_generate.json
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
{"model":"llama3.2:latest","created_at":"2026-10-10T03:55:03.747565489Z","response":"OK","done":true,"done_reason":"stop","context":[128006,9125,128007,271,38766,1303,33025,2696,25,6790,220,2366,18,271,128009,128006,882,128007,271,21509,449,832,3492,25,10619,128009,128006,78191,128007,271,4012],"total_duration":2482083673,"load_duration":2385723400,"prompt_eval_count":31,"prompt_eval_cached_count":0,"prompt_eval_duration":58650000,"eval_count":2,"eval_duration":35593000}
76 changes: 76 additions & 0 deletions tests/stub_server.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
"""Threaded local HTTP server for backend contract tests."""

from __future__ import annotations

import json
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from threading import Lock, Thread
from time import sleep


class _Handler(BaseHTTPRequestHandler):
def do_POST(self) -> None:
length = int(self.headers.get("Content-Length", "0"))
raw_body = self.rfile.read(length)
try:
decoded = json.loads(raw_body.decode("utf-8"))
except (UnicodeError, ValueError):
decoded = None

state = self.server.state # type: ignore[attr-defined]
with state.lock:
state.requests.append(
{
"path": self.path,
"headers": dict(self.headers.items()),
"body": raw_body,
"json": decoded,
}
)
status, body, delay_s = state.response
if delay_s:
sleep(delay_s)
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
try:
self.wfile.write(body)
except OSError:
pass

def log_message(self, format: str, *args: object) -> None:
return


class StubServer:
def __init__(self) -> None:
self.lock = Lock()
self.requests: list[dict] = []
self.response: tuple[int, bytes, float] = (200, b"{}", 0.0)
self.server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
self.server.state = self # type: ignore[attr-defined]
self.thread = Thread(target=self.server.serve_forever, daemon=True)
self.thread.start()

@property
def url(self) -> str:
host, port = self.server.server_address
return f"http://{host}:{port}"

def set_response(
self, status: int = 200, body: bytes | str | dict | list = b"{}", delay_s: float = 0
) -> None:
if isinstance(body, (dict, list)):
encoded = json.dumps(body).encode("utf-8")
elif isinstance(body, str):
encoded = body.encode("utf-8")
else:
encoded = body
with self.lock:
self.response = (status, encoded, delay_s)

def close(self) -> None:
self.server.shutdown()
self.server.server_close()
self.thread.join(timeout=1)
200 changes: 200 additions & 0 deletions tests/test_backends.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,200 @@
"""PRD-02 transport tests against a local standard-library HTTP stub."""

from __future__ import annotations

import json
import socket
import unittest
import warnings

from tests.stub_server import StubServer
from tinyagent.backends import LlamaCppBackend, OllamaBackend
from tinyagent.errors import BackendError


DEFAULT_STOP = ["\n\n", "\nQuery:", "\nReply:"]


class BackendTest(unittest.TestCase):
def setUp(self) -> None:
self.server = StubServer()
self.addCleanup(self.server.close)

def _ollama(self, *, timeout_s: float = 60.0, num_ctx: int | None = None):
return OllamaBackend(
"test-model", host=self.server.url, timeout_s=timeout_s, num_ctx=num_ctx
)

def _llamacpp(self, *, timeout_s: float = 60.0):
return LlamaCppBackend(base_url=self.server.url, timeout_s=timeout_s)

def test_ollama_request_contract(self) -> None:
self.server.set_response(200, {"response": " raw\ntext "})
result = self._ollama().complete("prompt", stop=DEFAULT_STOP, max_tokens=96)
request = self.server.requests[0]
self.assertEqual(request["path"], "/api/generate")
headers = {name.lower(): value for name, value in request["headers"].items()}
self.assertEqual(headers["content-type"], "application/json")
self.assertEqual(
request["json"],
{
"model": "test-model",
"prompt": "prompt",
"stream": False,
"raw": False,
"options": {
"temperature": 0.0,
"top_k": 1,
"top_p": 1.0,
"seed": 42,
"num_predict": 96,
"stop": DEFAULT_STOP,
},
},
)
self.assertEqual(result.text, " raw\ntext ")
self.assertEqual(result.backend, "ollama:test-model")
self.assertEqual(len(self.server.requests), 1)

def test_ollama_optional_context_and_token_counts(self) -> None:
self.server.set_response(
200,
{"response": "ok", "prompt_eval_count": 5, "eval_count": 2},
)
result = self._ollama(num_ctx=2048).complete(
"prompt", stop=(), max_tokens=3
)
self.assertEqual(self.server.requests[0]["json"]["options"]["num_ctx"], 2048)
self.assertEqual((result.prompt_tokens, result.completion_tokens), (5, 2))

def test_ollama_omits_context_and_invalid_token_counts(self) -> None:
self.server.set_response(
200,
{"response": "ok", "prompt_eval_count": "5", "eval_count": True},
)
result = self._ollama().complete("prompt", stop=(), max_tokens=1)
self.assertNotIn("num_ctx", self.server.requests[0]["json"]["options"])
self.assertEqual((result.prompt_tokens, result.completion_tokens), (None, None))

def test_llamacpp_request_contract(self) -> None:
self.server.set_response(
200,
{"content": "unchanged", "tokens_evaluated": 7, "tokens_predicted": 4},
)
result = self._llamacpp().complete("prompt", stop=DEFAULT_STOP, max_tokens=96)
request = self.server.requests[0]
self.assertEqual(request["path"], "/completion")
self.assertEqual(
request["json"],
{
"prompt": "prompt",
"stream": False,
"temperature": 0.0,
"top_k": 1,
"top_p": 1.0,
"seed": 42,
"n_predict": 96,
"stop": DEFAULT_STOP,
"cache_prompt": True,
},
)
self.assertEqual((result.text, result.backend), ("unchanged", "llamacpp"))
self.assertEqual((result.prompt_tokens, result.completion_tokens), (7, 4))

def test_constructor_and_call_validation(self) -> None:
for model in ("", " ", None):
with self.subTest(model=model), self.assertRaises(ValueError):
OllamaBackend(model) # type: ignore[arg-type]
backend = self._ollama()
for max_tokens in (0, 513, 1.5, True):
with self.subTest(max_tokens=max_tokens), self.assertRaises(ValueError):
backend.complete("prompt", stop=(), max_tokens=max_tokens)
with self.assertRaises(ValueError):
backend.complete("prompt", stop=None, max_tokens=96) # type: ignore[arg-type]
self.assertEqual(self.server.requests, [])

def test_sampling_has_no_public_override(self) -> None:
self.assertNotIn("temperature", OllamaBackend.__init__.__annotations__)
self.assertNotIn("temperature", LlamaCppBackend.__init__.__annotations__)
self.assertNotIn("temperature", OllamaBackend.complete.__annotations__)
self.assertNotIn("temperature", LlamaCppBackend.complete.__annotations__)

def test_unreachable_is_typed_and_uses_one_attempt(self) -> None:
probe = socket.socket()
probe.bind(("127.0.0.1", 0))
port = probe.getsockname()[1]
probe.close()
backend = OllamaBackend("test-model", host=f"http://127.0.0.1:{port}")
with self.assertRaises(BackendError) as raised:
backend.complete("secret prompt", stop=(), max_tokens=1)
self.assertEqual(raised.exception.code, "E_BACKEND_UNREACHABLE")
self.assertNotIn("secret prompt", raised.exception.detail)
self.assertLessEqual(len(raised.exception.detail), 200)

def test_timeout_is_typed(self) -> None:
self.server.set_response(200, {"response": "late"}, delay_s=0.4)
with self.assertRaises(BackendError) as raised:
self._ollama(timeout_s=0.05).complete(
"prompt", stop=(), max_tokens=1
)
self.assertEqual(raised.exception.code, "E_BACKEND_TIMEOUT")
self.assertIn("0.05s", raised.exception.detail)
self.assertEqual(len(self.server.requests), 1)

def test_http_error_precedes_url_error_mapping(self) -> None:
for count, status in enumerate((404, 500), 1):
with self.subTest(status=status):
self.server.set_response(status, "server says no")
with self.assertRaises(BackendError) as raised:
self._ollama().complete("secret prompt", stop=(), max_tokens=1)
self.assertEqual(raised.exception.code, "E_BACKEND_HTTP")
self.assertIn(f"HTTP {status}", raised.exception.detail)
self.assertNotIn("secret prompt", raised.exception.detail)
self.assertEqual(len(self.server.requests), count)

def test_ollama_error_field(self) -> None:
self.server.set_response(
200, {"error": "model not found: secret prompt " + "x" * 250}
)
with self.assertRaises(BackendError) as raised:
self._ollama().complete("secret prompt", stop=(), max_tokens=1)
self.assertEqual(raised.exception.code, "E_BACKEND_HTTP")
self.assertNotIn("secret prompt", raised.exception.detail)
self.assertLessEqual(len(raised.exception.detail), 200)

def test_protocol_errors(self) -> None:
cases = (
(b"not json", self._ollama()),
(b"[]", self._ollama()),
(b'{"response": 2}', self._ollama()),
(b'{"content": null}', self._llamacpp()),
)
for body, backend in cases:
with self.subTest(body=body, backend=backend):
self.server.set_response(200, body)
with self.assertRaises(BackendError) as raised:
backend.complete("prompt", stop=(), max_tokens=1)
self.assertEqual(raised.exception.code, "E_BACKEND_PROTOCOL")

def test_fifty_malformed_responses_only_raise_backend_error(self) -> None:
backend = self._ollama()
for index in range(50):
body = (b"\xff" * (index + 1)) if index % 2 else json.dumps(
{"wrong": index}
).encode("utf-8")
self.server.set_response(200, body)
with self.subTest(index=index), self.assertRaises(BackendError):
backend.complete("prompt", stop=(), max_tokens=1)
self.assertEqual(len(self.server.requests), 50)

def test_remote_host_warns_and_repr_does_not_hold_prompt(self) -> None:
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
backend = OllamaBackend("test-model", host="http://example.invalid:11434")
self.assertEqual(len(caught), 1)
self.assertIn("non-loopback", str(caught[0].message))
self.assertNotIn("prompt", repr(backend).lower())


if __name__ == "__main__":
unittest.main()
36 changes: 15 additions & 21 deletions tinyagent/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,18 +44,22 @@ def run(self, query: str) -> AgentResult:
current = original
outputs: list[str] = []
last_tool_result: ToolResult | None = None
attempt = 0

def finish(status, text, provenance, error_code=None) -> AgentResult:
# Reads the loop state (attempt, outputs, last_tool_result) at call time.
return AgentResult(
status, text, provenance, last_tool_result,
attempt + 1, tuple(outputs), error_code,
)

for attempt in range(1 + self.max_repairs):
calls = attempt + 1
try:
completion = self.backend.complete(
current, stop=STOP_SEQUENCES, max_tokens=self.max_tokens
)
except BackendError as error:
return AgentResult(
"error", f"{error.code}: {error.detail}", "none",
last_tool_result, calls, tuple(outputs), error.code,
)
return finish("error", f"{error.code}: {error.detail}", "none", error.code)

outputs.append(completion.text)
outcome = parse(completion.text)
Expand All @@ -65,33 +69,23 @@ def run(self, query: str) -> AgentResult:
)
last_tool_result = result
if result.ok:
return AgentResult(
"tool", result.output or "", "tool", result,
calls, tuple(outputs), None,
)
return finish("tool", result.output or "", "tool")
code = result.error_code or "E_TOOL_EXCEPTION"
if code in EXEC_TERMINAL:
return AgentResult(
"error", f"{code}: {result.error_detail or ''}".rstrip(": "),
"none", result, calls, tuple(outputs), code,
)
detail = result.error_detail or ""
if code in EXEC_TERMINAL:
return finish("error", f"{code}: {detail}".rstrip(": "), "none", code)
elif isinstance(outcome, ParsedFinal):
if not self.require_tool:
return AgentResult(
"final", outcome.text, "model", last_tool_result,
calls, tuple(outputs), None,
)
return finish("final", outcome.text, "model")
code, detail = "E_FINAL_NOT_ALLOWED", ""
elif isinstance(outcome, ParseFailure):
code, detail = outcome.code, ""
else:
raise TypeError("parser returned an invalid outcome")

if attempt == self.max_repairs:
return AgentResult(
"error", f"retry exhausted after {code}", "none",
last_tool_result, calls, tuple(outputs), "E_RETRY_EXHAUSTED",
return finish(
"error", f"retry exhausted after {code}", "none", "E_RETRY_EXHAUSTED"
)
current = render_repair_prompt(
original, completion.text, code, detail
Expand Down
6 changes: 6 additions & 0 deletions tinyagent/backends/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
"""Supported native local inference backends."""

from tinyagent.backends.llamacpp import LlamaCppBackend
from tinyagent.backends.ollama import OllamaBackend

__all__ = ["OllamaBackend", "LlamaCppBackend"]
Loading
Loading