diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..19be305 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""TinyAgent test suite.""" diff --git a/tests/fixtures/ollama_generate.json b/tests/fixtures/ollama_generate.json new file mode 100644 index 0000000..fd7ed85 --- /dev/null +++ b/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} \ No newline at end of file diff --git a/tests/stub_server.py b/tests/stub_server.py new file mode 100644 index 0000000..535706b --- /dev/null +++ b/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) diff --git a/tests/test_backends.py b/tests/test_backends.py new file mode 100644 index 0000000..c70c821 --- /dev/null +++ b/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() diff --git a/tinyagent/agent.py b/tinyagent/agent.py index bc30a61..371e499 100644 --- a/tinyagent/agent.py +++ b/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,23 +69,14 @@ 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, "" @@ -89,9 +84,8 @@ def run(self, query: str) -> AgentResult: 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 diff --git a/tinyagent/backends/__init__.py b/tinyagent/backends/__init__.py new file mode 100644 index 0000000..b844ead --- /dev/null +++ b/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"] diff --git a/tinyagent/backends/base.py b/tinyagent/backends/base.py new file mode 100644 index 0000000..bec6823 --- /dev/null +++ b/tinyagent/backends/base.py @@ -0,0 +1,113 @@ +"""Shared standard-library HTTP transport for local inference backends.""" + +from __future__ import annotations + +import json +import socket +import urllib.error +import urllib.request + +from tinyagent.errors import BackendError + + +def _single_line(value: object, *, sensitive_text: str | None = None) -> str: + text = str(value) + if sensitive_text: + text = text.replace(sensitive_text, "[redacted]") + return " ".join(text.split())[:200] + + +def _is_loopback_url(url: str) -> bool: + authority = url.split("://", 1)[-1].split("/", 1)[0] + authority = authority.split("?", 1)[0].split("#", 1)[0].rsplit("@", 1)[-1] + if authority.startswith("["): + host = authority[1:].split("]", 1)[0] + elif authority.count(":") == 1: + host = authority.rsplit(":", 1)[0] + else: + host = authority + + if host.lower() == "localhost": + return True + try: + return socket.inet_pton(socket.AF_INET, host)[0] == 127 + except OSError: + pass + try: + return socket.inet_pton(socket.AF_INET6, host) == (b"\x00" * 15 + b"\x01") + except OSError: + return False + + +def _http_error_body(error: urllib.error.HTTPError) -> str: + try: + body = error.read(200) + except Exception: + return "" + return body.decode("utf-8", errors="replace") + + +def _post_json( + url: str, + payload: dict, + timeout_s: float, + *, + sensitive_text: str | None = None, +) -> dict: + """POST one UTF-8 JSON request and return a decoded JSON object.""" + try: + body = json.dumps(payload, ensure_ascii=False).encode("utf-8") + request = urllib.request.Request( + url, + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(request, timeout=timeout_s) as response: + raw_response = response.read() + except urllib.error.HTTPError as exc: + detail = f"HTTP {exc.code}: {_http_error_body(exc)}" + raise BackendError( + "E_BACKEND_HTTP", _single_line(detail, sensitive_text=sensitive_text) + ) from exc + except Exception as exc: + # URLError wraps the underlying socket failure in .reason. + reason = exc.reason if isinstance(exc, urllib.error.URLError) else exc + if isinstance(reason, (socket.timeout, TimeoutError)): + code, detail = "E_BACKEND_TIMEOUT", f"no response within {timeout_s}s" + else: + code, detail = "E_BACKEND_UNREACHABLE", f"cannot reach {url}: {reason}" + raise BackendError( + code, _single_line(detail, sensitive_text=sensitive_text) + ) from exc + + try: + result = json.loads(raw_response.decode("utf-8")) + except (UnicodeError, ValueError) as exc: + raise _protocol_error("invalid JSON") from exc + if not isinstance(result, dict): + raise _protocol_error("expected an object") + return result + + +def _validate_call(prompt: str, stop: object, max_tokens: int) -> list[str]: + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") + if type(max_tokens) is not int or not 1 <= max_tokens <= 512: + raise ValueError("max_tokens must be an int in 1..512") + try: + return [str(value) for value in stop] # type: ignore[union-attr] + except TypeError as exc: + raise ValueError("stop must be a sequence of strings") from exc + + +def _token_count(value: object) -> int | None: + return value if type(value) is int else None + + +def _protocol_error(message: str) -> BackendError: + error = BackendError( + "E_BACKEND_PROTOCOL", _single_line(f"unexpected response shape: {message}") + ) + error.__cause__ = ValueError(message) + return error diff --git a/tinyagent/backends/llamacpp.py b/tinyagent/backends/llamacpp.py new file mode 100644 index 0000000..0eb0c76 --- /dev/null +++ b/tinyagent/backends/llamacpp.py @@ -0,0 +1,76 @@ +"""llama.cpp's non-streaming completion endpoint.""" + +from __future__ import annotations + +import time +import warnings +from typing import Sequence + +from tinyagent.backends.base import ( + _is_loopback_url, + _post_json, + _protocol_error, + _token_count, + _validate_call, +) +from tinyagent.types import Completion + + +class LlamaCppBackend: + def __init__( + self, + *, + base_url: str = "http://127.0.0.1:8080", + timeout_s: float = 60.0, + seed: int = 42, + ) -> None: + if not isinstance(base_url, str) or not base_url.strip(): + raise ValueError("base_url must be a non-empty string") + self.base_url = base_url.rstrip("/") + self.timeout_s = timeout_s + self.seed = seed + self.name = "llamacpp" + if not _is_loopback_url(self.base_url): + warnings.warn( + "llama.cpp host is non-loopback; prompts may leave this machine.", + RuntimeWarning, + stacklevel=2, + ) + + def complete( + self, prompt: str, *, stop: Sequence[str], max_tokens: int + ) -> Completion: + stop_list = _validate_call(prompt, stop, max_tokens) + payload = { + "prompt": prompt, + "stream": False, + "temperature": 0.0, + "top_k": 1, + "top_p": 1.0, + "seed": self.seed, + "n_predict": max_tokens, + "stop": stop_list, + "cache_prompt": True, + } + started = time.perf_counter() + response = _post_json( + f"{self.base_url}/completion", + payload, + self.timeout_s, + sensitive_text=prompt, + ) + latency_ms = (time.perf_counter() - started) * 1000 + + text = response.get("content") + if not isinstance(text, str): + raise _protocol_error("missing or non-string content field") + return Completion( + text=text, + latency_ms=latency_ms, + prompt_tokens=_token_count(response.get("tokens_evaluated")), + completion_tokens=_token_count(response.get("tokens_predicted")), + backend=self.name, + ) + + def __repr__(self) -> str: + return f"LlamaCppBackend(base_url={self.base_url!r}, timeout_s={self.timeout_s!r})" diff --git a/tinyagent/backends/ollama.py b/tinyagent/backends/ollama.py new file mode 100644 index 0000000..9542061 --- /dev/null +++ b/tinyagent/backends/ollama.py @@ -0,0 +1,99 @@ +"""Ollama's non-streaming local generate endpoint.""" + +from __future__ import annotations + +import time +import warnings +from typing import Sequence + +from tinyagent.backends.base import ( + _is_loopback_url, + _post_json, + _protocol_error, + _single_line, + _token_count, + _validate_call, +) +from tinyagent.errors import BackendError +from tinyagent.types import Completion + + +class OllamaBackend: + def __init__( + self, + model: str, + *, + host: str = "http://127.0.0.1:11434", + timeout_s: float = 60.0, + seed: int = 42, + num_ctx: int | None = None, + ) -> None: + if not isinstance(model, str) or not model.strip(): + raise ValueError("model must be a non-empty string") + if not isinstance(host, str) or not host.strip(): + raise ValueError("host must be a non-empty string") + self.model = model + self.host = host.rstrip("/") + self.timeout_s = timeout_s + self.seed = seed + self.num_ctx = num_ctx + self.name = f"ollama:{model}" + if not _is_loopback_url(self.host): + warnings.warn( + "Ollama host is non-loopback; prompts may leave this machine.", + RuntimeWarning, + stacklevel=2, + ) + + def complete( + self, prompt: str, *, stop: Sequence[str], max_tokens: int + ) -> Completion: + stop_list = _validate_call(prompt, stop, max_tokens) + options = { + "temperature": 0.0, + "top_k": 1, + "top_p": 1.0, + "seed": self.seed, + "num_predict": max_tokens, + "stop": stop_list, + } + if self.num_ctx is not None: + options["num_ctx"] = self.num_ctx + payload = { + "model": self.model, + "prompt": prompt, + "stream": False, + "raw": False, + "options": options, + } + started = time.perf_counter() + response = _post_json( + f"{self.host}/api/generate", + payload, + self.timeout_s, + sensitive_text=prompt, + ) + latency_ms = (time.perf_counter() - started) * 1000 + + server_error = response.get("error") + if isinstance(server_error, str) and server_error.strip(): + raise BackendError( + "E_BACKEND_HTTP", _single_line(server_error, sensitive_text=prompt) + ) from ValueError("Ollama reported an error") + + text = response.get("response") + if not isinstance(text, str): + raise _protocol_error("missing or non-string response field") + return Completion( + text=text, + latency_ms=latency_ms, + prompt_tokens=_token_count(response.get("prompt_eval_count")), + completion_tokens=_token_count(response.get("eval_count")), + backend=self.name, + ) + + def __repr__(self) -> str: + return ( + f"OllamaBackend(model={self.model!r}, host={self.host!r}, " + f"timeout_s={self.timeout_s!r})" + ) diff --git a/tinyagent/errors.py b/tinyagent/errors.py index 291556e..47209da 100644 --- a/tinyagent/errors.py +++ b/tinyagent/errors.py @@ -4,14 +4,17 @@ BACKEND_CODES = frozenset({"E_BACKEND_UNREACHABLE", "E_BACKEND_TIMEOUT", "E_BACKEND_HTTP", "E_BACKEND_PROTOCOL"}) AGENT_CODES = frozenset({"E_FINAL_NOT_ALLOWED", "E_RETRY_EXHAUSTED"}) + class TinyAgentError(Exception): """Base class. Only raised for programmer errors and backend failures, never for model output.""" + class BackendError(TinyAgentError): def __init__(self, code: str, detail: str) -> None: super().__init__(f"{code}: {detail}") self.code = code self.detail = detail + class CoercionError(TinyAgentError, ValueError): pass diff --git a/tinyagent/prompt.py b/tinyagent/prompt.py index 9eb85fd..1f12196 100644 --- a/tinyagent/prompt.py +++ b/tinyagent/prompt.py @@ -26,20 +26,26 @@ "E_FINAL_NOT_ALLOWED": "a tool call is required", } + def estimate_tokens(text: str) -> int: """Estimate token count for budgeting (not a real tokenizer).""" return math.ceil(len(text) / 4) + def _sanitize(text: str) -> str: - chars = [] - for c in text: - if unicodedata.category(c) == "Cc": - chars.append(" ") - else: - chars.append(c) - s = "".join(chars) - s = re.sub(r'\s+', ' ', s) - return s.strip() + """Replace control characters with spaces and collapse all whitespace runs.""" + spaced = "".join(" " if unicodedata.category(c) == "Cc" else c for c in text) + return re.sub(r"\s+", " ", spaced).strip() + + +def _render_default(value: object) -> str: + if isinstance(value, str): + value = _sanitize(value) + if len(value) > 20: + value = value[:20] + "..." + return f'"{value}"' + return repr(value) + def render_tool_line(spec: ToolSpec) -> str: rendered_params = [] @@ -47,32 +53,22 @@ def render_tool_line(spec: ToolSpec) -> str: if p.required: rendered_params.append(f"{p.name}:{p.type}") else: - if p.default is None: - default_str = "None" - elif isinstance(p.default, str): - default_val = _sanitize(p.default) - if len(default_val) > 20: - default_val = default_val[:20] + "..." - default_str = f'"{default_val}"' - else: - default_str = repr(p.default) - rendered_params.append(f"{p.name}:{p.type}={default_str}") - + rendered_params.append(f"{p.name}:{p.type}={_render_default(p.default)}") + params_str = ", ".join(rendered_params) return f"{spec.name}({params_str}) - {spec.description}" + def render_prompt(query: str, tools: Sequence[ToolSpec], *, max_query_chars: int = 2000) -> str: sanitized = _sanitize(query)[:max_query_chars] if not sanitized: raise ValueError("empty query") - - if not tools: - tool_lines = "(none)" - else: - tool_lines = "\n".join(render_tool_line(t) for t in tools) - + + tool_lines = "\n".join(render_tool_line(t) for t in tools) or "(none)" + return PROMPT_TEMPLATE.format(tool_lines=tool_lines, query=sanitized) + def render_repair_prompt(prompt: str, previous_output: str, code: str, detail: str) -> str: """ Renders a repair prompt for a model to try again. @@ -80,8 +76,8 @@ def render_repair_prompt(prompt: str, previous_output: str, code: str, detail: s """ sanitized_output = _sanitize(previous_output)[:200] sanitized_detail = _sanitize(detail)[:40] - + reason_template = _REPAIR_REASONS.get(code, "bad format") short_reason = reason_template.replace("{detail}", sanitized_detail) - + return f'{prompt} {sanitized_output}\nInvalid: {short_reason}. Reply with ONE line: ACTION: tool_name(arg="value") or FINAL: text\nReply:' diff --git a/tinyagent/reflection.py b/tinyagent/reflection.py index d635c0b..a421441 100644 --- a/tinyagent/reflection.py +++ b/tinyagent/reflection.py @@ -1,98 +1,90 @@ import inspect import re import types -import typing -from typing import Any, Callable, Union, get_args, get_origin +from typing import Any, Callable, Union, get_args, get_origin, get_type_hints from .types import ParamSpec, ToolSpec from .registry import ToolRegistry, default_registry -_ALLOWED_TYPES = { - str: "str", - int: "int", - float: "float", - bool: "bool" -} +_ALLOWED_TYPES = {str: "str", int: "int", float: "float", bool: "bool"} +_UNSUPPORTED_KINDS = ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + inspect.Parameter.POSITIONAL_ONLY, +) _NAME_REGEX = re.compile(r"^[A-Za-z_][A-Za-z0-9_]{0,63}$") + def _parse_type_annotation(param_name: str, fn_name: str, annotation: Any) -> tuple[str, bool]: if annotation is inspect.Parameter.empty: return "str", False - - origin = get_origin(annotation) - args = get_args(annotation) - + # Handle Union[T, None] and T | None - is_union = origin is Union or isinstance(annotation, types.UnionType) - if is_union: - none_type = type(None) - valid_args = [a for a in args if a is not none_type] - if len(valid_args) == 1 and type(None) in args: - inner_type = valid_args[0] - if inner_type in _ALLOWED_TYPES: - return _ALLOWED_TYPES[inner_type], True - raise TypeError(f"tiny_tool: parameter '{param_name}' of '{fn_name}' has unsupported type {annotation}; allowed: str,int,float,bool") - - if annotation in _ALLOWED_TYPES: + if get_origin(annotation) is Union or isinstance(annotation, types.UnionType): + args = get_args(annotation) + non_none = [a for a in args if a is not type(None)] + if len(non_none) == 1 and type(None) in args and non_none[0] in _ALLOWED_TYPES: + return _ALLOWED_TYPES[non_none[0]], True + elif annotation in _ALLOWED_TYPES: return _ALLOWED_TYPES[annotation], False - - raise TypeError(f"tiny_tool: parameter '{param_name}' of '{fn_name}' has unsupported type {annotation}; allowed: str,int,float,bool") + + raise TypeError( + f"tiny_tool: parameter '{param_name}' of '{fn_name}' has unsupported type " + f"{annotation}; allowed: str,int,float,bool" + ) + + +def _describe(func: Callable, fallback: str) -> str: + """First non-blank docstring line, whitespace-collapsed and capped at 80 chars.""" + doc = inspect.getdoc(func) or "" + first = next((line for line in doc.splitlines() if line.strip()), None) + if first is None: + return fallback + desc = " ".join(first.split()) + return desc[:77] + "..." if len(desc) > 80 else desc + def tiny_tool(func: Callable | None = None, *, name: str | None = None, registry: ToolRegistry | None = None): if func is None: def wrapper(f: Callable) -> Callable: return tiny_tool(f, name=name, registry=registry) return wrapper - + fn_name = func.__name__ tool_name = name if name is not None else fn_name - + if not _NAME_REGEX.match(tool_name): raise ValueError(f"Invalid tool name: '{tool_name}'") - + sig = inspect.signature(func) - type_hints = typing.get_type_hints(func) - + type_hints = get_type_hints(func) + params = [] for param_name, param in sig.parameters.items(): - if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD, inspect.Parameter.POSITIONAL_ONLY): - raise TypeError(f"tiny_tool: parameter '{param_name}' of '{fn_name}' has unsupported kind {param.kind}") - + if param.kind in _UNSUPPORTED_KINDS: + raise TypeError( + f"tiny_tool: parameter '{param_name}' of '{fn_name}' has unsupported kind {param.kind}" + ) + annotation = type_hints.get(param_name, inspect.Parameter.empty) param_type, is_optional = _parse_type_annotation(param_name, fn_name, annotation) - + required = param.default is inspect.Parameter.empty - default_val = None if required else param.default - + params.append(ParamSpec( name=param_name, type=param_type, required=required, - default=default_val, - optional=is_optional + default=None if required else param.default, + optional=is_optional, )) - - doc = inspect.getdoc(func) - if doc: - lines = [line.strip() for line in doc.splitlines() if line.strip()] - if lines: - desc = " ".join(lines[0].split()) - if len(desc) > 80: - desc = desc[:77] + "..." - else: - desc = tool_name - else: - desc = tool_name - + spec = ToolSpec( name=tool_name, - description=desc, + description=_describe(func, tool_name), params=tuple(params), - func=func + func=func, ) - - reg = registry if registry is not None else default_registry - reg.register(spec) - + (registry if registry is not None else default_registry).register(spec) return func diff --git a/tinyagent/registry.py b/tinyagent/registry.py index c4ba0de..8b74330 100644 --- a/tinyagent/registry.py +++ b/tinyagent/registry.py @@ -1,7 +1,8 @@ -from typing import Tuple +from __future__ import annotations from .types import ToolSpec + class ToolRegistry: def __init__(self) -> None: self._tools: dict[str, ToolSpec] = {} @@ -14,10 +15,10 @@ def register(self, spec: ToolSpec) -> None: def get(self, name: str) -> ToolSpec | None: return self._tools.get(name) - def all(self) -> Tuple[ToolSpec, ...]: + def all(self) -> tuple[ToolSpec, ...]: return tuple(self._tools.values()) - def names(self) -> Tuple[str, ...]: + def names(self) -> tuple[str, ...]: return tuple(self._tools.keys()) def __len__(self) -> int: @@ -26,4 +27,5 @@ def __len__(self) -> int: def __contains__(self, name: str) -> bool: return name in self._tools + default_registry = ToolRegistry() diff --git a/tinyagent/types.py b/tinyagent/types.py index 5fe0fc7..0392031 100644 --- a/tinyagent/types.py +++ b/tinyagent/types.py @@ -1,9 +1,11 @@ from __future__ import annotations + from dataclasses import dataclass from typing import Any, Callable, Literal, Mapping, Protocol, Sequence, Union ParamType = Literal["str", "int", "float", "bool"] + @dataclass(frozen=True) class ParamSpec: name: str @@ -12,6 +14,7 @@ class ParamSpec: default: Any = None optional: bool = False # True when annotated Optional[T]; "None"/"null" coerces to None + @dataclass(frozen=True) class ToolSpec: name: str @@ -19,6 +22,7 @@ class ToolSpec: params: tuple[ParamSpec, ...] func: Callable[..., Any] + @dataclass(frozen=True) class Completion: text: str @@ -27,10 +31,12 @@ class Completion: completion_tokens: int | None backend: str + class Backend(Protocol): name: str def complete(self, prompt: str, *, stop: Sequence[str], max_tokens: int) -> Completion: ... + @dataclass(frozen=True) class ParsedAction: tool: str @@ -38,18 +44,22 @@ class ParsedAction: raw: str warnings: tuple[str, ...] = () + @dataclass(frozen=True) class ParsedFinal: text: str raw: str + @dataclass(frozen=True) class ParseFailure: code: str # one of errors.PARSE_CODES raw: str + ParseOutcome = Union[ParsedAction, ParsedFinal, ParseFailure] + @dataclass(frozen=True) class ToolResult: ok: bool @@ -60,6 +70,7 @@ class ToolResult: error_detail: str | None # short, single line, safe to show to the model duration_ms: float + @dataclass(frozen=True) class AgentResult: status: Literal["tool", "final", "error"]