lite-agent-sdk 0.1.0__py3-none-any.whl

lite_agent_sdk/__init__.py

@@ -0,0 +1,41 @@
+"""Minimal, lightweight, provider-agnostic LLM agent SDK."""
+
+from lite_agent_sdk.agent import Agent
+from lite_agent_sdk.errors import AgentError
+from lite_agent_sdk.tools import BaseTool, ToolContext, tool
+from lite_agent_sdk.types import (
+    AgentEvent,
+    Message,
+    RunComplete,
+    RunResult,
+    RunStarted,
+    TextDelta,
+    ToolCall,
+    ToolCallEnd,
+    ToolCallRecord,
+    ToolCallStart,
+    ToolResult,
+    Usage,
+    UsageUpdate,
+)
+
+__all__ = [
+    "Agent",
+    "AgentError",
+    "AgentEvent",
+    "BaseTool",
+    "Message",
+    "RunComplete",
+    "RunResult",
+    "RunStarted",
+    "TextDelta",
+    "ToolCall",
+    "ToolCallEnd",
+    "ToolCallRecord",
+    "ToolCallStart",
+    "ToolContext",
+    "ToolResult",
+    "Usage",
+    "UsageUpdate",
+    "tool",
+]

lite_agent_sdk/agent.py

@@ -0,0 +1,562 @@
+import json
+import os
+import time
+from collections.abc import AsyncIterator, Sequence
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, ValidationError
+
+from lite_agent_sdk.errors import AgentError
+from lite_agent_sdk.providers import _ADAPTER_FACTORIES, create_adapter
+from lite_agent_sdk.providers.base import (
+    ModelFinish,
+    ModelRequest,
+    ModelResponse,
+    ProviderAdapter,
+    ProviderFinish,
+    ProviderTextDelta,
+    ProviderToolCall,
+    ProviderToolCallStart,
+    ProviderUsage,
+    ToolSchema,
+)
+from lite_agent_sdk.tools.base import BaseTool, ToolContext
+from lite_agent_sdk.types import (
+    AgentEvent,
+    FinishReason,
+    Message,
+    RunComplete,
+    RunResult,
+    RunStarted,
+    TextDelta,
+    ToolCall,
+    ToolCallEnd,
+    ToolCallRecord,
+    ToolCallStart,
+    ToolResult,
+    Usage,
+    UsageUpdate,
+)
+
+
+def _resolve_api_key(provider_name: str, api_key: str | None) -> str:
+    if api_key is not None:
+        return api_key
+    env_var = "OPENAI_API_KEY" if provider_name == "openai" else "ANTHROPIC_API_KEY"
+    resolved = os.environ.get(env_var)
+    if not resolved:
+        raise AgentError(
+            "missing_api_key",
+            f"missing api key for provider {provider_name!r} (set api_key= or {env_var})",
+        )
+    return resolved
+
+
+@dataclass
+class _StreamHopDone:
+    response: ModelResponse
+
+
+class Agent:
+    def __init__(
+        self,
+        model: str,
+        system_prompt: str,
+        api_key: str | None = None,
+        tools: Sequence[BaseTool] | None = None,
+        output_schema: type[BaseModel] | None = None,
+        temperature: float | None = None,
+        top_p: float | None = None,
+        max_tokens: int | None = None,
+        stop: list[str] | None = None,
+        provider_options: dict[str, Any] | None = None,
+        *,
+        adapter: ProviderAdapter | None = None,
+    ) -> None:
+        if output_schema is not None and not issubclass(output_schema, BaseModel):
+            raise TypeError("output_schema must be a BaseModel subclass")
+
+        provider_name, sep, model_id = model.partition("/")
+        if not sep or not provider_name or not model_id:
+            raise AgentError("unknown_provider", f"invalid model id: {model!r}")
+
+        if adapter is not None:
+            self._adapter = adapter
+        elif provider_name in _ADAPTER_FACTORIES:
+            resolved_key = _resolve_api_key(provider_name, api_key)
+            self._adapter = create_adapter(provider_name, resolved_key)
+        else:
+            raise AgentError("unknown_provider", f"unknown provider {provider_name!r}")
+
+        tool_list = list(tools or [])
+        names = [t.name for t in tool_list]
+        if len(names) != len(set(names)):
+            raise ValueError("duplicate tool names")
+
+        self.model = model
+        self._model_id = model_id
+        self.system_prompt = system_prompt
+        self.api_key = api_key
+        self.tools = tool_list
+        self.output_schema = output_schema
+        self.temperature = temperature
+        self.top_p = top_p
+        self.max_tokens = max_tokens
+        self.stop = stop
+        self.provider_options = provider_options
+        self._tool_map = {t.name: t for t in self.tools}
+
+    async def run(self, messages: list[Message]) -> RunResult:
+        run_id = str(uuid4())
+        started = time.monotonic()
+        transcript = list(messages)
+        usage_total = Usage()
+
+        pending, resolved, _ = await self._resolve_unresolved(transcript, run_id)
+        if pending:
+            return self._build_result(
+                transcript,
+                usage_total,
+                run_id,
+                "waiting_approval",
+                output="",
+                started=started,
+            )
+
+        need_hop = resolved or not _unresolved_calls(transcript)
+
+        while need_hop:
+            response = await self._model_hop(transcript, run_id)
+            usage_total = _merge_usage(usage_total, response.usage)
+
+            if response.finish == "length":
+                return self._build_result(
+                    transcript,
+                    usage_total,
+                    run_id,
+                    "length",
+                    output=response.message.content or "",
+                    started=started,
+                )
+
+            assistant = response.message
+            if assistant.tool_calls:
+                stamped = _stamp_approvals(assistant.tool_calls, self._tool_map)
+                assistant = assistant.model_copy(update={"tool_calls": stamped})
+            transcript.append(assistant)
+
+            if not assistant.tool_calls:
+                output = self._final_output(assistant.content or "", run_id)
+                return self._build_result(
+                    transcript,
+                    usage_total,
+                    run_id,
+                    "stop",
+                    output=output,
+                    started=started,
+                )
+
+            pending, resolved, _ = await self._resolve_unresolved(transcript, run_id)
+            if pending:
+                return self._build_result(
+                    transcript,
+                    usage_total,
+                    run_id,
+                    "waiting_approval",
+                    output="",
+                    started=started,
+                )
+
+            need_hop = resolved
+
+        raise RuntimeError("agent loop ended without a terminal state")
+
+    def stream(self, messages: list[Message]) -> AsyncIterator[AgentEvent]:
+        async def _gen() -> AsyncIterator[AgentEvent]:
+            run_id = str(uuid4())
+            started = time.monotonic()
+            transcript = list(messages)
+            usage_total = Usage()
+
+            yield RunStarted(run_id=run_id, model=self.model)
+
+            pending, resolved, resolve_events = await self._resolve_unresolved(
+                transcript, run_id, emit_tool_results=True
+            )
+            for event in resolve_events:
+                yield event
+            if pending:
+                yield RunComplete(
+                    result=self._build_result(
+                        transcript,
+                        usage_total,
+                        run_id,
+                        "waiting_approval",
+                        output="",
+                        started=started,
+                    )
+                )
+                return
+
+            need_hop = resolved or not _unresolved_calls(transcript)
+
+            while need_hop:
+                response: ModelResponse | None = None
+                async for event in self._adapter_stream_hop(transcript, usage_total):
+                    if isinstance(event, _StreamHopDone):
+                        response = event.response
+                        continue
+                    if isinstance(event, UsageUpdate):
+                        usage_total = event.usage_total
+                    yield event
+
+                if response is None:
+                    raise RuntimeError("adapter stream ended without a model response")
+
+                if response.finish == "length":
+                    yield RunComplete(
+                        result=self._build_result(
+                            transcript,
+                            usage_total,
+                            run_id,
+                            "length",
+                            output=response.message.content or "",
+                            started=started,
+                        )
+                    )
+                    return
+
+                assistant = response.message
+                if assistant.tool_calls:
+                    stamped = _stamp_approvals(assistant.tool_calls, self._tool_map)
+                    assistant = assistant.model_copy(update={"tool_calls": stamped})
+                transcript.append(assistant)
+
+                for tc in assistant.tool_calls or []:
+                    tool = self._tool_map.get(tc.name)
+                    yield ToolCallEnd(
+                        tool_call_id=tc.id,
+                        name=tc.name,
+                        args=tc.args,
+                        requires_approval=tool.requires_approval if tool else False,
+                        approval=tc.approval,
+                    )
+
+                if not assistant.tool_calls:
+                    output = self._final_output(assistant.content or "", run_id)
+                    yield RunComplete(
+                        result=self._build_result(
+                            transcript,
+                            usage_total,
+                            run_id,
+                            "stop",
+                            output=output,
+                            started=started,
+                        )
+                    )
+                    return
+
+                pending, resolved, resolve_events = await self._resolve_unresolved(
+                    transcript, run_id, emit_tool_results=True
+                )
+                for event in resolve_events:
+                    yield event
+                if pending:
+                    yield RunComplete(
+                        result=self._build_result(
+                            transcript,
+                            usage_total,
+                            run_id,
+                            "waiting_approval",
+                            output="",
+                            started=started,
+                        )
+                    )
+                    return
+
+                need_hop = resolved
+
+            raise RuntimeError("agent loop ended without a terminal state")
+
+        return _gen()
+
+    async def _adapter_stream_hop(
+        self,
+        transcript: list[Message],
+        usage_total: Usage,
+    ) -> AsyncIterator[AgentEvent | _StreamHopDone]:
+        request = self._build_model_request(transcript)
+        text_parts: list[str] = []
+        tool_calls: list[ToolCall] = []
+        hop_usage = Usage()
+        finish: ModelFinish | None = None
+
+        async for event in self._adapter.stream(request):
+            if isinstance(event, ProviderTextDelta):
+                text_parts.append(event.delta)
+                yield TextDelta(delta=event.delta)
+            elif isinstance(event, ProviderToolCallStart):
+                yield ToolCallStart(tool_call_id=event.tool_call_id, name=event.name)
+            elif isinstance(event, ProviderToolCall):
+                tool_calls.append(ToolCall(id=event.tool_call_id, name=event.name, args=event.args))
+            elif isinstance(event, ProviderUsage):
+                hop_usage = event.usage
+                merged = _merge_usage(usage_total, hop_usage)
+                yield UsageUpdate(usage=event.usage, usage_total=merged)
+            elif isinstance(event, ProviderFinish):
+                finish = event.finish
+
+        if finish is None:
+            raise ValueError("adapter stream ended without finish")
+
+        content = "".join(text_parts) if text_parts else None
+        if content == "" and tool_calls:
+            content = None
+
+        message = Message(
+            id=str(uuid4()),
+            role="assistant",
+            content=content,
+            tool_calls=tool_calls or None,
+            created_at=datetime.now(UTC),
+        )
+        yield _StreamHopDone(
+            response=ModelResponse(message=message, usage=hop_usage, finish=finish)
+        )
+
+    def _build_model_request(self, transcript: list[Message]) -> ModelRequest:
+        return ModelRequest(
+            model=self._model_id,
+            system=self.system_prompt,
+            messages=list(transcript),
+            tools=[
+                ToolSchema(
+                    name=t.name,
+                    description=t.description,
+                    parameters=t.input_schema().model_json_schema(),
+                )
+                for t in self.tools
+            ],
+            temperature=self.temperature,
+            top_p=self.top_p,
+            max_tokens=self.max_tokens,
+            stop=self.stop,
+            output_schema=(
+                self.output_schema.model_json_schema() if self.output_schema is not None else None
+            ),
+            provider_options=self.provider_options,
+        )
+
+    async def _model_hop(self, transcript: list[Message], run_id: str) -> ModelResponse:
+        request = self._build_model_request(transcript)
+        return await self._adapter.complete(request)
+
+    async def _resolve_unresolved(
+        self,
+        transcript: list[Message],
+        run_id: str,
+        *,
+        emit_tool_results: bool = False,
+    ) -> tuple[bool, bool, list[ToolResult]]:
+        still_pending = False
+        results_appended = False
+        resolve_events: list[ToolResult] = []
+
+        for _assistant_msg, tc in _unresolved_calls(transcript):
+            if tc.approval == "pending":
+                still_pending = True
+                continue
+
+            if tc.approval == "denied":
+                body = self._append_tool_result(
+                    transcript,
+                    tc,
+                    {"status": "denied", "reason": "User denied this tool call"},
+                    is_error=False,
+                )
+                results_appended = True
+                if emit_tool_results:
+                    resolve_events.append(
+                        ToolResult(
+                            tool_call_id=tc.id,
+                            name=tc.name,
+                            result=body,
+                            is_error=False,
+                        )
+                    )
+                continue
+
+            tool = self._tool_map.get(tc.name)
+            if tool is None:
+                raise AgentError("unknown_tool", f"unknown tool: {tc.name}", run_id=run_id)
+
+            if tc.approval not in (None, "approved"):
+                continue
+
+            schema = tool.input_schema()
+            try:
+                args = schema.model_validate(tc.args)
+            except ValidationError as e:
+                body = self._append_tool_result(
+                    transcript,
+                    tc,
+                    {"status": "error", "error": str(e)},
+                    is_error=True,
+                )
+                results_appended = True
+                if emit_tool_results:
+                    resolve_events.append(
+                        ToolResult(
+                            tool_call_id=tc.id,
+                            name=tc.name,
+                            result=body,
+                            is_error=True,
+                        )
+                    )
+                continue
+
+            ctx = ToolContext(run_id=run_id)
+            try:
+                result = await tool.execute(args, ctx)
+                content: Any = result
+                if not isinstance(result, str):
+                    content = json.dumps(result)
+                body = self._append_tool_result(transcript, tc, content, is_error=False)
+                is_error = False
+            except Exception as e:
+                body = self._append_tool_result(
+                    transcript,
+                    tc,
+                    {"status": "error", "error": str(e)},
+                    is_error=True,
+                )
+                is_error = True
+            results_appended = True
+            if emit_tool_results:
+                resolve_events.append(
+                    ToolResult(
+                        tool_call_id=tc.id,
+                        name=tc.name,
+                        result=body,
+                        is_error=is_error,
+                    )
+                )
+
+        return still_pending, results_appended, resolve_events
+
+    def _append_tool_result(
+        self,
+        transcript: list[Message],
+        tc: ToolCall,
+        content: Any,
+        *,
+        is_error: bool,
+    ) -> str:
+        body = json.dumps(content) if isinstance(content, dict) else str(content)
+        transcript.append(
+            Message(
+                id=str(uuid4()),
+                role="tool",
+                content=body,
+                tool_call_id=tc.id,
+                tool_name=tc.name,
+                tool_is_error=is_error,
+                created_at=datetime.now(UTC),
+            )
+        )
+        return body
+
+    def _final_output(self, content: str, run_id: str) -> Any:
+        if self.output_schema is None:
+            return content
+        try:
+            return self.output_schema.model_validate_json(content)
+        except ValidationError as e:
+            raise AgentError("invalid_output", str(e), run_id=run_id) from e
+
+    def _build_result(
+        self,
+        transcript: list[Message],
+        usage: Usage,
+        run_id: str,
+        finish_reason: FinishReason,
+        *,
+        output: Any,
+        started: float,
+    ) -> RunResult:
+        pending = [tc for _msg, tc in _unresolved_calls(transcript) if tc.approval == "pending"]
+        return RunResult(
+            run_id=run_id,
+            output=output,
+            messages=transcript,
+            usage=usage,
+            tool_calls=_tool_call_records(transcript),
+            pending_approvals=pending,
+            finish_reason=finish_reason,
+            model=self.model,
+            latency_ms=int((time.monotonic() - started) * 1000),
+        )
+
+
+def _unresolved_calls(transcript: list[Message]) -> list[tuple[Message, ToolCall]]:
+    result_ids = {m.tool_call_id for m in transcript if m.role == "tool" and m.tool_call_id}
+    out: list[tuple[Message, ToolCall]] = []
+    for msg in transcript:
+        if msg.role != "assistant" or not msg.tool_calls:
+            continue
+        for tc in msg.tool_calls:
+            if tc.id not in result_ids:
+                out.append((msg, tc))
+    return out
+
+
+def _stamp_approvals(
+    tool_calls: list[ToolCall],
+    tool_map: dict[str, BaseTool],
+) -> list[ToolCall]:
+    stamped: list[ToolCall] = []
+    for tc in tool_calls:
+        tool = tool_map.get(tc.name)
+        if tool is not None and tool.requires_approval and tc.approval is None:
+            stamped.append(tc.model_copy(update={"approval": "pending"}))
+        else:
+            stamped.append(tc)
+    return stamped
+
+
+def _merge_usage(total: Usage, hop: Usage) -> Usage:
+    return Usage(
+        input_tokens=total.input_tokens + hop.input_tokens,
+        output_tokens=total.output_tokens + hop.output_tokens,
+        cache_read_tokens=total.cache_read_tokens + hop.cache_read_tokens,
+        cache_write_tokens=total.cache_write_tokens + hop.cache_write_tokens,
+        raw={**total.raw, **hop.raw},
+    )
+
+
+def _tool_call_records(transcript: list[Message]) -> list[ToolCallRecord]:
+    results_by_id: dict[str, Message] = {}
+    for msg in transcript:
+        if msg.role == "tool" and msg.tool_call_id:
+            results_by_id[msg.tool_call_id] = msg
+
+    records: list[ToolCallRecord] = []
+    for msg in transcript:
+        if msg.role != "assistant" or not msg.tool_calls:
+            continue
+        for tc in msg.tool_calls:
+            result_msg = results_by_id.get(tc.id)
+            records.append(
+                ToolCallRecord(
+                    id=tc.id,
+                    name=tc.name,
+                    args=tc.args,
+                    approval=tc.approval,
+                    result=result_msg.content if result_msg else None,
+                    is_error=result_msg.tool_is_error if result_msg else False,
+                )
+            )
+    return records

lite_agent_sdk/errors.py

@@ -0,0 +1,19 @@
+from typing import Literal
+
+
+class AgentError(Exception):
+    def __init__(
+        self,
+        code: Literal[
+            "unknown_provider",
+            "unknown_tool",
+            "invalid_output",
+            "missing_api_key",
+        ],
+        message: str,
+        run_id: str | None = None,
+    ) -> None:
+        self.code = code
+        self.message = message
+        self.run_id = run_id
+        super().__init__(message)

lite_agent_sdk/providers/__init__.py

@@ -0,0 +1,29 @@
+from collections.abc import Callable
+
+from lite_agent_sdk.providers.anthropic_provider import AnthropicAdapter, create_anthropic_adapter
+from lite_agent_sdk.providers.base import ProviderAdapter
+from lite_agent_sdk.providers.openai_provider import OpenAIAdapter, create_openai_adapter
+
+AdapterFactory = Callable[[str | None], ProviderAdapter]
+
+_ADAPTER_FACTORIES: dict[str, AdapterFactory] = {
+    "openai": create_openai_adapter,
+    "anthropic": create_anthropic_adapter,
+}
+
+
+def create_adapter(provider_name: str, api_key: str | None) -> ProviderAdapter:
+    factory = _ADAPTER_FACTORIES.get(provider_name)
+    if factory is None:
+        raise KeyError(provider_name)
+    return factory(api_key)
+
+
+__all__ = [
+    "_ADAPTER_FACTORIES",
+    "AnthropicAdapter",
+    "OpenAIAdapter",
+    "create_adapter",
+    "create_anthropic_adapter",
+    "create_openai_adapter",
+]