Merge pull request #2 from LLM-project-Team/dynamo/transport-refactor

Dynamo/transport refactor

Author: DYNAMOatGHUB

tests/__init__.py

@@ -0,0 +1 @@
+"""TinyAgent test suite."""

tests/fixtures/ollama_generate.json

@@ -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}

tests/stub_server.py

@@ -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)

tests/test_backends.py

@@ -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()

tinyagent/agent.py

@@ -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)
@@ -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

tinyagent/backends/__init__.py

@@ -0,0 +1,6 @@
+"""Supported native local inference backends."""
+
+from tinyagent.backends.llamacpp import LlamaCppBackend
+from tinyagent.backends.ollama import OllamaBackend
+
+__all__ = ["OllamaBackend", "LlamaCppBackend"]