toolstream 0.1.0__py3-none-any.whl

toolstream/_invoke.py

@@ -0,0 +1,126 @@
+"""Agent invocation helpers.
+
+Build a SessionConfig from an AgentDefinition and yield a ready-to-use
+session (async or sync).
+"""
+
+from __future__ import annotations
+
+import contextlib
+from collections.abc import AsyncIterator, Iterator
+from typing import Callable
+
+from ._agent import AgentDefinition, ToolRef, resolve_prompt
+from ._schema import _generate_schema
+from ._session import AsyncSession, SyncSession
+from ._tools import Tool
+from .config import SessionConfig
+
+__all__ = ["invoke_agent", "invoke_agent_sync"]
+
+
+def _filter_tools(
+    definition: AgentDefinition,
+    available_tools: dict[str, Callable] | None,
+) -> list[Tool] | None:
+    """Match ToolRefs in *definition* against *available_tools* handlers.
+
+    Returns None when the definition declares no tools or no tools dict
+    is provided.  Raises ValueError if any declared tool name is missing
+    from *available_tools*.
+    """
+    if definition.tools is None or available_tools is None:
+        return None
+
+    missing: list[str] = []
+    matched: list[Tool] = []
+
+    for ref in definition.tools:
+        handler = available_tools.get(ref.name)
+        if handler is None:
+            missing.append(ref.name)
+            continue
+
+        if hasattr(handler, "_tool"):
+            matched.append(handler._tool)
+        else:
+            # Generate schema on the fly for plain callables.
+            input_schema = _generate_schema(handler, inject=set())
+            doc = handler.__doc__
+            description = doc.strip().split("\n")[0].strip() if doc else ""
+            matched.append(
+                Tool(
+                    name=ref.name,
+                    description=description,
+                    input_schema=input_schema,
+                    handler=handler,
+                    inject=[],
+                )
+            )
+
+    if missing:
+        raise ValueError(
+            f"Missing tool handlers: {', '.join(missing)}"
+        )
+
+    return matched
+
+
+def _build_invocation_config(
+    definition: AgentDefinition,
+    config: SessionConfig,
+    *,
+    variables: dict[str, str] | None = None,
+    available_tools: dict[str, Callable] | None = None,
+) -> SessionConfig:
+    """Create a SessionConfig tailored to *definition*."""
+    system_prompt = resolve_prompt(definition.prompt_template, variables or {})
+    tools = _filter_tools(definition, available_tools)
+
+    return SessionConfig(
+        model=definition.model or config.model,
+        api_key=config.api_key,
+        base_url=config.base_url,
+        system_prompt=system_prompt,
+        cwd=config.cwd,
+        tools=tools,
+        tool_context=config.tool_context,
+        tool_env=config.tool_env,
+        max_completion_tokens=config.max_completion_tokens,
+        sandbox=config.sandbox,
+        metadata=config.metadata,
+    )
+
+
+@contextlib.asynccontextmanager
+async def invoke_agent(
+    definition: AgentDefinition,
+    config: SessionConfig,
+    *,
+    variables: dict[str, str] | None = None,
+    available_tools: dict[str, Callable] | None = None,
+) -> AsyncIterator[AsyncSession]:
+    """Async context manager that yields an AsyncSession for *definition*."""
+    invocation_config = _build_invocation_config(
+        definition, config, variables=variables, available_tools=available_tools,
+    )
+    session = AsyncSession(invocation_config)
+    async with session:
+        yield session
+
+
+@contextlib.contextmanager
+def invoke_agent_sync(
+    definition: AgentDefinition,
+    config: SessionConfig,
+    *,
+    variables: dict[str, str] | None = None,
+    available_tools: dict[str, Callable] | None = None,
+) -> Iterator[SyncSession]:
+    """Sync context manager that yields a SyncSession for *definition*."""
+    invocation_config = _build_invocation_config(
+        definition, config, variables=variables, available_tools=available_tools,
+    )
+    session = SyncSession(invocation_config)
+    with session:
+        yield session

toolstream/_protocol.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import json
+
+from .events import Error, StepFinish, StepStart, Text, ToolUse
+
+Event = StepStart | Text | ToolUse | StepFinish | Error
+
+
+def parse_event(line: str) -> Event | None:
+    """Parse a single NDJSON line into a typed event.
+
+    Returns None for unparseable or unrecognized lines.
+    """
+    line = line.strip()
+    if not line:
+        return None
+
+    try:
+        data = json.loads(line)
+    except json.JSONDecodeError:
+        return None
+
+    if not isinstance(data, dict):
+        return None
+
+    event_type = data.get("type")
+    session_id = data.get("sessionID", "")
+    timestamp = data.get("timestamp", 0)
+    part = data.get("part", {})
+
+    if event_type == "step_start":
+        return StepStart(
+            session_id=session_id,
+            message_id=part.get("messageID", ""),
+            timestamp=timestamp,
+        )
+
+    if event_type == "text":
+        return Text(
+            session_id=session_id,
+            message_id=part.get("messageID", ""),
+            text=part.get("text", ""),
+            timestamp=timestamp,
+        )
+
+    if event_type == "tool_use":
+        state = part.get("state", {})
+        return ToolUse(
+            session_id=session_id,
+            message_id=part.get("messageID", ""),
+            tool=part.get("tool", ""),
+            call_id=part.get("callID", ""),
+            status=state.get("status", ""),
+            input=state.get("input", {}),
+            output=state.get("output", ""),
+            title=part.get("title", ""),
+            timestamp=timestamp,
+        )
+
+    if event_type == "step_finish":
+        tokens = part.get("tokens", {})
+        cache = tokens.get("cache", {})
+        return StepFinish(
+            session_id=session_id,
+            message_id=part.get("messageID", ""),
+            reason=part.get("reason", ""),
+            input_tokens=tokens.get("input", 0),
+            output_tokens=tokens.get("output", 0),
+            reasoning_tokens=tokens.get("reasoning", 0),
+            cache_read_tokens=cache.get("read", 0),
+            cache_write_tokens=cache.get("write", 0),
+            cost=part.get("cost", 0.0),
+            timestamp=timestamp,
+        )
+
+    if event_type == "error":
+        error = data.get("error", {})
+        return Error(
+            session_id=session_id,
+            name=error.get("name", ""),
+            message=error.get("message", ""),
+            data=error.get("data", {}),
+            timestamp=timestamp,
+        )
+
+    return None

toolstream/_schema.py

@@ -0,0 +1,259 @@
+"""Generate JSON Schema from Python function type hints.
+
+Foundation for the tool registration system. Converts function signatures
+(parameters, type annotations, docstrings) into JSON Schema objects
+suitable for LLM tool definitions.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import enum
+import inspect
+import re
+import typing
+
+
+def _type_to_schema(annotation: type) -> dict:
+    """Convert a Python type hint to a JSON Schema fragment."""
+
+    # Handle None / NoneType directly
+    if annotation is type(None):
+        return {"type": "null"}
+
+    # Handle missing annotation
+    if annotation is inspect.Parameter.empty:
+        return {"type": "object"}
+
+    origin = typing.get_origin(annotation)
+    args = typing.get_args(annotation)
+
+    # Optional[T] is Union[T, None]; T | None has origin types.UnionType
+    if origin is typing.Union or _is_union_type(annotation):
+        non_none = [a for a in args if a is not type(None)]
+        if len(non_none) == 1 and len(args) == 2:
+            # Optional[T] -- produce nullable schema
+            inner = _type_to_schema(non_none[0])
+            if "type" in inner:
+                t = inner["type"]
+                if isinstance(t, list):
+                    if "null" not in t:
+                        inner["type"] = t + ["null"]
+                else:
+                    inner["type"] = [t, "null"]
+            elif "enum" in inner:
+                inner["enum"] = inner["enum"] + [None]
+            else:
+                inner["type"] = ["object", "null"]
+            return inner
+        # General Union -- not handled beyond Optional, fall through
+        return {"type": "object"}
+
+    # Literal["a", "b"]
+    if origin is typing.Literal:
+        values = list(args)
+        # Infer JSON Schema type from the literal values
+        value_types = {type(v) for v in values}
+        if value_types == {int}:
+            schema_type = "integer"
+        elif value_types == {float}:
+            schema_type = "number"
+        elif value_types == {bool}:
+            schema_type = "boolean"
+        else:
+            schema_type = "string"
+        return {"type": schema_type, "enum": values}
+
+    # Enum subclass
+    if isinstance(annotation, type) and issubclass(annotation, enum.Enum):
+        return {"enum": [member.value for member in annotation]}
+
+    # list[T]
+    if origin is list:
+        if args:
+            return {"type": "array", "items": _type_to_schema(args[0])}
+        return {"type": "array"}
+
+    # dict[str, T]
+    if origin is dict:
+        if args and len(args) == 2:
+            return {
+                "type": "object",
+                "additionalProperties": _type_to_schema(args[1]),
+            }
+        return {"type": "object"}
+
+    # dataclass
+    if dataclasses.is_dataclass(annotation) and isinstance(annotation, type):
+        properties = {}
+        required = []
+        # Resolve stringified annotations for dataclass fields
+        try:
+            resolved = typing.get_type_hints(annotation)
+        except Exception:
+            resolved = {}
+        for field in dataclasses.fields(annotation):
+            field_type = resolved.get(field.name, field.type)
+            prop = _type_to_schema(field_type)
+            properties[field.name] = prop
+            # Fields with default or default_factory are optional
+            has_default = (
+                field.default is not dataclasses.MISSING
+                or field.default_factory is not dataclasses.MISSING
+            )
+            if not has_default:
+                required.append(field.name)
+        schema: dict = {"type": "object", "properties": properties}
+        if required:
+            schema["required"] = required
+        return schema
+
+    # TypedDict
+    if _is_typed_dict(annotation):
+        properties = {}
+        hints = typing.get_type_hints(annotation)
+        req_keys = getattr(annotation, "__required_keys__", frozenset())
+        for name, hint in hints.items():
+            properties[name] = _type_to_schema(hint)
+        schema = {"type": "object", "properties": properties}
+        req = [k for k in hints if k in req_keys]
+        if req:
+            schema["required"] = req
+        return schema
+
+    # Primitive types
+    primitives = {
+        str: {"type": "string"},
+        int: {"type": "integer"},
+        float: {"type": "number"},
+        bool: {"type": "boolean"},
+        dict: {"type": "object"},
+        list: {"type": "array"},
+    }
+    if annotation in primitives:
+        return dict(primitives[annotation])
+
+    return {"type": "object"}
+
+
+def _is_union_type(annotation: type) -> bool:
+    """Check if annotation is a PEP 604 union (X | Y)."""
+    import types
+
+    return isinstance(annotation, types.UnionType)
+
+
+def _is_typed_dict(annotation: type) -> bool:
+    """Check if annotation is a TypedDict subclass."""
+    return (
+        isinstance(annotation, type)
+        and issubclass(annotation, dict)
+        and hasattr(annotation, "__annotations__")
+        and hasattr(annotation, "__required_keys__")
+    )
+
+
+# Regex for Google-style docstring param lines:
+#   Args:
+#       name: description text
+#       name (type): description text
+_GOOGLE_ARGS_RE = re.compile(
+    r"^\s{2,}(\w+)(?:\s*\([^)]*\))?\s*:\s*(.+)", re.MULTILINE
+)
+
+# Regex for reST-style docstring param lines:
+#   :param name: description text
+_REST_PARAM_RE = re.compile(
+    r"^\s*:param\s+(\w+)\s*:\s*(.+)", re.MULTILINE
+)
+
+
+def _parse_param_descriptions(fn: typing.Callable) -> dict[str, str]:
+    """Extract parameter descriptions from a function's docstring.
+
+    Handles Google-style (Args: section) and reST-style (:param:) formats.
+    Follows __wrapped__ for functools.wraps-decorated functions.
+    """
+    # Follow __wrapped__ chain to find the original docstring
+    target = fn
+    while hasattr(target, "__wrapped__"):
+        target = target.__wrapped__
+
+    doc = inspect.getdoc(target)
+    if not doc:
+        return {}
+
+    descriptions: dict[str, str] = {}
+
+    # Try Google-style: find the Args: section
+    args_match = re.search(r"^\s*Args?\s*:\s*$", doc, re.MULTILINE)
+    if args_match:
+        # Extract the block after "Args:" until the next section or end
+        after_args = doc[args_match.end() :]
+        # Stop at next section header (word followed by colon at start of line,
+        # or end of string)
+        section_end = re.search(r"^\S", after_args, re.MULTILINE)
+        args_block = after_args[: section_end.start()] if section_end else after_args
+        for m in _GOOGLE_ARGS_RE.finditer(args_block):
+            descriptions[m.group(1)] = m.group(2).strip()
+
+    # Try reST-style
+    for m in _REST_PARAM_RE.finditer(doc):
+        # Don't overwrite Google-style if both present
+        if m.group(1) not in descriptions:
+            descriptions[m.group(1)] = m.group(2).strip()
+
+    return descriptions
+
+
+def _generate_schema(
+    fn: typing.Callable,
+    inject: set[str] | None = None,
+) -> dict:
+    """Generate a JSON Schema object from a function's signature.
+
+    Args:
+        fn: The function to generate a schema for.
+        inject: Parameter names to exclude (they will be injected at call time).
+
+    Returns:
+        A JSON Schema dict with "type", "properties", and "required" keys.
+    """
+    inject = inject or set()
+    sig = inspect.signature(fn)
+    descriptions = _parse_param_descriptions(fn)
+
+    # Resolve stringified annotations (from `from __future__ import annotations`)
+    try:
+        resolved_hints = typing.get_type_hints(fn)
+    except Exception:
+        resolved_hints = {}
+
+    properties: dict[str, dict] = {}
+    required: list[str] = []
+
+    for name, param in sig.parameters.items():
+        # Skip self/cls
+        if name in ("self", "cls"):
+            continue
+        # Skip injected parameters
+        if name in inject:
+            continue
+
+        annotation = resolved_hints.get(name, param.annotation)
+        prop = _type_to_schema(annotation)
+
+        # Add description if available
+        if name in descriptions:
+            prop["description"] = descriptions[name]
+
+        properties[name] = prop
+
+        # Parameters with no default are required
+        if param.default is inspect.Parameter.empty:
+            required.append(name)
+
+    schema: dict = {"type": "object", "properties": properties}
+    if required:
+        schema["required"] = required
+    return schema

toolstream/_session.py

@@ -0,0 +1,176 @@
+from __future__ import annotations
+
+import asyncio
+import queue
+import threading
+from collections.abc import AsyncIterator, Iterator
+from typing import Any
+
+from ._direct import DirectClient
+from .config import SessionConfig
+from .events import Result, StepFinish
+
+
+class AsyncSession:
+    """Async session wrapping the direct LLM API client."""
+
+    def __init__(self, config: SessionConfig):
+        self._config = config
+        self._direct = DirectClient(
+            config,
+            tools=config.tools,
+            tool_context=config.tool_context,
+            max_completion_tokens=config.max_completion_tokens,
+        )
+        self._turn_count = 0
+        self._total_cost = 0.0
+        self._total_input_tokens = 0
+        self._total_output_tokens = 0
+
+    async def __aenter__(self) -> AsyncSession:
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        await self.close()
+
+    async def send(self, message: str) -> AsyncIterator[Any]:
+        """Send a message and yield events. Yields a Result summary at the end."""
+        async for event in self._send_direct(message):
+            yield event
+
+    async def _send_direct(self, message: str) -> AsyncIterator[Any]:
+        """Send via direct API client."""
+        self._turn_count += 1
+
+        turn_input_tokens = 0
+        turn_output_tokens = 0
+        turn_reasoning_tokens = 0
+        turn_cache_read_tokens = 0
+        turn_cache_write_tokens = 0
+        turn_cost = 0.0
+        steps = 0
+
+        async for event in self._direct.send(message):
+            if isinstance(event, StepFinish):
+                turn_input_tokens += event.input_tokens
+                turn_output_tokens += event.output_tokens
+                turn_reasoning_tokens += event.reasoning_tokens
+                turn_cache_read_tokens += event.cache_read_tokens
+                turn_cache_write_tokens += event.cache_write_tokens
+                turn_cost += event.cost
+                steps += 1
+            yield event
+
+        # Update totals
+        self._total_input_tokens += turn_input_tokens
+        self._total_output_tokens += turn_output_tokens
+        self._total_cost += turn_cost
+
+        # Yield result summary
+        yield Result(
+            session_id=self._direct.session_id,
+            total_input_tokens=turn_input_tokens,
+            total_output_tokens=turn_output_tokens,
+            total_cost=turn_cost,
+            steps=steps,
+        )
+
+    async def close(self) -> None:
+        """Close the session and clean up."""
+        await self._direct.close()
+
+    @property
+    def session_id(self) -> str | None:
+        return self._direct.session_id
+
+    @property
+    def total_cost(self) -> float:
+        return self._total_cost
+
+    @property
+    def turn_count(self) -> int:
+        return self._turn_count
+
+
+class SyncSession:
+    """Sync wrapper around AsyncSession using a dedicated event loop thread."""
+
+    def __init__(self, config: SessionConfig):
+        self._config = config
+        self._async_session: AsyncSession | None = None
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._thread: threading.Thread | None = None
+
+    def _start_loop(self) -> None:
+        """Run the event loop in a background thread."""
+        assert self._loop is not None
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+
+    def __enter__(self) -> SyncSession:
+        self._loop = asyncio.new_event_loop()
+        self._thread = threading.Thread(target=self._start_loop, daemon=True)
+        self._thread.start()
+        self._async_session = AsyncSession(self._config)
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self.close()
+
+    def _run_coroutine(self, coro: Any) -> Any:
+        """Run a coroutine on the background event loop and return the result."""
+        assert self._loop is not None
+        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
+        return future.result()
+
+    def send(self, message: str) -> Iterator[Any]:
+        """Send a message and yield events incrementally via a queue bridge."""
+        assert self._async_session is not None
+        assert self._loop is not None
+
+        q: queue.Queue = queue.Queue()
+        sentinel = object()
+
+        async def _produce() -> None:
+            try:
+                async for event in self._async_session.send(message):  # type: ignore[union-attr]
+                    q.put(event)
+            except Exception as e:
+                q.put(e)
+            finally:
+                q.put(sentinel)
+
+        asyncio.run_coroutine_threadsafe(_produce(), self._loop)
+
+        while True:
+            item = q.get()
+            if item is sentinel:
+                break
+            if isinstance(item, Exception):
+                raise item
+            yield item
+
+    def close(self) -> None:
+        """Close the session and shut down the event loop."""
+        if self._async_session is not None:
+            self._run_coroutine(self._async_session.close())
+            self._async_session = None
+        if self._loop is not None:
+            self._loop.call_soon_threadsafe(self._loop.stop)
+            if self._thread is not None:
+                self._thread.join(timeout=5.0)
+            self._loop.close()
+            self._loop = None
+            self._thread = None
+
+    @property
+    def session_id(self) -> str | None:
+        return self._async_session.session_id if self._async_session else None
+
+    @property
+    def total_cost(self) -> float:
+        return self._async_session.total_cost if self._async_session else 0.0
+
+    @property
+    def turn_count(self) -> int:
+        return self._async_session.turn_count if self._async_session else 0