toolstream 0.1.0__py3-none-any.whl

toolstream/__init__.py

@@ -0,0 +1,44 @@
+from ._agent import (
+    AgentDefinition,
+    AgentSandbox,
+    ToolRef,
+    discover_agents,
+    load_agent,
+    resolve_prompt,
+)
+from ._context import ToolContext
+from ._invoke import invoke_agent, invoke_agent_sync
+from ._session import AsyncSession, SyncSession
+from ._tools import Tool, collect_tools, tool
+from .config import SessionConfig
+from .events import Error, Result, StepFinish, StepStart, Text, ToolUse
+
+__all__ = [
+    # Agent definitions
+    "AgentDefinition",
+    "AgentSandbox",
+    "ToolRef",
+    "load_agent",
+    "discover_agents",
+    "resolve_prompt",
+    # Agent invocation
+    "invoke_agent",
+    "invoke_agent_sync",
+    # Tools
+    "tool",
+    "Tool",
+    "collect_tools",
+    # Context
+    "ToolContext",
+    # Sessions
+    "AsyncSession",
+    "SyncSession",
+    "SessionConfig",
+    # Events
+    "StepStart",
+    "Text",
+    "ToolUse",
+    "StepFinish",
+    "Error",
+    "Result",
+]

toolstream/_agent.py

@@ -0,0 +1,215 @@
+"""Agent definition loading and discovery."""
+
+from __future__ import annotations
+
+import importlib.resources
+import json
+import logging
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "ToolRef",
+    "AgentSandbox",
+    "AgentDefinition",
+    "load_agent",
+    "discover_agents",
+    "resolve_prompt",
+]
+
+
+@dataclass(frozen=True)
+class ToolRef:
+    name: str
+
+
+@dataclass(frozen=True)
+class AgentSandbox:
+    tools: list[str] | None = None
+    skip_permissions: bool = False
+
+
+@dataclass(frozen=True)
+class AgentDefinition:
+    name: str
+    prompt_template: str
+    version: str
+    description: str = ""
+    tools: list[ToolRef] | None = None
+    sandbox: AgentSandbox | None = None
+    model: str | None = None
+
+
+def _parse_agent_data(data: dict, source: str) -> AgentDefinition:
+    """Parse a dict (from JSON) into an AgentDefinition.
+
+    Args:
+        data: Parsed JSON dict.
+        source: Human-readable origin for error messages (file path, package).
+    """
+    tools: list[ToolRef] | None = None
+    if "tools" in data:
+        tools = [ToolRef(name=entry["name"]) for entry in (data["tools"] or [])]
+
+    sandbox: AgentSandbox | None = None
+    if "sandbox" in data:
+        sandbox_data = data["sandbox"]
+        sandbox = AgentSandbox(
+            tools=sandbox_data.get("tools"),
+            skip_permissions=sandbox_data.get("skip_permissions", False),
+        )
+
+    return AgentDefinition(
+        name=data["name"],
+        prompt_template=data["prompt_template"],
+        version=data["version"],
+        description=data.get("description", ""),
+        tools=tools,
+        sandbox=sandbox,
+        model=data.get("model"),
+    )
+
+
+def load_agent(path: str | Path) -> AgentDefinition:
+    """Load an agent definition from a JSON file.
+
+    Args:
+        path: Path to a ``.agent.json`` file. Bare names (no path separator,
+              no ``.json`` suffix) are rejected -- use :func:`discover_agents`
+              for name-based lookup.
+
+    Raises:
+        FileNotFoundError: If *path* is a bare name or the file does not exist.
+    """
+    if isinstance(path, str) and "/" not in path and os.sep not in path and not path.endswith(".json"):
+        raise FileNotFoundError(
+            f"Agent '{path}' not found (bare name resolution requires discover_agents)"
+        )
+
+    p = Path(path)
+    data = json.loads(p.read_text(encoding="utf-8"))
+    return _parse_agent_data(data, source=str(p))
+
+
+def _try_load(
+    file_path: Path,
+    seen: dict[str, AgentDefinition],
+) -> None:
+    """Attempt to load a single agent file, deduplicating by name."""
+    try:
+        agent = load_agent(file_path)
+    except (json.JSONDecodeError, KeyError, TypeError) as exc:
+        logger.warning("Skipping %s: %s", file_path, exc)
+        return
+    except FileNotFoundError as exc:
+        logger.warning("Skipping %s: %s", file_path, exc)
+        return
+
+    if agent.name in seen:
+        logger.warning(
+            "Duplicate agent '%s' in %s (already loaded); skipping",
+            agent.name,
+            file_path,
+        )
+        return
+
+    seen[agent.name] = agent
+
+
+def discover_agents(
+    cwd: str | Path | None = None,
+    paths: list[str | Path] | None = None,
+    packages: list[str] | None = None,
+) -> list[AgentDefinition]:
+    """Discover agent definitions from directories and packages.
+
+    Search order (first occurrence of a name wins):
+
+    1. ``<cwd>/.toolstream/agents/*.agent.json``
+    2. Each directory in *paths*: ``<dir>/*.agent.json``
+    3. Each Python package in *packages*: resources matching ``*.agent.json``
+
+    Returns:
+        Agent definitions sorted alphabetically by name.
+    """
+    seen: dict[str, AgentDefinition] = {}
+
+    # 1. cwd-local agents
+    if cwd is not None:
+        agents_dir = Path(cwd) / ".toolstream" / "agents"
+        if agents_dir.is_dir():
+            for f in sorted(agents_dir.glob("*.agent.json")):
+                _try_load(f, seen)
+
+    # 2. Explicit directory paths
+    if paths is not None:
+        for dir_path in paths:
+            d = Path(dir_path)
+            if not d.is_dir():
+                logger.warning("Agent path %s is not a directory; skipping", d)
+                continue
+            for f in sorted(d.glob("*.agent.json")):
+                _try_load(f, seen)
+
+    # 3. Python packages
+    if packages is not None:
+        for package_name in packages:
+            try:
+                pkg_files = importlib.resources.files(package_name)
+            except ModuleNotFoundError:
+                logger.warning("Package '%s' not found; skipping", package_name)
+                continue
+
+            for item in pkg_files.iterdir():
+                if not item.name.endswith(".agent.json"):
+                    continue
+                try:
+                    data = json.loads(item.read_text(encoding="utf-8"))
+                    agent = _parse_agent_data(data, source=f"{package_name}/{item.name}")
+                except (json.JSONDecodeError, KeyError, TypeError) as exc:
+                    logger.warning(
+                        "Skipping %s/%s: %s", package_name, item.name, exc
+                    )
+                    continue
+
+                if agent.name in seen:
+                    logger.warning(
+                        "Duplicate agent '%s' in %s/%s (already loaded); skipping",
+                        agent.name,
+                        package_name,
+                        item.name,
+                    )
+                    continue
+
+                seen[agent.name] = agent
+
+    return sorted(seen.values(), key=lambda a: a.name)
+
+
+def resolve_prompt(template: str, variables: dict[str, str]) -> str:
+    """Substitute ``{key}`` placeholders in a prompt template.
+
+    Uses :meth:`str.replace` (not :meth:`str.format`) so only the
+    supplied *variables* are touched; literal braces in the template
+    are preserved.
+
+    Raises:
+        ValueError: If any original placeholder has no matching variable.
+    """
+    # Find all original placeholders before substitution.
+    original_placeholders = set(re.findall(r"\{(\w+)\}", template))
+
+    result = template
+    for key, value in variables.items():
+        result = result.replace(f"{{{key}}}", value)
+
+    # Only check placeholders that were in the original template.
+    unresolved = [p for p in sorted(original_placeholders) if p not in variables]
+    if unresolved:
+        raise ValueError(f"Unresolved template variables: {', '.join(unresolved)}")
+
+    return result

toolstream/_builtin_tools.py

@@ -0,0 +1,115 @@
+"""Standalone implementations of the built-in tools for direct-mode LLM sessions.
+
+Each function is an async def decorated with @tool that receives injectable
+parameters (excluded from the JSON Schema sent to the LLM, passed explicitly
+by the DirectClient's dispatch loop).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import shlex
+from glob import glob as _stdlib_glob
+from pathlib import Path
+
+from ._tools import tool
+
+_MAX_OUTPUT_CHARS = 50_000
+
+
+def _resolve_path(file_path: str, cwd: str) -> Path:
+    """Resolve *file_path*, making relative paths absolute against *cwd*."""
+    path = Path(file_path)
+    if not path.is_absolute():
+        path = Path(cwd) / path
+    return path
+
+
+@tool(inject=["cwd"])
+async def read(file_path: str, cwd: str, offset: int = 0, limit: int = 2000) -> str:
+    """Read a file and return its contents with line numbers."""
+    path = _resolve_path(file_path, cwd)
+    lines = path.read_text().splitlines()
+    selected = lines[offset : offset + limit]
+    return "\n".join(f"{i + offset + 1}: {line}" for i, line in enumerate(selected))
+
+
+@tool(inject=["cwd"])
+async def write(file_path: str, content: str, cwd: str) -> str:
+    """Write content to a file, creating parent directories as needed."""
+    path = _resolve_path(file_path, cwd)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(content)
+    return f"Wrote {len(content)} bytes to {file_path}"
+
+
+@tool(inject=["cwd", "env"])
+async def bash(command: str, cwd: str, env: dict[str, str], timeout: int = 120) -> str:
+    """Run a shell command and return stdout+stderr combined."""
+    proc = await asyncio.create_subprocess_shell(
+        command,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.STDOUT,
+        cwd=cwd,
+        env=env,
+    )
+    try:
+        stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=timeout)
+        output = stdout.decode(errors="replace")
+        if len(output) > _MAX_OUTPUT_CHARS:
+            output = output[:_MAX_OUTPUT_CHARS] + f"\n... (truncated, {len(output)} total chars)"
+        return output
+    except asyncio.TimeoutError:
+        proc.kill()
+        await proc.wait()
+        return f"Command timed out after {timeout}s"
+
+
+@tool(inject=["cwd"])
+async def edit(file_path: str, old_string: str, new_string: str, cwd: str) -> str:
+    """Edit a file by replacing the first occurrence of old_string with new_string."""
+    path = _resolve_path(file_path, cwd)
+    content = path.read_text()
+    if old_string not in content:
+        return f"Error: old_string not found in {file_path}"
+    content = content.replace(old_string, new_string, 1)
+    path.write_text(content)
+    return f"Edited {file_path}"
+
+
+@tool(inject=["cwd"])
+async def grep(pattern: str, path: str, cwd: str, include: str | None = None) -> str:
+    """Search for a pattern in files using grep -rn, piped to head -50."""
+    cmd = f"grep -rn {shlex.quote(pattern)} {shlex.quote(path)}"
+    if include:
+        cmd += f" --include={shlex.quote(include)}"
+    cmd += " | head -50"
+    grep_env = {"PATH": os.environ.get("PATH", "/usr/bin:/bin")}
+    proc = await asyncio.create_subprocess_shell(
+        cmd,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.STDOUT,
+        cwd=cwd,
+        env=grep_env,
+    )
+    try:
+        stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=30)
+        output = stdout.decode(errors="replace")
+        if len(output) > _MAX_OUTPUT_CHARS:
+            output = output[:_MAX_OUTPUT_CHARS] + f"\n... (truncated, {len(output)} total chars)"
+        return output
+    except asyncio.TimeoutError:
+        proc.kill()
+        await proc.wait()
+        return "Command timed out after 30s"
+
+
+@tool(name="glob", inject=["cwd"])
+async def glob_files(pattern: str, cwd: str) -> str:
+    """Find files matching a glob pattern. Returns up to 100 matches."""
+    if os.path.isabs(pattern):
+        matches = _stdlib_glob(pattern, recursive=True)
+    else:
+        matches = _stdlib_glob(pattern, root_dir=cwd, recursive=True)
+    return "\n".join(matches[:100])

toolstream/_context.py

@@ -0,0 +1,14 @@
+"""ToolContext -- base class for tool dependency injection."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class ToolContext:
+    """Base class for tool dependency injection.
+
+    Subclass with your application's injectable attributes.
+    Tools declare inject parameters that are resolved via getattr on this object.
+    """

toolstream/_direct.py

@@ -0,0 +1,292 @@
+"""Direct LLM API client via AI Gateway."""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+import uuid
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+
+from . import _builtin_tools
+from ._tools import Tool, collect_tools
+from .config import SessionConfig
+from .events import Error, StepFinish, StepStart, Text, ToolUse
+
+_MAX_RETRIES = 1
+
+
+def _strip_provider(model: str) -> str:
+    """Strip provider prefix from model name.
+
+    'azure-cognitive-services/gpt-5.4' -> 'gpt-5.4'
+    'gpt-5.4' -> 'gpt-5.4'
+    """
+    if "/" in model:
+        return model.rsplit("/", 1)[1]
+    return model
+
+
+def _timestamp_ms() -> int:
+    return int(time.time() * 1000)
+
+
+def _build_tool_definitions(tools: dict[str, Tool]) -> list[dict]:
+    """Build OpenAI-format tool definitions from Tool objects."""
+    defs = []
+    for t in tools.values():
+        defs.append({
+            "type": "function",
+            "function": {
+                "name": t.name,
+                "description": t.description,
+                "parameters": t.input_schema,
+            },
+        })
+    return defs
+
+
+class DirectClient:
+    """Direct LLM API client (via AI Gateway) with tool calling."""
+
+    def __init__(
+        self,
+        config: SessionConfig,
+        *,
+        tools: list[Tool] | None = None,
+        tool_context: object | None = None,
+        max_completion_tokens: int = 16384,
+        http_client: httpx.AsyncClient | None = None,
+    ) -> None:
+        if not config.api_key:
+            raise ValueError("api_key is required for direct backend")
+        if not config.base_url:
+            raise ValueError("base_url is required for direct backend")
+
+        self._config = config
+        self._base_url = config.base_url.rstrip("/")
+        self._api_key = config.api_key
+        self._model = _strip_provider(config.model)
+        self._messages: list[dict] = []
+        self._session_id = str(uuid.uuid4())
+        self._cwd = config.cwd or os.getcwd()
+        self._max_completion_tokens = max_completion_tokens
+        self._tool_context = tool_context
+        self._owns_client = http_client is None
+
+        # Builtin context for inject resolution
+        self._builtin_context: dict[str, Any] = {
+            "cwd": self._cwd,
+            "env": config.tool_env,
+        }
+
+        # HTTP client: reuse injected or create one
+        if http_client is not None:
+            self._client = http_client
+        else:
+            self._client = httpx.AsyncClient(
+                timeout=httpx.Timeout(connect=10, read=120, write=10, pool=10),
+            )
+
+        # Build tool registry: start with built-ins, override with user tools
+        builtin_tools = collect_tools(_builtin_tools)
+        self._builtin_names: set[str] = {t.name for t in builtin_tools}
+        self._tools: dict[str, Tool] = {t.name: t for t in builtin_tools}
+
+        if tools is not None:
+            for t in tools:
+                self._builtin_names.discard(t.name)
+                self._tools[t.name] = t
+
+        # Pre-compute API tool definitions
+        self._tool_definitions = _build_tool_definitions(self._tools)
+
+        # Build system prompt
+        self._messages.append({"role": "system", "content": config.system_prompt})
+
+    @property
+    def session_id(self) -> str:
+        return self._session_id
+
+    async def send(self, message: str) -> AsyncIterator[StepStart | Text | ToolUse | StepFinish | Error]:
+        """Send a message and yield events. Handles the tool-calling loop internally."""
+        self._messages.append({"role": "user", "content": message})
+        msg_id = str(uuid.uuid4())
+
+        yield StepStart(
+            session_id=self._session_id,
+            message_id=msg_id,
+            timestamp=_timestamp_ms(),
+        )
+
+        total_input_tokens = 0
+        total_output_tokens = 0
+
+        while True:
+            response = await self._chat_completion(self._messages)
+            usage = response.get("usage", {})
+            total_input_tokens += usage.get("prompt_tokens", 0)
+            total_output_tokens += usage.get("completion_tokens", 0)
+
+            choice = response["choices"][0]
+            assistant_msg = choice["message"]
+
+            # Build the message dict to append to conversation
+            msg_to_append: dict = {"role": "assistant"}
+            if assistant_msg.get("content"):
+                msg_to_append["content"] = assistant_msg["content"]
+            if assistant_msg.get("tool_calls"):
+                msg_to_append["tool_calls"] = assistant_msg["tool_calls"]
+            self._messages.append(msg_to_append)
+
+            # Yield text if present
+            if assistant_msg.get("content"):
+                yield Text(
+                    session_id=self._session_id,
+                    message_id=msg_id,
+                    text=assistant_msg["content"],
+                    timestamp=_timestamp_ms(),
+                )
+
+            # Check for tool calls
+            tool_calls = assistant_msg.get("tool_calls", [])
+            if not tool_calls:
+                # No tool calls -- conversation turn is done
+                yield StepFinish(
+                    session_id=self._session_id,
+                    message_id=msg_id,
+                    reason="stop",
+                    input_tokens=total_input_tokens,
+                    output_tokens=total_output_tokens,
+                    reasoning_tokens=usage.get("completion_tokens_details", {}).get("reasoning_tokens", 0),
+                    cache_read_tokens=usage.get("prompt_tokens_details", {}).get("cached_tokens", 0),
+                    cache_write_tokens=0,
+                    cost=0.0,
+                    timestamp=_timestamp_ms(),
+                )
+                break
+
+            # Execute tool calls and add results
+            for tc in tool_calls:
+                func_name = tc["function"]["name"]
+                try:
+                    func_args = json.loads(tc["function"]["arguments"])
+                except json.JSONDecodeError:
+                    func_args = {}
+
+                result = await self._dispatch_tool(func_name, func_args)
+
+                yield ToolUse(
+                    session_id=self._session_id,
+                    message_id=msg_id,
+                    tool=func_name,
+                    call_id=tc["id"],
+                    status="completed",
+                    input=func_args,
+                    output=result,
+                    title=func_name,
+                    timestamp=_timestamp_ms(),
+                )
+
+                self._messages.append({
+                    "role": "tool",
+                    "tool_call_id": tc["id"],
+                    "content": result,
+                })
+
+            # Yield StepFinish for this tool-calling round
+            yield StepFinish(
+                session_id=self._session_id,
+                message_id=msg_id,
+                reason="tool-calls",
+                input_tokens=total_input_tokens,
+                output_tokens=total_output_tokens,
+                reasoning_tokens=0,
+                cache_read_tokens=0,
+                cache_write_tokens=0,
+                cost=0.0,
+                timestamp=_timestamp_ms(),
+            )
+
+    async def _chat_completion(self, messages: list[dict]) -> dict:
+        """Call LLM chat completions via AI Gateway with timeout and retry."""
+        url = self._base_url
+        headers = {
+            "x-api-key": self._api_key,
+            "Content-Type": "application/json",
+        }
+        body: dict[str, Any] = {
+            "model": self._model,
+            "messages": messages,
+            "tools": self._tool_definitions,
+            "max_completion_tokens": self._max_completion_tokens,
+        }
+        if self._config.metadata:
+            body["metadata"] = self._config.metadata
+
+        last_error: Exception | None = None
+        for attempt in range(_MAX_RETRIES + 1):
+            try:
+                response = await self._client.post(url, json=body, headers=headers)
+                response.raise_for_status()
+                return response.json()
+            except httpx.ReadTimeout as e:
+                last_error = e
+                if attempt < _MAX_RETRIES:
+                    continue
+                raise
+            except httpx.HTTPStatusError:
+                raise
+
+        # Should never reach here, but satisfy type checker
+        raise last_error  # type: ignore[misc]
+
+    async def _dispatch_tool(self, name: str, args: dict[str, Any]) -> str:
+        """Dispatch a tool call by name. Returns the result string."""
+        tool_obj = self._tools.get(name)
+        if tool_obj is None:
+            return f"Error: unknown tool '{name}'"
+
+        if name in self._builtin_names:
+            # Resolve inject params from builtin context
+            kwargs = {p: self._builtin_context[p] for p in tool_obj.inject}
+            try:
+                result = await tool_obj.handler(**args, **kwargs)
+            except Exception as e:
+                return f"Error: {e}"
+        else:
+            # User tools: inject context params, then call handler
+            for param_name in tool_obj.inject:
+                if self._tool_context is None:
+                    raise RuntimeError(
+                        f"Tool '{name}' requires tool_context "
+                        f"(inject=['{param_name}']) but tool_context is None"
+                    )
+                try:
+                    args[param_name] = getattr(self._tool_context, param_name)
+                except AttributeError:
+                    raise AttributeError(
+                        f"tool_context ({type(self._tool_context).__name__}) "
+                        f"has no attribute '{param_name}' "
+                        f"required by tool '{name}'"
+                    )
+            try:
+                result = await tool_obj.handler(**args)
+            except Exception as e:
+                return f"Error: {e}"
+
+        return result
+
+    async def close(self) -> None:
+        """Close the HTTP client if we own it."""
+        if self._owns_client:
+            await self._client.aclose()
+
+    async def __aenter__(self) -> DirectClient:
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        await self.close()