a2a-lite 0.1.0__py3-none-any.whl

a2a_lite/streaming.py

@@ -0,0 +1,89 @@
+"""
+Streaming support for A2A Lite agents.
+
+Enables generator-based streaming for LLM-style responses:
+
+    @agent.skill("chat", streaming=True)
+    async def chat(message: str):
+        async for chunk in llm.stream(message):
+            yield chunk
+"""
+from __future__ import annotations
+
+from typing import Any, AsyncGenerator, Callable, Generator, Union
+import asyncio
+import inspect
+
+
+def is_generator_function(func: Callable) -> bool:
+    """Check if a function is a generator (sync or async)."""
+    return (
+        inspect.isgeneratorfunction(func) or
+        inspect.isasyncgenfunction(func)
+    )
+
+
+async def collect_generator(
+    gen: Union[Generator, AsyncGenerator]
+) -> list[Any]:
+    """Collect all items from a generator into a list."""
+    items = []
+    if inspect.isasyncgen(gen):
+        async for item in gen:
+            items.append(item)
+    else:
+        for item in gen:
+            items.append(item)
+    return items
+
+
+async def stream_generator(
+    gen: Union[Generator, AsyncGenerator],
+    event_queue,
+) -> None:
+    """
+    Stream generator output through the A2A event queue.
+
+    Each yielded item becomes a separate event in the stream.
+    """
+    from a2a.utils import new_agent_text_message
+
+    if inspect.isasyncgen(gen):
+        async for chunk in gen:
+            text = str(chunk) if not isinstance(chunk, str) else chunk
+            await event_queue.enqueue_event(new_agent_text_message(text))
+    else:
+        for chunk in gen:
+            text = str(chunk) if not isinstance(chunk, str) else chunk
+            await event_queue.enqueue_event(new_agent_text_message(text))
+
+
+class StreamingResponse:
+    """
+    Helper class for building streaming responses.
+
+    Example:
+        @agent.skill("count")
+        async def count(n: int):
+            stream = StreamingResponse()
+            for i in range(n):
+                stream.write(f"Count: {i}")
+            return stream
+    """
+
+    def __init__(self):
+        self._chunks: list[str] = []
+
+    def write(self, chunk: str) -> None:
+        """Add a chunk to the stream."""
+        self._chunks.append(chunk)
+
+    def __iter__(self):
+        return iter(self._chunks)
+
+    def __aiter__(self):
+        return self._async_iter()
+
+    async def _async_iter(self):
+        for chunk in self._chunks:
+            yield chunk

a2a_lite/tasks.py

@@ -0,0 +1,221 @@
+"""
+Task lifecycle management (OPTIONAL).
+
+By default, A2A Lite just returns results (simple mode).
+Enable task tracking when you need:
+- Progress updates
+- Long-running tasks
+- Task status visibility
+
+Example (simple - no task tracking needed):
+    @agent.skill("greet")
+    async def greet(name: str) -> str:
+        return f"Hello, {name}!"
+
+Example (with task tracking - opt-in):
+    agent = Agent(name="Bot", task_store="memory")  # Enable tracking
+
+    @agent.skill("process")
+    async def process(data: str, task: TaskContext) -> str:
+        await task.update("working", "Starting...", progress=0.0)
+
+        for i in range(10):
+            await task.update("working", f"Step {i+1}/10", progress=i/10)
+            await do_work(i)
+
+        return "Done!"
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any, Callable, Dict, List, Optional
+from uuid import uuid4
+import asyncio
+
+
+class TaskState(str, Enum):
+    """A2A Protocol task states."""
+    SUBMITTED = "submitted"
+    WORKING = "working"
+    INPUT_REQUIRED = "input-required"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELED = "canceled"
+    AUTH_REQUIRED = "auth-required"
+
+
+@dataclass
+class TaskStatus:
+    """Current status of a task."""
+    state: TaskState
+    message: Optional[str] = None
+    progress: Optional[float] = None  # 0.0 to 1.0
+    timestamp: datetime = field(default_factory=datetime.utcnow)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "state": self.state.value,
+            "message": self.message,
+            "progress": self.progress,
+            "timestamp": self.timestamp.isoformat(),
+        }
+
+
+@dataclass
+class Task:
+    """Represents an A2A task."""
+    id: str
+    skill: str
+    params: Dict[str, Any]
+    status: TaskStatus
+    result: Any = None
+    error: Optional[str] = None
+    artifacts: List[Any] = field(default_factory=list)
+    history: List[TaskStatus] = field(default_factory=list)
+    created_at: datetime = field(default_factory=datetime.utcnow)
+    updated_at: datetime = field(default_factory=datetime.utcnow)
+
+    def update_status(
+        self,
+        state: TaskState,
+        message: Optional[str] = None,
+        progress: Optional[float] = None,
+    ) -> None:
+        """Update task status."""
+        self.history.append(self.status)
+        self.status = TaskStatus(state=state, message=message, progress=progress)
+        self.updated_at = datetime.utcnow()
+
+
+class TaskContext:
+    """
+    Context passed to skills when task tracking is enabled.
+
+    Provides methods to update task status and request user input.
+
+    Example:
+        @agent.skill("process")
+        async def process(data: str, task: TaskContext) -> str:
+            await task.update("working", "Processing...", progress=0.5)
+            result = await heavy_computation(data)
+            return result
+    """
+
+    def __init__(self, task: Task, event_queue=None, input_handler=None):
+        self._task = task
+        self._event_queue = event_queue
+        self._input_handler = input_handler
+        self._status_callbacks: List[Callable] = []
+
+    @property
+    def task_id(self) -> str:
+        return self._task.id
+
+    @property
+    def state(self) -> TaskState:
+        return self._task.status.state
+
+    async def update(
+        self,
+        state: str = "working",
+        message: Optional[str] = None,
+        progress: Optional[float] = None,
+    ) -> None:
+        """
+        Update task status.
+
+        Args:
+            state: Task state (working, completed, failed, etc.)
+            message: Human-readable status message
+            progress: Progress from 0.0 to 1.0
+
+        Example:
+            await task.update("working", "Processing item 5/10", progress=0.5)
+        """
+        task_state = TaskState(state) if isinstance(state, str) else state
+        self._task.update_status(task_state, message, progress)
+
+        # Notify callbacks
+        for callback in self._status_callbacks:
+            try:
+                if asyncio.iscoroutinefunction(callback):
+                    await callback(self._task.status)
+                else:
+                    callback(self._task.status)
+            except Exception:
+                pass
+
+        # Send SSE event if streaming
+        if self._event_queue:
+            await self._send_status_event()
+
+    async def _send_status_event(self) -> None:
+        """Send status update via SSE."""
+        if self._event_queue:
+            from a2a.utils import new_agent_text_message
+            import json
+            status_msg = json.dumps({
+                "_type": "status_update",
+                "task_id": self._task.id,
+                "status": self._task.status.to_dict(),
+            })
+            await self._event_queue.enqueue_event(new_agent_text_message(status_msg))
+
+    def on_status_change(self, callback: Callable) -> None:
+        """Register callback for status changes."""
+        self._status_callbacks.append(callback)
+
+
+class TaskStore:
+    """
+    In-memory task store.
+
+    For production, extend this with Redis/DB backend.
+    """
+
+    def __init__(self):
+        self._tasks: Dict[str, Task] = {}
+
+    def create(self, skill: str, params: Dict[str, Any]) -> Task:
+        """Create a new task."""
+        task = Task(
+            id=uuid4().hex,
+            skill=skill,
+            params=params,
+            status=TaskStatus(state=TaskState.SUBMITTED),
+        )
+        self._tasks[task.id] = task
+        return task
+
+    def get(self, task_id: str) -> Optional[Task]:
+        """Get task by ID."""
+        return self._tasks.get(task_id)
+
+    def update(self, task: Task) -> None:
+        """Update task in store."""
+        self._tasks[task.id] = task
+
+    def list(
+        self,
+        state: Optional[TaskState] = None,
+        skill: Optional[str] = None,
+        limit: int = 100,
+    ) -> List[Task]:
+        """List tasks with optional filters."""
+        tasks = list(self._tasks.values())
+
+        if state:
+            tasks = [t for t in tasks if t.status.state == state]
+        if skill:
+            tasks = [t for t in tasks if t.skill == skill]
+
+        return sorted(tasks, key=lambda t: t.created_at, reverse=True)[:limit]
+
+    def delete(self, task_id: str) -> bool:
+        """Delete a task."""
+        if task_id in self._tasks:
+            del self._tasks[task_id]
+            return True
+        return False

a2a_lite/testing.py

@@ -0,0 +1,268 @@
+"""
+Testing utilities for A2A Lite agents.
+
+Makes testing agents as simple as:
+
+    from a2a_lite.testing import TestClient
+
+    def test_my_agent():
+        client = TestClient(agent)
+        result = client.call("greet", name="World")
+        assert result == "Hello, World!"
+"""
+from __future__ import annotations
+
+import json
+from typing import Any, Dict, Optional
+from uuid import uuid4
+
+
+class TestClient:
+    """
+    Simple test client for A2A Lite agents.
+
+    Example:
+        agent = Agent(name="Test", description="Test")
+
+        @agent.skill("add")
+        async def add(a: int, b: int) -> int:
+            return a + b
+
+        # In your test
+        client = TestClient(agent)
+        assert client.call("add", a=2, b=3) == 5
+    """
+
+    def __init__(self, agent):
+        """
+        Create a test client for an agent.
+
+        Args:
+            agent: The A2A Lite Agent instance to test
+        """
+        self.agent = agent
+        self._app = None
+        self._client = None
+
+    def _get_client(self):
+        """Lazily create the test client."""
+        if self._client is None:
+            from starlette.testclient import TestClient as StarletteTestClient
+            self._app = self.agent.get_app()
+            self._client = StarletteTestClient(self._app)
+        return self._client
+
+    def call(self, skill: str, **params) -> Any:
+        """
+        Call a skill and return the result.
+
+        Args:
+            skill: Name of the skill to call
+            **params: Parameters to pass to the skill
+
+        Returns:
+            The skill's return value (parsed from JSON if possible)
+
+        Example:
+            result = client.call("greet", name="World")
+        """
+        client = self._get_client()
+
+        message = json.dumps({"skill": skill, "params": params})
+        request_body = {
+            "jsonrpc": "2.0",
+            "method": "message/send",
+            "id": uuid4().hex,
+            "params": {
+                "message": {
+                    "role": "user",
+                    "parts": [{"type": "text", "text": message}],
+                    "messageId": uuid4().hex,
+                }
+            }
+        }
+
+        response = client.post("/", json=request_body)
+        response.raise_for_status()
+        data = response.json()
+
+        # Extract the actual result from A2A response
+        return self._extract_result(data)
+
+    def _extract_result(self, response: Dict) -> Any:
+        """Extract the skill result from A2A response."""
+        if "error" in response:
+            raise TestClientError(response["error"])
+
+        result = response.get("result", {})
+
+        # Get text from message parts
+        parts = result.get("parts", [])
+        for part in parts:
+            if part.get("kind") == "text" or part.get("type") == "text":
+                text = part.get("text", "")
+                # Try to parse as JSON
+                try:
+                    return json.loads(text)
+                except json.JSONDecodeError:
+                    return text
+
+        return result
+
+    def get_agent_card(self) -> Dict[str, Any]:
+        """
+        Fetch the agent card.
+
+        Returns:
+            The agent card as a dictionary
+        """
+        client = self._get_client()
+        response = client.get("/.well-known/agent.json")
+        response.raise_for_status()
+        return response.json()
+
+    def list_skills(self) -> list[str]:
+        """
+        Get list of available skill names.
+
+        Returns:
+            List of skill names
+        """
+        card = self.get_agent_card()
+        return [s.get("name", s.get("id")) for s in card.get("skills", [])]
+
+    def stream(self, skill: str, **params) -> list[Any]:
+        """
+        Call a streaming skill and collect all results.
+
+        Args:
+            skill: Name of the skill to call
+            **params: Parameters to pass to the skill
+
+        Returns:
+            List of all streamed values
+
+        Example:
+            results = client.stream("count", limit=3)
+            assert len(results) == 3
+        """
+        import asyncio
+
+        # Access skills directly from agent
+        skill_def = self.agent._skills.get(skill)
+
+        if not skill_def:
+            raise TestClientError(f"Unknown skill: {skill}")
+
+        # Call handler directly and collect results
+        results = []
+
+        async def run_handler():
+            handler = skill_def.handler
+            gen = handler(**params)
+
+            # Handle both async and sync generators
+            if hasattr(gen, '__anext__'):
+                async for value in gen:
+                    results.append(value)
+            elif hasattr(gen, '__next__'):
+                for value in gen:
+                    results.append(value)
+            else:
+                # Not a generator, just a coroutine
+                result = await gen
+                results.append(result)
+
+        asyncio.get_event_loop().run_until_complete(run_handler())
+        return results
+
+
+class TestClientError(Exception):
+    """Error from test client."""
+    pass
+
+
+# Async version for async tests
+class AsyncTestClient:
+    """
+    Async test client for A2A Lite agents.
+
+    Example:
+        async def test_my_agent():
+            client = AsyncTestClient(agent)
+            result = await client.call("greet", name="World")
+            assert result == "Hello, World!"
+    """
+
+    def __init__(self, agent):
+        self.agent = agent
+        self._app = None
+        self._client = None
+
+    async def _get_client(self):
+        """Lazily create the async test client."""
+        if self._client is None:
+            import httpx
+            self._app = self.agent.get_app()
+            self._client = httpx.AsyncClient(
+                app=self._app,
+                base_url="http://testserver"
+            )
+        return self._client
+
+    async def call(self, skill: str, **params) -> Any:
+        """
+        Call a skill and return the result.
+
+        Args:
+            skill: Name of the skill to call
+            **params: Parameters to pass to the skill
+
+        Returns:
+            The skill's return value
+        """
+        client = await self._get_client()
+
+        message = json.dumps({"skill": skill, "params": params})
+        request_body = {
+            "jsonrpc": "2.0",
+            "method": "message/send",
+            "id": uuid4().hex,
+            "params": {
+                "message": {
+                    "role": "user",
+                    "parts": [{"type": "text", "text": message}],
+                    "messageId": uuid4().hex,
+                }
+            }
+        }
+
+        response = await client.post("/", json=request_body)
+        response.raise_for_status()
+        data = response.json()
+
+        return self._extract_result(data)
+
+    def _extract_result(self, response: Dict) -> Any:
+        """Extract the skill result from A2A response."""
+        if "error" in response:
+            raise TestClientError(response["error"])
+
+        result = response.get("result", {})
+        parts = result.get("parts", [])
+
+        for part in parts:
+            if part.get("kind") == "text" or part.get("type") == "text":
+                text = part.get("text", "")
+                try:
+                    return json.loads(text)
+                except json.JSONDecodeError:
+                    return text
+
+        return result
+
+    async def close(self):
+        """Close the client."""
+        if self._client:
+            await self._client.aclose()
+            self._client = None

a2a_lite/utils.py

@@ -0,0 +1,117 @@
+"""
+Helper functions for A2A Lite.
+"""
+from typing import Any, Dict, Type, get_origin, get_args, Union
+import inspect
+
+
+def type_to_json_schema(python_type: Type) -> Dict[str, Any]:
+    """
+    Convert Python type to JSON Schema.
+
+    Handles basic types, generics (List, Dict, Optional), and Pydantic models.
+    """
+    # Handle None type
+    if python_type is type(None):
+        return {"type": "null"}
+
+    # Basic type mapping
+    type_map = {
+        str: {"type": "string"},
+        int: {"type": "integer"},
+        float: {"type": "number"},
+        bool: {"type": "boolean"},
+        list: {"type": "array"},
+        dict: {"type": "object"},
+        Any: {"type": "object"},
+    }
+
+    # Check basic types first
+    if python_type in type_map:
+        return type_map[python_type]
+
+    # Handle generic types
+    origin = get_origin(python_type)
+    args = get_args(python_type)
+
+    # Handle Optional (Union[X, None])
+    if origin is Union:
+        non_none_args = [a for a in args if a is not type(None)]
+        if len(non_none_args) == 1:
+            # This is Optional[X]
+            return type_to_json_schema(non_none_args[0])
+        # Union of multiple types
+        return {"oneOf": [type_to_json_schema(a) for a in args]}
+
+    # Handle List[X]
+    if origin is list and args:
+        return {
+            "type": "array",
+            "items": type_to_json_schema(args[0])
+        }
+
+    # Handle Dict[K, V]
+    if origin is dict and len(args) >= 2:
+        return {
+            "type": "object",
+            "additionalProperties": type_to_json_schema(args[1])
+        }
+
+    # Handle Pydantic models
+    if hasattr(python_type, 'model_json_schema'):
+        return python_type.model_json_schema()
+
+    # Handle dataclasses
+    if hasattr(python_type, '__dataclass_fields__'):
+        properties = {}
+        required = []
+        for field_name, field_info in python_type.__dataclass_fields__.items():
+            properties[field_name] = type_to_json_schema(field_info.type)
+            if field_info.default is inspect.Parameter.empty and field_info.default_factory is inspect.Parameter.empty:
+                required.append(field_name)
+        return {
+            "type": "object",
+            "properties": properties,
+            "required": required,
+        }
+
+    # Fallback for unknown types
+    return {"type": "object"}
+
+
+def extract_function_schemas(func) -> tuple[Dict[str, Any], Dict[str, Any]]:
+    """
+    Extract input and output JSON schemas from a function's type hints.
+
+    Returns:
+        Tuple of (input_schema, output_schema)
+    """
+    sig = inspect.signature(func)
+    hints = getattr(func, '__annotations__', {})
+
+    # Build input schema from parameters
+    properties = {}
+    required = []
+
+    for param_name, param in sig.parameters.items():
+        if param_name in ('self', 'cls'):
+            continue
+
+        param_type = hints.get(param_name, Any)
+        properties[param_name] = type_to_json_schema(param_type)
+
+        # Parameter is required if it has no default value
+        if param.default is inspect.Parameter.empty:
+            required.append(param_name)
+
+    input_schema = {
+        "type": "object",
+        "properties": properties,
+        "required": required,
+    }
+
+    # Build output schema from return type
+    return_type = hints.get('return', Any)
+    output_schema = type_to_json_schema(return_type)
+
+    return input_schema, output_schema