tyrion-cli 0.1.0__py3-none-any.whl

tyrion_agent/events.py

@@ -0,0 +1,89 @@
+"""Events emitted by the portable agent layer."""
+
+from __future__ import annotations
+from typing import Annotated, Literal
+
+from pydantic import Field
+
+from tyrion_agent.messages import AgentMessage, ToolResultMessage, WireModel
+from tyrion_agent.provider_events import AssistantMessageEvent
+from tyrion_agent.tools import AgentToolResult
+from tyrion_agent.types import JSONValue
+
+class AgentStartEvent(WireModel):
+    """The agent has started processing."""
+    type: Literal["agent_start"] = "agent_start"
+
+class AgentEndEvent(WireModel):
+    """The agent has finished processing."""
+
+    type: Literal["agent_end"] = "agent_end"
+    messages: list[AgentMessage] = Field(default_factory=list)
+
+class TurnStartEvent(WireModel):
+    """A new model turn has started."""
+
+    type: Literal["turn_start"] = "turn_start"
+
+class TurnEndEvent(WireModel):
+    """A model turn has completed."""
+
+    type: Literal["turn_end"] = "turn_end"
+    message: AgentMessage
+    tool_results: list[ToolResultMessage] = Field(default_factory=list)
+
+class MessageStartEvent(WireModel):
+    """A message has started (user, assistant or tool result)."""
+    type: Literal["message_start"] = "message_start"
+    message: AgentMessage
+
+class MessageUpdateEvent(WireModel):
+    """A streaming delta has arrived for the current message."""
+
+    type: Literal["message_update"] = "message_update"
+    message: AgentMessage
+    assistant_message_event: AssistantMessageEvent
+
+class MessageEndEvent(WireModel):
+    """A message has been finalized."""
+    type: Literal["message_end"] = "message_end"
+    message: AgentMessage
+
+class ToolExecutionStartEvent(WireModel):
+    """A tool has started executing."""
+
+    type: Literal["tool_execution_start"] = "tool_execution_start"
+    tool_call_id: str
+    tool_name : str
+    args: dict[str, JSONValue] = Field(default_factory=dict)
+
+class ToolExecutionUpdateEvent(WireModel):
+    """A tool has started executing."""
+
+    type: Literal["tool_execution_update"] = "tool_execution_update"
+    tool_call_id: str
+    tool_name : str
+    args: dict[str, JSONValue] = Field(default_factory=dict)
+    partial_result: AgentToolResult
+
+class ToolExecutionEndEvent(WireModel):
+    """A tool has finished executing."""
+    type: Literal["tool_execution_end"] = "tool_execution_end"
+    tool_call_id: str
+    tool_name: str
+    result: AgentToolResult
+    is_error: bool
+# Union of all agent events
+type AgentEvent = Annotated[
+    AgentStartEvent
+    | AgentEndEvent
+    | TurnStartEvent
+    | TurnEndEvent
+    | MessageStartEvent
+    | MessageUpdateEvent
+    | MessageEndEvent
+    | ToolExecutionStartEvent
+    | ToolExecutionUpdateEvent
+    | ToolExecutionEndEvent,
+    Field(discriminator="type"),
+]

tyrion_agent/harness.py

@@ -0,0 +1,273 @@
+"""Stateful reusable agent brain built on top of the pure agent loop."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
+from contextlib import suppress
+from dataclasses import dataclass, field
+from inspect import isawaitable
+
+from tyrion_agent.events import AgentEvent, MessageEndEvent, MessageStartEvent
+from tyrion_agent.loop import run_agent_loop
+from tyrion_agent.messages import (
+    AgentMessage,
+    AssistantMessage,
+    TextContent,
+    ToolResultMessage,
+    UserMessage,
+)
+from tyrion_agent.provider import ModelProvider
+from tyrion_agent.tools import AgentTool
+
+# A listener is any function that receives an event (sync or async)
+EventListener = Callable[[AgentEvent], Awaitable[None] | None]
+
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+
+@dataclass(slots=True)
+class AgentHarnessConfig:
+    """Everything the harness needs to run the loop."""
+
+    provider: ModelProvider
+    model: str
+    system: str
+    tools: list[AgentTool] = field(default_factory=list)
+    max_turns: int | None = None
+
+
+# ---------------------------------------------------------------------------
+# Cancellation
+# ---------------------------------------------------------------------------
+
+
+class SimpleCancellationToken:
+    """A simple flag that can be checked to see if work should stop."""
+
+    def __init__(self) -> None:
+        self._cancelled = False
+
+    def cancel(self) -> None:
+        self._cancelled = True
+
+    def is_cancelled(self) -> bool:
+        return self._cancelled
+
+
+# ---------------------------------------------------------------------------
+# The Harness
+# ---------------------------------------------------------------------------
+
+
+class AgentHarness:
+    """Reusable stateful agent brain.
+
+    Owns the transcript, delegates to the pure agent loop,
+    and provides a clean API for frontends.
+
+    Usage:
+        harness = AgentHarness(AgentHarnessConfig(
+            provider=provider,
+            model="gpt-4",
+            system="You are a coding agent.",
+            tools=[read_tool, bash_tool],
+        ))
+
+        async for event in harness.prompt("Read README.md"):
+            print(event)
+    """
+
+    def __init__(
+        self,
+        config: AgentHarnessConfig,
+        *,
+        messages: Sequence[AgentMessage] = (),
+    ) -> None:
+        self._config = config
+        self._messages: list[AgentMessage] = list(messages)
+        self._listeners: list[EventListener] = []
+        self._current_signal: SimpleCancellationToken | None = None
+        self._running = False
+
+    # --- Public properties ---
+
+    @property
+    def messages(self) -> tuple[AgentMessage, ...]:
+        """Immutable snapshot of the current transcript."""
+        return tuple(self._messages)
+
+    @property
+    def config(self) -> AgentHarnessConfig:
+        return self._config
+
+    @property
+    def is_running(self) -> bool:
+        return self._running
+
+    # --- Core API ---
+
+    def prompt(self, content: str) -> AsyncIterator[AgentEvent]:
+        """Send a user message and run the agent loop.
+
+        This is the main entry point. It:
+        1. Appends a UserMessage to the transcript
+        2. Runs the agent loop
+        3. Yields events as they happen
+        """
+        return self.prompt_message(UserMessage(content=content))
+
+    def prompt_message(self, message: AgentMessage) -> AsyncIterator[AgentEvent]:
+        """Send an arbitrary message and run the agent loop."""
+        self._ensure_not_running()
+        self._running = True
+        return self._run(prompts=(message,))
+
+    def continue_(self) -> AsyncIterator[AgentEvent]:
+        """Run the agent loop without appending a new message.
+
+        Useful for:
+        - Resuming after restoring a session
+        - Continuing after an interrupted run
+        """
+        self._ensure_not_running()
+        self._running = True
+        return self._run()
+
+    # --- Cancellation ---
+
+    def cancel(self) -> None:
+        """Cancel the current run.
+
+        The loop will stop after the current tool finishes.
+        Any unanswered tool calls get synthetic error results.
+        """
+        if self._current_signal is not None:
+            self._current_signal.cancel()
+
+    # --- Event listeners ---
+
+    def subscribe(self, listener: EventListener) -> Callable[[], None]:
+        """Subscribe to agent events. Returns an unsubscribe function.
+
+        Listeners receive the same events as the prompt()/continue_()
+        consumer. This lets persistence, logging, and other observers
+        watch runs without being the main consumer.
+
+        Usage:
+            def on_event(event):
+                print(f"Got: {event.type}")
+
+            unsubscribe = harness.subscribe(on_event)
+            # ... later ...
+            unsubscribe()
+        """
+        self._listeners.append(listener)
+
+        def unsubscribe() -> None:
+            with suppress(ValueError):
+                self._listeners.remove(listener)
+
+        return unsubscribe
+
+    # --- Transcript management ---
+
+    def append_message(self, message: AgentMessage) -> None:
+        """Manually append a message (used when restoring sessions)."""
+        self._messages.append(message)
+
+    def replace_messages(self, messages: Sequence[AgentMessage]) -> None:
+        """Replace the entire transcript (used when restoring sessions)."""
+        self._messages = list(messages)
+
+    # --- Internal machinery ---
+
+    async def _run(
+        self,
+        *,
+        prompts: Sequence[AgentMessage] = (),
+    ) -> AsyncIterator[AgentEvent]:
+        """Run the agent loop and handle cleanup."""
+        signal = SimpleCancellationToken()
+        self._current_signal = signal
+
+        try:
+            # Repair any dangling tool calls from a previous interrupted run
+            self._repair_interrupted_tools()
+
+            # Run the loop
+            async for event in run_agent_loop(
+                provider=self._config.provider,
+                model=self._config.model,
+                system=self._config.system,
+                messages=self._messages,
+                prompts=prompts,
+                tools=self._config.tools,
+                max_turns=self._config.max_turns,
+                signal=signal,
+            ):
+                # Notify all listeners
+                await self._notify(event)
+                # Yield to the consumer (CLI, TUI, etc.)
+                yield event
+
+        finally:
+            # If we were cancelled, repair any new dangling tool calls
+            if signal.is_cancelled():
+                before = len(self._messages)
+                self._repair_interrupted_tools()
+                # Push repair messages to listeners so persistence sees them
+                for message in self._messages[before:]:
+                    with suppress(Exception):
+                        await self._notify(MessageStartEvent(message=message))
+                        await self._notify(MessageEndEvent(message=message))
+
+            if self._current_signal is signal:
+                self._current_signal = None
+            self._running = False
+
+    async def _notify(self, event: AgentEvent) -> None:
+        """Send an event to all listeners."""
+        for listener in list(self._listeners):
+            result = listener(event)
+            if isawaitable(result):
+                await result
+
+    def _ensure_not_running(self) -> None:
+        """Prevent overlapping runs."""
+        if self._running:
+            raise RuntimeError(
+                "AgentHarness is already running. "
+                "Wait for the current run to finish before starting a new one."
+            )
+
+    def _repair_interrupted_tools(self) -> None:
+        """Append synthetic error results for any unanswered tool calls.
+
+        This happens when a run is cancelled mid-tool-execution.
+        The model asked to call a tool, but we stopped before getting
+        the result. We need to add a "Tool call interrupted" result
+        so the transcript stays valid.
+        """
+        returned_ids: set[str] = {
+            msg.tool_call_id
+            for msg in self._messages
+            if isinstance(msg, ToolResultMessage)
+        }
+
+        for msg in tuple(self._messages):
+            if not isinstance(msg, AssistantMessage):
+                continue
+            for call in msg.tool_calls:
+                if call.id not in returned_ids:
+                    returned_ids.add(call.id)
+                    self._messages.append(
+                        ToolResultMessage(
+                            tool_call_id=call.id,
+                            tool_name=call.name,
+                            content=[TextContent(text="Tool call interrupted by user")],
+                            is_error=True,
+                        )
+                    )

tyrion_agent/loop.py

@@ -0,0 +1,273 @@
+"""the pure agent loop - the engine that drives model <-> tool interaction."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator, Mapping, Sequence
+
+from tyrion_agent.events import (
+    AgentEndEvent,
+    AgentEvent,
+    AgentStartEvent,
+    MessageEndEvent,
+    MessageStartEvent,
+    MessageUpdateEvent,
+    ToolExecutionEndEvent,
+    ToolExecutionStartEvent,
+    TurnEndEvent,
+    TurnStartEvent,
+)
+from tyrion_agent.messages import (
+    AgentMessage,
+    AssistantMessage,
+    TextContent,
+    ToolCall,
+    ToolResultMessage,
+)
+from tyrion_agent.provider import CancellationToken, ModelProvider
+from tyrion_agent.provider_events import (
+    AssistantErrorEvent,
+    AssistantMessageEvent,
+    AssistantStartEvent,
+    AssistantDoneEvent
+)
+from tyrion_agent.tool_history import repair_tool_history
+from tyrion_agent.tools import AgentTool, AgentToolResult
+
+async def run_agent_loop(
+        *,
+        provider: ModelProvider,
+        model: str,
+        system: str,
+        messages :list[AgentMessage],
+        tools : list[AgentTool],
+        prompts: Sequence[AgentMessage] =(),
+        max_turns: int | None = None,
+        signal: CancellationToken | None = None,
+) -> AsyncIterator[AgentEvent]:
+    """Run the provider/tool loop and emit the agent events.
+    Args:
+        provider: The model provider to stream responses from.
+        model: Model name to use.
+        system: System prompt text.
+        messages: The mutable transcript. The loop APPENDS to this list.
+        tools: Available tools the model can call.
+        prompts: Initial messages to append (e.g. a new user message).
+        max_turns: Optional safety cap on model turns.
+        signal: Optional cancellation token.
+    """
+    new_messages: list[AgentMessage] = list(prompts)
+    if prompts:
+        messages.extend(prompts)
+
+    tool_by_name = {tool.name: tool for tool in tools}
+    turn =1
+
+    ## Emit start events
+    yield AgentStartEvent()
+    yield TurnStartEvent()
+
+    # Emit events for initial prompts messages
+    for prompt in prompts:
+        yield MessageStartEvent(message = prompt)
+        yield MessageEndEvent(message = prompt)
+
+    #Validate max_turns
+    if max_turns is not None and max_turns < 1:
+        error = _error_message(model, "max turns must be atleat 1")
+        messages.append(error)
+        new_messages.append(error)
+        yield MessageStartEvent(message=error)
+        yield MessageEndEvent(message=error)
+        yield TurnEndEvent(message=error)
+        yield AgentEndEvent(messages=new_messages)
+        return
+
+    ## --- Main Loop ---
+    while True:
+        # check max turns
+        if max_turns is not None and turn > max_turns:
+            error = _error_message(model, f"Reached the max turns limit ({max_turns})")
+            messages.append(error)
+            new_messages.append(error)
+            yield MessageStartEvent(message=error)
+            yield MessageEndEvent(message=error)
+            yield TurnEndEvent(message=error)
+            yield AgentEndEvent(messages=new_messages)
+            return
+        assistant : AssistantMessage | None = None
+
+        async for event in _stream_assistant(
+            provider = provider,
+            model = model,
+            system = system,
+            messages = messages,
+            tools =tools,
+            signal= signal,
+        ):
+            yield event
+            # capture the final assistant message
+
+            if isinstance(event, MessageEndEvent) and isinstance(event.message, AssistantMessage):
+                assistant = event.message
+
+        # The provider should always produce a message, but incase it doesn't
+        if assistant is None:
+            assistant = _error_message(model, "Provider produced no assistant message")
+            yield MessageStartEvent(message=assistant)
+            yield MessageEndEvent(message=assistant)
+
+        ## Add assistant message to transcript
+        messages.append(assistant)
+        new_messages.append(assistant)
+
+        if assistant.stop_reason in {"error", "aborted"}:
+            yield TurnEndEvent(message=assistant)
+            yield AgentEndEvent(messages=new_messages)
+            return
+        ## Execute tool calls (if any)
+        calls = list(assistant.tool_calls)
+        tool_results: list[ToolResultMessage] =[]
+
+        for call in calls:
+            async for event in _execute_tool_call(call, tool_by_name, signal):
+                yield event
+
+                if isinstance(event, MessageEndEvent) and isinstance(event.message, ToolResultMessage):
+                    tool_results.append(event.message)
+                    messages.append(event.message)
+                    new_messages.append(event.message)
+
+        # Emit turn end
+        yield TurnEndEvent(message = assistant, tool_results=tool_results)
+        turn +=1
+
+        if not calls:
+            break
+
+        # Cancelled while tools were running: stop here instead of asking the
+        # model for another turn.
+        if signal is not None and signal.is_cancelled():
+            yield AgentEndEvent(messages=new_messages)
+            return
+        yield TurnStartEvent()
+    yield AgentEndEvent(messages=new_messages)
+
+
+### Stream assistant response from porvider
+
+async def _stream_assistant(
+        *,
+        provider:  ModelProvider,
+        model: str,
+        system : str,
+        messages: list[AgentMessage],
+        tools: list[AgentTool],
+        signal: CancellationToken | None = None,
+) -> AsyncIterator[AgentEvent]:
+    """Strema one model's response and convert provider events to agent events."""
+
+    clean = repair_tool_history(messages)
+
+    source : AsyncIterator[AssistantMessageEvent] = provider.stream_response(
+        model = model,
+        system = system,
+        messages = clean.messages,
+        tools=tools,
+        signal = signal,
+    )
+    started = False
+    async for event in source:
+        if isinstance(event,AssistantStartEvent):
+            started = True
+            yield MessageStartEvent(message = event.partial)
+
+        elif isinstance(event, AssistantDoneEvent):
+            if not started:
+                yield MessageStartEvent(message = event.message)
+            yield MessageEndEvent(message = event.message)
+
+        elif isinstance(event, AssistantErrorEvent):
+            if not started:
+                yield MessageStartEvent(message=event.error)
+            yield MessageEndEvent(message=event.error)
+        else :
+            yield MessageUpdateEvent(
+                message = event.partial,
+                assistant_message_event=event,
+            )
+
+async def _execute_tool_call(
+        call: ToolCall,
+        tools: Mapping[str, AgentTool],
+        signal: CancellationToken | None,
+) -> AsyncIterator[AgentEvent]:
+    """Execute one tool call and emit events."""
+
+    # Announce the tool execution
+    yield ToolExecutionStartEvent(
+        tool_call_id = call.id,
+        tool_name = call.name,
+        args = call.arguments,
+    )
+
+    #check if cancelled
+    if signal is not None and signal.is_cancelled():
+        result = _error_result("Operation aborted")
+        is_error= True
+
+    else:
+        ## Find the tool
+        tool = tools.get(call.name)
+        if tool is None:
+            result = _error_result(f"Unknown Tool: {call.name}")
+            is_error = True
+        else:
+            result, is_error = await _run_tool(tool, call, signal)
+
+    ## Announce the result
+    yield ToolExecutionEndEvent(
+        tool_call_id = call.id,
+        tool_name = call.name,
+        result = result,
+        is_error = is_error
+    )
+
+    # Emit tool result message
+    message = ToolResultMessage(
+        tool_call_id=call.id,
+        tool_name=call.name,
+        content=result.content,
+        is_error=is_error,
+    )
+    yield MessageStartEvent(message=message)
+    yield MessageEndEvent(message=message)
+
+async def _run_tool(
+        tool: AgentTool,
+        call: ToolCall,
+        signal: CancellationToken | None,
+) -> tuple[AgentToolResult, bool]:
+    """Run a tool catching any exceptions"""
+    try: 
+        result = await tool.execute(call.id, call.arguments, signal)
+        return result, False
+    except asyncio.CancelledError:
+        raise
+    except Exception as exc:
+        return _error_result(str(exc)), True
+
+## Helpers
+def _error_result(message: str) -> AgentToolResult:
+    """Create an error tool result."""
+    return AgentToolResult(content=[TextContent(text=message)])
+
+
+def _error_message(model: str, message: str) -> AssistantMessage:
+    """Create an error assistant message."""
+    return AssistantMessage(
+        model=model,
+        content=[],
+        stop_reason="error",
+        error_message=message,
+    )