framework-agents 0.1.0__py3-none-any.whl

framework_agents/tools.py

@@ -0,0 +1,114 @@
+"""Tools: pydantic-validated args. Parity: src/tools/index.ts."""
+
+from __future__ import annotations
+
+import itertools
+from dataclasses import dataclass, field
+from typing import Any, Awaitable, Callable, Optional, Type
+
+from pydantic import BaseModel, ValidationError
+
+from .errors import AgentError
+from .state import AgentState
+
+
+@dataclass
+class ToolContext:
+    """Arg #2 for tools that need state (agent name / agent state / invocation
+    state). The agent loop wires its AgentState and the LIVE per-run
+    invocation_state (by reference) at each dispatch, like the TS ``toolCtx``."""
+
+    agent_name: str = "agent"
+    # The agent's AgentState (dict-like via .get/.set) — defaults to a fresh
+    # AgentState for standalone/registry calls.
+    state: AgentState = field(default_factory=AgentState)
+    invocation_state: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class ToolResult:
+    """Split raw output: model-visible content vs downstream-only artifact."""
+    content: Any
+    artifact: Any = None
+
+
+async def _default_to_result(output: Any) -> ToolResult:
+    return ToolResult(content=output if isinstance(output, str) else _jsonish(output))
+
+
+def _jsonish(value: Any) -> str:
+    import json
+
+    return json.dumps(value, default=str)
+
+
+@dataclass
+class Tool:
+    name: str
+    description: str
+    schema: Type[BaseModel]  # pydantic model = the validated arg shape
+    execute: Callable[[Any, ToolContext | None], Awaitable[Any]]
+    to_result: Callable[[Any], Awaitable[ToolResult]] = _default_to_result
+
+
+def define_tool(
+    *,
+    name: str,
+    description: str,
+    schema: Type[BaseModel],
+    execute: Callable[[Any, ToolContext | None], Awaitable[Any]],
+    to_result: Optional[Callable[[Any], Awaitable[ToolResult]]] = None,
+) -> Tool:
+    if not name:
+        raise ValueError("Tool must have a name")
+    return Tool(
+        name=name,
+        description=description,
+        schema=schema,
+        execute=execute,
+        to_result=to_result or _default_to_result,
+    )
+
+
+def parameters_schema(tool: Tool) -> dict[str, Any] | None:
+    """JSON Schema for the wire — tools without it are advertised zero-arg.
+
+    Parity with the TS fix: the model must see real argument shapes.
+    """
+    try:
+        schema = tool.schema.model_json_schema()
+        return schema
+    except Exception:  # pragma: no cover - exotic pydantic setups
+        return None
+
+
+class ToolRegistry:
+    """Validating registry — single place tools are looked up + arg-checked."""
+
+    def __init__(self) -> None:
+        self._tools: dict[str, Tool] = {}
+
+    def register(self, tool: Tool) -> "ToolRegistry":
+        if tool.name in self._tools:
+            raise ValueError(f"Tool already registered: {tool.name}")
+        self._tools[tool.name] = tool
+        return self
+
+    def get(self, name: str) -> Optional[Tool]:
+        return self._tools.get(name)
+
+    def list(self) -> list[Tool]:
+        return list(self._tools.values())
+
+    async def call(self, name: str, raw_args: Any, ctx: ToolContext | None = None) -> Any:
+        """Validate args and execute. ``ctx`` is optional (None stays None —
+        callers may pass a ToolContext so registry calls don't force None on
+        tools that read ctx)."""
+        tool = self._tools.get(name)
+        if tool is None:
+            raise AgentError("UNKNOWN_TOOL", f"Unknown tool: {name}", retryable=False)
+        try:
+            args = tool.schema.model_validate(raw_args if isinstance(raw_args, dict) else {})
+        except ValidationError as err:
+            raise AgentError("TOOL_ARG_INVALID", f'Invalid args for tool "{name}": {err}', retryable=False) from err
+        return await tool.execute(tool.schema.model_validate(args.model_dump()), ctx)

framework_agents/tracing.py

@@ -0,0 +1,131 @@
+"""OTel-shaped tracing exporter. Parity: src/tracing/index.ts (span half).
+
+Spans from AgentEvents, sunk to console/JSONL/memory. No OTel SDK — the
+Span shape is transport-compatible, the sinks are local-first.
+
+Pinned (TS post-fix):
+- Spans are emitted DURATION-COMPLETE (started_at + ended_at both set at
+  emit time): no dangling open-span bookkeeping.
+- Span ids are unique hex (secrets) — sequential per sink, no name-keyed map.
+- run.done metrics + tool.result parent linkage. CAVEAT (parity with the TS
+  caveat): the parent is the LAST tool.call span emitted on THIS instance —
+  there is no run-root, so concurrent interleaved runs through one sink can
+  cross-link a tool.result to another run's tool.call. Fixing that (a run-root
+  span) is pending in the TS track too; keep it simple here.
+- Coverage: ALL events (step/token/thinking/tool calls) — the TS bug lost
+  everything but tool calls / partial data; emit spans for every hook.
+"""
+
+from __future__ import annotations
+
+import json
+import secrets
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Optional, Protocol
+
+
+@dataclass
+class Span:
+    trace_id: str
+    span_id: str
+    name: str
+    started_at: str  # ISO timestamp
+    ended_at: Optional[str] = None  # ISO timestamp; duration-complete spans set both
+    parent_id: Optional[str] = None
+    attributes: dict[str, Any] = field(default_factory=dict)
+
+
+class SpanSink(Protocol):
+    def emit(self, span: Span) -> None: ...
+
+
+class ConsoleSpanSink:
+    def emit(self, span: Span) -> None:
+        print(f"[trace] {span.name}", json.dumps(span.attributes))
+
+
+class JsonlSpanSink:
+    """Per-span JSON lines, appending; the parent dir is created (mkdir -p)."""
+
+    def __init__(self, path: str) -> None:
+        self.path = path
+        Path(path).parent.mkdir(parents=True, exist_ok=True)
+
+    def emit(self, span: Span) -> None:
+        record = {
+            "traceId": span.trace_id,
+            "spanId": span.span_id,
+            "parentId": span.parent_id,
+            "name": span.name,
+            "startedAt": span.started_at,
+            "endedAt": span.ended_at,
+            "attributes": span.attributes,
+        }
+        with open(self.path, "a", encoding="utf-8") as handle:
+            handle.write(json.dumps(record) + "\n")
+
+
+class MemorySpanSink:
+    def __init__(self) -> None:
+        self.spans: list[Span] = []
+
+    def emit(self, span: Span) -> None:
+        self.spans.append(span)
+
+
+def _now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+class TracingEvents:
+    """AgentEvents-shaped projections of run/step/tool lifecycle into spans."""
+
+    def __init__(self, sink: SpanSink, trace_id: Optional[str] = None) -> None:
+        self.sink = sink
+        self.trace_id = trace_id or secrets.token_hex(16)
+        # Sequential unique span ids only — no name-keyed map on this instance
+        # beyond the LAST tool.call (name-keyed maps leak across runs).
+        self._last_tool_call_span_id: Optional[str] = None
+
+    def _emit(self, name: str, attributes: dict[str, Any], parent_id: Optional[str] = None) -> Span:
+        now = _now_iso()
+        span = Span(
+            trace_id=self.trace_id,
+            span_id=secrets.token_hex(8),  # 16 hex chars, unique per emission
+            parent_id=parent_id,
+            name=name,
+            started_at=now,
+            ended_at=now,  # duration-complete at emit time
+            attributes=attributes,
+        )
+        self.sink.emit(span)
+        return span
+
+    def on_step(self, step: int) -> None:
+        self._emit("step.start", {"step": step})
+
+    def on_token(self, delta: str) -> None:
+        self._emit("model.token", {"delta": delta[:200]})
+
+    def on_thinking(self, delta: str) -> None:
+        self._emit("model.thinking", {"delta": delta[:200]})
+
+    def on_tool_call(self, name: str, args: Any) -> None:
+        span = self._emit("tool.call", {"tool": name, "args": args})
+        self._last_tool_call_span_id = span.span_id
+
+    def on_tool_result(self, name: str, output: Any) -> None:
+        # Parent = the most recent tool.call span on this instance (see the
+        # module caveat: no run root, so concurrent runs through one sink can
+        # cross-link — pending in the TS track as well).
+        self._emit("tool.result", {"tool": name, "output": output}, parent_id=self._last_tool_call_span_id)
+
+    def on_done(self, text: str, tool_calls: Any = None) -> None:
+        count = len(tool_calls) if tool_calls is not None else 0
+        self._emit("run.done", {"text": text[:500], "toolCalls": count})
+
+
+def tracing_events(sink: SpanSink, trace_id: Optional[str] = None) -> TracingEvents:
+    return TracingEvents(sink=sink, trace_id=trace_id)

framework_agents/types.py

@@ -0,0 +1,262 @@
+"""Message + usage models. Parity: packages/typescript/src/types/*.
+
+Pinned semantics (these are the FIXED ones — the python port encodes them
+from day one, not the bugs the TS sweep found):
+
+- TokenUsage.reasoning_tokens is a SUBSET of output_tokens on every supported
+  provider (OpenAI reasoning_tokens C completion_tokens; Together/Anthropic
+  think inside completion/output tokens). Costs/budgets NEVER add it on top.
+- cost_usd is set only when a price is known — never guessed (never a fake $0).
+- run()-created user turns carry no id (editMessage parity addresses by index).
+"""
+
+from __future__ import annotations
+
+import itertools
+import json
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator
+from pydantic import TypeAdapter
+
+Role = Literal["system", "user", "assistant", "tool"]
+
+
+class TokenUsage(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+
+    input_tokens: int = Field(default=0, serialization_alias="inputTokens", validation_alias="inputTokens")
+    output_tokens: int = Field(default=0, serialization_alias="outputTokens", validation_alias="outputTokens")
+    reasoning_tokens: int | None = Field(
+        default=None,
+        description="SUBSET of output_tokens on every supported provider: report it, never add it",
+        serialization_alias="reasoningTokens",
+        validation_alias="reasoningTokens",
+    )
+    cost_usd: float | None = Field(
+        default=None,
+        description="USD cost when a price is known (see pricing). Never guessed.",
+        serialization_alias="costUsd",
+        validation_alias="costUsd",
+    )
+
+    def __add__(self, other: "TokenUsage") -> "TokenUsage":
+        return TokenUsage(
+            input_tokens=self.input_tokens + other.input_tokens,
+            output_tokens=self.output_tokens + other.output_tokens,
+            reasoning_tokens=(self.reasoning_tokens or 0) + (other.reasoning_tokens or 0),
+        )
+
+
+class ToolCall(BaseModel):
+    id: str
+    name: str
+    args: Any = None
+
+
+class ExecutedToolCall(BaseModel):
+    id: str
+    name: str
+    args: Any = None
+    output: Any = None
+
+
+class TextBlock(BaseModel):
+    type: Literal["text"] = "text"
+    text: str
+
+
+class ReasoningBlock(BaseModel):
+    type: Literal["reasoning"] = "reasoning"
+    reasoning: str
+
+
+class ImageBlock(BaseModel):
+    type: Literal["image"] = "image"
+    url: str | None = None
+    data: str | None = None
+    mime_type: str | None = Field(default=None, serialization_alias="mimeType", validation_alias="mimeType")
+    file_id: str | None = Field(default=None, serialization_alias="fileId", validation_alias="fileId")
+
+
+class AudioBlock(BaseModel):
+    type: Literal["audio"] = "audio"
+    url: str | None = None
+    data: str | None = None
+    mime_type: str | None = Field(default=None, serialization_alias="mimeType", validation_alias="mimeType")
+    file_id: str | None = Field(default=None, serialization_alias="fileId", validation_alias="fileId")
+
+
+class FileBlock(BaseModel):
+    type: Literal["file"] = "file"
+    url: str | None = None
+    data: str | None = None
+    mime_type: str | None = Field(default=None, serialization_alias="mimeType", validation_alias="mimeType")
+    file_id: str | None = Field(default=None, serialization_alias="fileId", validation_alias="fileId")
+
+
+class ToolCallBlock(BaseModel):
+    type: Literal["tool_call"] = "tool_call"
+    id: str
+    name: str
+    args: Any = None
+
+
+class NonStandardBlock(BaseModel):
+    type: Literal["non_standard"] = "non_standard"
+    value: Any = None
+
+
+ContentBlock = (
+    TextBlock
+    | ReasoningBlock
+    | ImageBlock
+    | AudioBlock
+    | FileBlock
+    | ToolCallBlock
+    | NonStandardBlock
+)
+
+_BLOCK_TYPES: dict[str, type[BaseModel]] = {
+    "text": TextBlock,
+    "reasoning": ReasoningBlock,
+    "image": ImageBlock,
+    "audio": AudioBlock,
+    "file": FileBlock,
+    "tool_call": ToolCallBlock,
+    "non_standard": NonStandardBlock,
+}
+
+
+def parse_content_block(raw: Any) -> BaseModel:
+    if not isinstance(raw, dict):
+        # A bare string (or any other shape) in a content list used to crash
+        # later text_of() work with AttributeError; fail loudly, naming it.
+        raise ValueError(
+            f"content list elements must be block dicts/models (got {type(raw).__name__})"
+        )
+    kind = raw.get("type")
+    cls = _BLOCK_TYPES.get(str(kind))
+    if cls is None:
+        raise ValueError(f"Unknown content block type: {kind!r}")
+    return cls.model_validate(raw)
+
+
+# Parsed blocks are memoized per message INSTANCE (id-keyed): text_of /
+# wire_text_of / providers re-parse the same content repeatedly per run.
+# The entry is cleared whenever the model fields are (re)validated, and a
+# new instance parses fresh (a fresh id never collides with a live entry).
+_BLOCKS_CACHE: dict[int, list[BaseModel]] = {}
+
+
+class ChatMessage(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+
+    id: str | None = None
+    role: Role
+    content: str | list[Any]  # string or parsed ContentBlock models (validated lazily)
+    name: str | None = None
+    tool_call_id: str | None = Field(default=None, serialization_alias="toolCallId", validation_alias="toolCallId")
+    tool_name: str | None = Field(default=None, serialization_alias="toolName", validation_alias="toolName")
+    pinned: bool | None = None
+
+    @model_validator(mode="after")
+    def _clear_block_cache(self) -> "ChatMessage":
+        _BLOCKS_CACHE.pop(id(self), None)
+        return self
+
+    def blocks(self) -> list[BaseModel]:
+        """Parsed content blocks (cached per instance; str content -> [])."""
+        if isinstance(self.content, str):
+            return []
+        cached = _BLOCKS_CACHE.get(id(self))
+        if cached is not None:
+            return list(cached)
+        parsed: list[BaseModel] = []
+        for block in self.content:
+            if isinstance(block, BaseModel):
+                parsed.append(block)
+            else:
+                parsed.append(parse_content_block(block))
+        _BLOCKS_CACHE[id(self)] = parsed
+        return list(parsed)
+
+
+_message_seq = itertools.count(1)
+
+
+def system(content: str | list[Any]) -> ChatMessage:
+    return ChatMessage(id=f"msg_{next(_message_seq)}", role="system", content=content)
+
+
+def user(content: str | list[Any], id: str | None = None, name: str | None = None) -> ChatMessage:
+    msg_id = id or f"msg_{next(_message_seq)}"
+    return ChatMessage(id=msg_id, role="user", content=content, name=name)
+
+
+def assistant(content: str | list[Any]) -> ChatMessage:
+    return ChatMessage(id=f"msg_{next(_message_seq)}", role="assistant", content=content)
+
+
+def tool_result(tool_call_id: str, tool_name: str, content: str | list[Any]) -> ChatMessage:
+    return ChatMessage(
+        id=f"msg_{next(_message_seq)}",
+        role="tool",
+        content=content,
+        tool_call_id=tool_call_id,
+        tool_name=tool_name,
+    )
+
+
+_CHAT_MESSAGE_LIST = TypeAdapter(list[ChatMessage])
+
+
+def serialize_messages(messages: list[ChatMessage]) -> list[dict[str, Any]]:
+    """JSON-safe history for persistence."""
+    return [m.model_dump(by_alias=True, exclude_none=True) for m in messages]
+
+
+def deserialize_messages(raw: list[dict[str, Any]]) -> list[ChatMessage]:
+    """Validated history restore. Raises on bad shape (never silently drops)."""
+    try:
+        return _CHAT_MESSAGE_LIST.validate_python(raw)
+    except ValidationError as err:
+        raise ValueError(f"Invalid message history: {err}") from err
+
+
+def text_of(content: str | list[Any]) -> str:
+    """Plain-text view of any content shape."""
+    if isinstance(content, str):
+        return content
+    out: list[str] = []
+    for b in content if isinstance(content, list) else []:
+        block = b if isinstance(b, BaseModel) else parse_content_block(b)
+        if isinstance(block, TextBlock):
+            out.append(block.text)
+        elif isinstance(block, ReasoningBlock):
+            out.append(block.reasoning)
+    return "".join(out)
+
+
+def wire_text_of(content: str | list[Any]) -> str:
+    """Text as re-sent to a provider: reasoning blocks are DROPPED
+    (a model's own past thinking must not come back as ordinary text)."""
+    if isinstance(content, str):
+        return content
+    out: list[str] = []
+    for b in content if isinstance(content, list) else []:
+        block = b if isinstance(b, BaseModel) else parse_content_block(b)
+        if isinstance(block, TextBlock):
+            out.append(block.text)
+    return "".join(out)
+
+
+def has_images(content: str | list[Any]) -> bool:
+    return isinstance(content, list) and any(
+        (b if isinstance(b, dict) else b.model_dump()).get("type") == "image"
+        for b in content
+    )
+
+
+def to_json_shallow(value: Any) -> str:
+    return json.dumps(value, default=str)

framework_agents/workflows.py

@@ -0,0 +1,130 @@
+"""DAG workflows. Parity: packages/typescript/src/workflows/index.ts.
+
+Topological levels run in order; nodes within a level run concurrently.
+Duplicate depends entries ("a", "a") are ONE edge — counting them twice
+wedges indegree at 1 and misreports a legal graph as a cycle (pinned TS
+regression). Abort exceptions always rethrow, never retry.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import Any, Awaitable, Callable, Optional
+
+from .errors import AgentError
+
+
+@dataclass
+class RetryPolicy:
+    attempts: int = 1
+    backoff_ms: float = 0
+
+
+@dataclass
+class WorkflowNode:
+    id: str
+    run: Callable[[dict[str, Any], Optional[dict[str, Any]]], Awaitable[Any]]
+    depends_on: list[str] = field(default_factory=list)
+    """Retry policy. Default: no retry (attempts=1). Abort always rethrows."""
+    retry: Optional[RetryPolicy] = None
+
+
+@dataclass
+class WorkflowOptions:
+    """Shared per-run data passed to every node (never sent to models by itself)."""
+    invocation_state: Optional[dict[str, Any]] = None
+
+
+def _sort_topo(nodes: list[WorkflowNode]) -> list[list[WorkflowNode]]:
+    by_id: dict[str, WorkflowNode] = {}
+    for n in nodes:
+        if n.id in by_id:
+            raise AgentError("WORKFLOW_DUPLICATE_ID", "Workflow node ids must be unique", retryable=False)
+        by_id[n.id] = n
+    for n in nodes:
+        for dep in n.depends_on or []:
+            if dep not in by_id:
+                raise AgentError("WORKFLOW_UNKNOWN_DEP", f'Node "{n.id}" depends on unknown node "{dep}"', retryable=False)
+            if dep == n.id:
+                raise AgentError("WORKFLOW_CYCLE", f'Node "{n.id}" depends on itself', retryable=False)
+    # Deduplicate first: what matters for the indegree is the number of
+    # DISTINCT dependencies, not the raw list length.
+    indegree = {n.id: len(set(n.depends_on or [])) for n in nodes}
+    levels: list[list[WorkflowNode]] = []
+    ready = [n for n in nodes if indegree[n.id] == 0]
+    while ready:
+        levels.append(ready)
+        ready_ids = {n.id for n in ready}
+        deps: dict[str, set[str]] = {m.id: set(m.depends_on or []) for m in nodes}
+        nxt: list[WorkflowNode] = []
+        for done in nodes:
+            if done.id not in ready_ids:
+                continue
+            for m in nodes:
+                if done.id not in deps[m.id]:
+                    continue
+                indegree[m.id] -= 1
+                if indegree[m.id] == 0:
+                    nxt.append(m)
+        ready = nxt
+    if sum(len(level) for level in levels) != len(nodes):
+        raise AgentError("WORKFLOW_CYCLE", "Workflow has a dependency cycle", retryable=False)
+    return levels
+
+
+def _policy_of(node: WorkflowNode) -> RetryPolicy:
+    retry = node.retry
+    if isinstance(retry, dict):
+        return RetryPolicy(**retry)
+    if retry is None:
+        return RetryPolicy()
+    return retry
+
+
+async def _run_with_retry(
+    node: WorkflowNode,
+    ctx: dict[str, Any],
+    invocation_state: Optional[dict[str, Any]],
+) -> Any:
+    policy = _policy_of(node)
+    attempts = max(1, policy.attempts or 1)
+    backoff_ms = policy.backoff_ms or 0
+    last: BaseException | None = None
+    for i in range(attempts):
+        try:
+            return await node.run(ctx, invocation_state)
+        except asyncio.CancelledError:
+            raise
+        except AgentError as err:
+            if err.code == "ABORTED":
+                raise
+            last = err
+        except Exception as err:  # noqa: BLE001 - node body errors retry per policy
+            last = err
+        if i < attempts - 1 and backoff_ms > 0:
+            await asyncio.sleep(backoff_ms / 1000)
+    assert last is not None  # pragma: no cover - loop always sets it on fallthrough
+    raise last
+
+
+async def run_workflow(
+    nodes: list[WorkflowNode],
+    opts: Optional[WorkflowOptions] = None,
+) -> dict[str, Any]:
+    """Execute a DAG. Returns id -> result. Errors (incl. aborts) propagate."""
+    invocation_state = opts.invocation_state if opts is not None else None
+    ctx: dict[str, Any] = {}
+    for level in _sort_topo(nodes):
+        values = await asyncio.gather(*(_run_with_retry(n, ctx, invocation_state) for n in level))
+        for node, value in zip(level, values):
+            ctx[node.id] = value
+    return ctx
+
+
+__all__ = [
+    "RetryPolicy",
+    "WorkflowNode",
+    "WorkflowOptions",
+    "run_workflow",
+]