relay-code 0.1.0__py3-none-any.whl

agent/__init__.py

@@ -0,0 +1,6 @@
+"""Agent module for the Relay AI coding agent."""
+
+from agent.agent import Agent
+from agent.events import AgentEvent, AgentEventType
+
+__all__ = ['Agent', 'AgentEvent', 'AgentEventType']

agent/agent.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+import json
+from typing import AsyncGenerator, Awaitable, Callable
+from agent.session import Session
+from config.config import Config
+from agent.events import AgentEvent, AgentEventType
+from client.response import StreamEventType, ToolCall, ToolResultMessage
+from tools.base import ToolResult
+
+class Agent:
+    def __init__(
+        self,
+        config: Config,
+        confirmation_handler: Callable[[str, dict], Awaitable[bool]] | None = None,
+    ):
+        self.config = config
+        self.confirmation_handler = confirmation_handler
+        self.session: Session | None = Session(self.config)
+
+    async def run(self, message: str):
+        yield AgentEvent.agent_start(message)
+        self.session.reset_turn_usage()
+        self.session.context_manager.add_user_message(message)
+
+        final_response: str | None = None
+        async for event in self._agentic_loop():
+            yield event
+
+            if event.type == AgentEventType.TEXT_COMPLETE:
+                final_response = event.data.get("content")
+
+        yield AgentEvent.agent_end(final_response, self.session.last_usage)
+    
+    async def _agentic_loop(self) -> AsyncGenerator[AgentEvent, None]:
+
+        max_turns = self.config.max_turns
+
+        for i in range(max_turns):
+
+            self.session.inc_turn()
+
+            response_text = ""
+
+            tool_schemas = self.session.tool_registery.get_schemas()
+            tool_calls: list[ToolCall] = []
+            
+            async for event in self.session.client.chat_completion(
+                self.session.context_manager.get_messages(), 
+                tools=tool_schemas if tool_schemas else None, 
+                stream=True
+            ):
+                if event.type == StreamEventType.TEXT_DELTA:
+                    if not event.text_delta:
+                        continue
+                    content = event.text_delta.content
+                    response_text += content
+                    yield AgentEvent.text_delta(content)
+                elif event.type == StreamEventType.TOOL_CALL_COMPLETE:
+                    if event.tool_call:
+                        tool_calls.append(event.tool_call)
+                elif event.type == StreamEventType.MESSAGE_COMPLETE:
+                    if event.usage:
+                        self.session.last_usage = event.usage
+                        self.session.turn_usage += event.usage
+                        yield AgentEvent.usage(self.session.turn_usage)
+                    if event.text_delta:
+                        response_text += event.text_delta.content
+                        yield AgentEvent.text_delta(event.text_delta.content)
+                    if event.tool_calls:
+                        tool_calls.extend(event.tool_calls)
+                elif event.type == StreamEventType.ERROR:
+                    yield AgentEvent.agent_error(event.error or "Unknown error occurred.")
+                    return
+            
+            self.session.context_manager.add_assistant_message(
+                response_text or None,
+                [
+                    {
+                        'id': tc.call_id,
+                        'type': 'function',
+                        'function': {
+                            'name': tc.name,
+                            'arguments': json.dumps(tc.arguments),
+                        },
+                    }
+                    for tc in tool_calls
+                ]
+                if tool_calls
+                else None
+            )
+        
+            if response_text:
+                yield AgentEvent.text_complete(response_text)
+            
+            if not tool_calls:
+                return
+            
+            tool_call_results: list[ToolResultMessage] = []
+
+            for tool_call in tool_calls:
+                yield AgentEvent.tool_call_start(
+                    tool_call.call_id,
+                    tool_call.name,
+                    tool_call.arguments
+                )
+
+                tool = self.session.tool_registery.get(tool_call.name)
+                approved = True
+                if tool and tool.is_mutating(tool_call.arguments):
+                    if self.confirmation_handler is None:
+                        approved = False
+                    else:
+                        approved = await self.confirmation_handler(
+                            tool_call.name,
+                            tool_call.arguments,
+                        )
+
+                if approved:
+                    result = await self.session.tool_registery.invoke(
+                        tool_call.name,
+                        tool_call.arguments,
+                        self.config.cwd,
+                    )
+                else:
+                    result = ToolResult.error_result(
+                        f"Tool execution denied: {tool_call.name}",
+                        metadata={"tool_name": tool_call.name, "denied": True},
+                    )
+
+                yield AgentEvent.tool_call_complete(
+                    tool_call.call_id,
+                    tool_call.name,
+                    result,
+                )
+
+                tool_call_results.append(
+                    ToolResultMessage(
+                        tool_call_id=tool_call.call_id,
+                        content=result.to_model_output(),
+                        is_error=not result.success
+                    )
+                )
+
+            for tool_result in tool_call_results:
+                self.session.context_manager.add_tool_result(
+                    tool_result.tool_call_id,
+                    tool_result.content,
+                )
+    
+    async def __aenter__(self) -> Agent:
+        return self
+    
+    async def __aexit__(
+            self, 
+            exc_type, 
+            exc_val, 
+            exc_tb
+            ) -> None:
+
+        if self.session and self.session.client:
+            await self.session.client.close()
+            self.session = None # Close session, everything else will be garbage collected

agent/events.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+from enum import Enum
+from typing import Any
+from dataclasses import dataclass, field
+from client.response import TokenUsage
+from tools.base import ToolResult
+
+class AgentEventType(str, Enum):
+
+    # Agentic loops, and err
+    AGENT_START="agent_start"
+    AGENT_END="agent_end"
+    AGENT_ERROR="agent_error"
+
+    # Token accounting, emitted once per completion
+    USAGE = "usage"
+
+    # Text streaming
+    TEXT_DELTA = "text_delta"
+    TEXT_COMPLETE = "text_complete"
+
+    # Tool Calls
+    TOOL_CALL_START = 'tool_call_start'
+    TOOL_CALL_COMPLETE = 'tool_call_complete'
+
+@dataclass
+class AgentEvent:
+    type: AgentEventType
+    data: dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def agent_start(
+        cls, 
+        message: str
+    ) -> AgentEvent:
+        return cls(
+            type=AgentEventType.AGENT_START,
+            data={
+                "message": message
+            }
+        )
+
+    @classmethod
+    def agent_end(
+        cls, 
+        response: str | None = None, 
+        usage: TokenUsage | None = None
+    ) -> AgentEvent:
+        return cls(
+            type=AgentEventType.AGENT_END,
+            data={
+                "response": response,
+                'usage': usage.__dict__ if usage else None
+            },
+        )
+    
+    @classmethod
+    def usage(
+        cls,
+        turn_usage: TokenUsage
+    ) -> AgentEvent:
+        return cls(
+            type=AgentEventType.USAGE,
+            data={
+                'usage': turn_usage.__dict__,
+            }
+        )
+
+    @classmethod
+    def agent_error(
+        cls, 
+        error: str, 
+        details: dict[str, Any] | None = None
+    ) -> AgentEvent:
+        return cls(
+            type=AgentEventType.AGENT_ERROR,
+            data={
+                "error": error,
+                'details': details or {} 
+            }
+        )
+
+    @classmethod
+    def text_delta(
+        cls, 
+        content: str
+    ) -> AgentEvent:
+        return cls(
+            type=AgentEventType.TEXT_DELTA,
+            data={
+                "content": content
+            }
+        )
+
+    @classmethod
+    def text_complete(
+        cls, 
+        content: str
+    ) -> AgentEvent:
+        return cls(
+            type=AgentEventType.TEXT_COMPLETE,
+            data={
+                "content": content
+            }
+        )
+    
+    @classmethod
+    def tool_call_start(cls, call_id: str, name: str, arguments: dict[str, Any]):
+        return cls(
+            type=AgentEventType.TOOL_CALL_START,
+            data={
+                'call_id': call_id,
+                'name': name,
+                'arguments': arguments
+            }
+        )
+    
+    @classmethod
+    def tool_call_complete(cls, call_id: str, name: str, result: ToolResult):
+        return cls(
+            type=AgentEventType.TOOL_CALL_COMPLETE,
+            data={
+                'call_id': call_id,
+                'name': name,
+                'success': result.success,
+                'output': result.output,
+                'error': result.error,
+                'metadata': result.metadata,
+                'diff': result.diff.create_diff() if result.diff else None,
+                'truncated': result.truncated,
+                'exit_code': result.exit_code,
+            }
+        )

agent/session.py

@@ -0,0 +1,71 @@
+import json
+import uuid
+
+from client.llm_client import LLMClient
+from client.response import TokenUsage
+from config.config import Config
+from config.loader import get_data_dir
+from context.manager import ContextManager
+from tools.registry import create_default_registery
+from datetime import datetime
+
+class Session:
+    def __init__(self, config: Config):
+        self.config = config
+        self.client = LLMClient(
+            config=config,
+        )
+        self.context_manager = ContextManager(
+            config=config, 
+            user_memory=self._load_memory()
+        )
+        self.tool_registery = create_default_registery(config)
+        self.session_id = str(uuid.uuid4()) # Unique identifiers to resume sessions
+        self.created_at = datetime.now()
+        self.updated_at = datetime.now()
+
+        self.last_usage: TokenUsage | None = None
+
+        self.turn_usage = TokenUsage()
+
+        self._turn_count = 0
+    
+    def _load_memory(self) -> str | None:
+
+        data_dir = get_data_dir()
+
+        data_dir.mkdir(
+            parents=True, 
+            exist_ok=True
+        )
+
+        path = data_dir / 'user_memory.json'
+
+        if not path.exists():
+            return None
+        
+        try:
+            content = path.read_text(
+                encoding='utf-8'
+            )
+            data = json.loads(content)
+            entries = data.get('entries')
+            if not entries:
+                return None
+            
+            lines = ["User prefrences and notes:"]
+            for key, value in entries.items():
+                lines.append(f"- {key}: {value}")
+            
+            return "\n".join(lines)
+        except Exception:
+            return None
+
+    def reset_turn_usage(self) -> None:
+        self.turn_usage = TokenUsage()
+    
+    def inc_turn(self) -> int:
+        self._turn_count += 1
+        self.updated_at = datetime.now()
+
+        return self._turn_count

client/__init__.py

@@ -0,0 +1,23 @@
+"""Client module for LLM interactions."""
+
+from client.llm_client import LLMClient
+from client.response import (
+    StreamEvent,
+    StreamEventType,
+    TextDelta,
+    TokenUsage,
+    ToolCall,
+    ToolCallDelta,
+    ToolResultMessage,
+)
+
+__all__ = [
+    'LLMClient',
+    'StreamEvent',
+    'StreamEventType',
+    'TextDelta',
+    'TokenUsage',
+    'ToolCall',
+    'ToolCallDelta',
+    'ToolResultMessage',
+]

client/llm_client.py

@@ -0,0 +1,236 @@
+import math
+import asyncio
+from openai import AsyncOpenAI, RateLimitError, APIConnectionError, APIError
+from typing import Any, AsyncGenerator
+from client.response import TextDelta, TokenUsage, StreamEvent, StreamEventType, ToolCallDelta, ToolCall, parse_tool_call_arguments
+from config.config import Config
+
+
+class LLMClient:
+    def __init__(self, config: Config) -> None:
+        self._client: AsyncOpenAI | None = None
+        self._max_retries: int = 3
+        self.config = config
+
+    def get_client(self) -> AsyncOpenAI:
+        if self._client is None:
+            # TODO: wire up to your own config system instead of hardcoding
+            # DONE !
+            self._client = AsyncOpenAI(
+                api_key=self.config.api_key,
+                base_url=self.config.base_url # "https://openrouter.ai/api/v1",
+            )
+        return self._client
+
+    async def close(self) -> None:
+        if self._client:
+            await self._client.close()
+            self._client = None
+
+    def _build_tools(
+            self,
+            tools: list[dict[str, Any]]
+    ):
+        return [
+            {
+                'type': 'function',
+                'function': {
+                    'name': tool['name'],
+                    'description': tool.get('description', ""),
+                    'parameters': tool.get('parameters', {
+                        'type': 'object',
+                        'properties': {}
+                    })
+                },
+            }
+            for tool in tools
+        ]
+
+    async def chat_completion(
+            self,
+            messages: list[dict[str, Any]],
+            tools: list[dict[str, Any]] | None = None,
+            stream: bool = True
+    ) -> AsyncGenerator[StreamEvent, None]:
+
+        client = self.get_client()
+
+        kwargs = {
+            "model": self.config.model_name,
+            "messages": messages,
+            "stream": stream,
+        }
+
+        # Without this the stream carries no usage chunk at all, and the
+        # context gauge on the prompt rule has nothing to show.
+        if stream:
+            kwargs["stream_options"] = {"include_usage": True}
+
+        if tools:
+            kwargs['tools'] = self._build_tools(tools)
+            kwargs['tool_choice'] = 'auto'
+
+        for attempt in range(self._max_retries + 1):
+            emitted = False
+            try:
+                if stream:
+                    async for event in self._stream_response(client, kwargs):
+                        emitted = True
+                        yield event
+                else:
+                    event = await self._non_stream_response(client, kwargs)
+                    yield event
+                return
+            except RateLimitError as e:
+                if not emitted and attempt < self._max_retries:
+                    wait_time = math.pow(2, attempt)
+                    await asyncio.sleep(wait_time)
+                else:
+                    yield StreamEvent(
+                        type=StreamEventType.ERROR,
+                        error=f"Rate limit exceeded: {e}",
+                    )
+                    return
+            except APIConnectionError as e:
+                if not emitted and attempt < self._max_retries:
+                    wait_time = math.pow(2, attempt)
+                    await asyncio.sleep(wait_time)
+                else:
+                    yield StreamEvent(
+                        type=StreamEventType.ERROR,
+                        error=f"Connection Error exceeded: {e}",
+                    )
+                    return
+            # No retries here, solely because there's no point in retrying if there is a hardstuck API Error
+            except APIError as e:
+                yield StreamEvent(
+                    type=StreamEventType.ERROR,
+                    error=f"API Error: {e}",
+                )
+                return
+
+    async def _stream_response(self, client: AsyncOpenAI, kwargs: dict[str, Any]) -> AsyncGenerator[StreamEvent, None]:
+        response = await client.chat.completions.create(**kwargs)
+
+        finish_reason: str | None = None
+        usage: TokenUsage | None = None
+        tool_calls: dict[int, dict[str, Any]] = {}
+
+        async for chunk in response:
+            if hasattr(chunk, "usage") and chunk.usage:
+                usage = TokenUsage(
+                    prompt_tokens=chunk.usage.prompt_tokens,
+                    completion_tokens=chunk.usage.completion_tokens,
+                    total_tokens=chunk.usage.total_tokens,
+                    cached_tokens=(
+                        getattr(chunk.usage.prompt_tokens_details, "cached_tokens", 0)
+                        if chunk.usage.prompt_tokens_details else 0
+                    ),
+                )
+
+            if not chunk.choices:
+                continue
+
+            choice = chunk.choices[0]
+            delta = choice.delta
+
+            if choice.finish_reason:
+                finish_reason = choice.finish_reason
+
+            if delta.content:
+                yield StreamEvent(
+                    type=StreamEventType.TEXT_DELTA,
+                    text_delta=TextDelta(delta.content),
+                )
+
+            if delta.tool_calls:
+                for tool_call_delta in delta.tool_calls:
+                    index = tool_call_delta.index
+
+                    if index not in tool_calls:
+                        tool_calls[index] = {
+                            'id': tool_call_delta.id or "",
+                            'name': '',
+                            'arguments': '',
+                        }
+
+                    # NOTE: this used to be nested inside "if index not in tool_calls",
+                    # which meant argument fragments after the first chunk for a given
+                    # index were silently dropped. Now runs on every chunk.
+                    if tool_call_delta.function:
+                        if tool_call_delta.function.name:
+                            tool_calls[index]['name'] = tool_call_delta.function.name
+                            yield StreamEvent(
+                                type=StreamEventType.TOOL_CALL_START,
+                                tool_call_delta=ToolCallDelta(
+                                    call_id=tool_calls[index]['id'],
+                                    name=tool_call_delta.function.name,
+                                )
+                            )
+
+                        if tool_call_delta.function.arguments:
+                            tool_calls[index]['arguments'] += tool_call_delta.function.arguments
+                            yield StreamEvent(
+                                type=StreamEventType.TOOL_CALL_DELTA,
+                                tool_call_delta=ToolCallDelta(
+                                    call_id=tool_calls[index]['id'],
+                                    name=tool_calls[index]['name'],
+                                    arguments_delta=tool_call_delta.function.arguments,
+                                )
+                            )
+
+        for index, toolcall in tool_calls.items():
+            yield StreamEvent(
+                type=StreamEventType.TOOL_CALL_COMPLETE,
+                tool_call=ToolCall(
+                    call_id=toolcall['id'],
+                    name=toolcall['name'],
+                    arguments=parse_tool_call_arguments(toolcall['arguments']),
+                )
+            )
+
+        yield StreamEvent(
+            type=StreamEventType.MESSAGE_COMPLETE,
+            finish_reason=finish_reason,
+            usage=usage,
+        )
+
+    async def _non_stream_response(self, client: AsyncOpenAI, kwargs: dict[str, Any]) -> StreamEvent:
+        response = await client.chat.completions.create(**kwargs)
+        choice = response.choices[0]  # only interested in first index, first message
+        message = choice.message
+
+        text_delta = None
+        if message.content:
+            text_delta = TextDelta(content=message.content)
+
+        tool_calls: list[ToolCall] = []
+        if message.tool_calls:
+            for toolcall in message.tool_calls:
+                tool_calls.append(ToolCall(
+                    call_id=toolcall.id,
+                    name=toolcall.function.name,
+                    arguments=parse_tool_call_arguments(
+                        toolcall.function.arguments
+                    )
+                ))
+
+        usage = None
+        if response.usage:
+            usage = TokenUsage(
+                prompt_tokens=response.usage.prompt_tokens,
+                completion_tokens=response.usage.completion_tokens,
+                total_tokens=response.usage.total_tokens,
+                cached_tokens=(
+                    getattr(response.usage.prompt_tokens_details, "cached_tokens", 0)
+                    if response.usage.prompt_tokens_details else 0
+                ),
+            )
+
+        return StreamEvent(
+            type=StreamEventType.MESSAGE_COMPLETE,
+            text_delta=text_delta,
+            tool_calls=tool_calls,
+            finish_reason=choice.finish_reason,
+            usage=usage,
+        )