a2a-lite 0.1.0__py3-none-any.whl

a2a_lite/discovery.py

@@ -0,0 +1,148 @@
+"""
+mDNS-based local agent discovery using Zeroconf.
+"""
+from __future__ import annotations
+
+import asyncio
+import socket
+from typing import Dict, List, Optional
+from dataclasses import dataclass
+
+from zeroconf import ServiceBrowser, ServiceListener, Zeroconf, ServiceInfo
+from zeroconf.asyncio import AsyncZeroconf, AsyncServiceBrowser
+
+
+SERVICE_TYPE = "_a2a-agent._tcp.local."
+
+
+@dataclass
+class DiscoveredAgent:
+    """Information about a discovered A2A agent."""
+    name: str
+    host: str
+    port: int
+    url: str
+    properties: Dict[str, str]
+
+
+class AgentDiscovery:
+    """
+    Discover A2A agents on the local network using mDNS.
+    """
+
+    def __init__(self):
+        self._discovered: Dict[str, DiscoveredAgent] = {}
+        self._zeroconf: Optional[Zeroconf] = None
+        self._service_info: Optional[ServiceInfo] = None
+
+    async def discover(self, timeout: float = 5.0) -> List[DiscoveredAgent]:
+        """
+        Discover agents on the local network.
+
+        Args:
+            timeout: How long to wait for discoveries
+
+        Returns:
+            List of discovered agents
+        """
+        self._discovered.clear()
+
+        async with AsyncZeroconf() as azc:
+            listener = _DiscoveryListener(self._discovered)
+            browser = AsyncServiceBrowser(
+                azc.zeroconf, SERVICE_TYPE, listener
+            )
+
+            await asyncio.sleep(timeout)
+            await browser.async_cancel()
+
+        return list(self._discovered.values())
+
+    def register(
+        self,
+        name: str,
+        port: int,
+        properties: Optional[Dict[str, str]] = None,
+    ) -> ServiceInfo:
+        """
+        Register this agent for discovery.
+
+        Args:
+            name: Agent name
+            port: Port the agent is running on
+            properties: Additional properties to advertise
+
+        Returns:
+            ServiceInfo for unregistration
+        """
+        hostname = socket.gethostname()
+        local_ip = self._get_local_ip()
+
+        # Sanitize name for mDNS (remove spaces and special chars)
+        safe_name = "".join(c if c.isalnum() or c == "-" else "-" for c in name)
+
+        info = ServiceInfo(
+            SERVICE_TYPE,
+            f"{safe_name}.{SERVICE_TYPE}",
+            addresses=[socket.inet_aton(local_ip)],
+            port=port,
+            properties=properties or {},
+            server=f"{hostname}.local.",
+        )
+
+        self._zeroconf = Zeroconf()
+        self._zeroconf.register_service(info)
+        self._service_info = info
+
+        return info
+
+    def unregister(self) -> None:
+        """Unregister this agent from discovery."""
+        if self._zeroconf and self._service_info:
+            self._zeroconf.unregister_service(self._service_info)
+            self._zeroconf.close()
+            self._zeroconf = None
+            self._service_info = None
+
+    def _get_local_ip(self) -> str:
+        """Get local IP address."""
+        try:
+            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+            s.connect(("8.8.8.8", 80))
+            ip = s.getsockname()[0]
+            s.close()
+            return ip
+        except Exception:
+            return "127.0.0.1"
+
+
+class _DiscoveryListener(ServiceListener):
+    """Internal listener for mDNS service discovery."""
+
+    def __init__(self, discovered: Dict[str, DiscoveredAgent]):
+        self.discovered = discovered
+
+    def add_service(self, zc: Zeroconf, type_: str, name: str) -> None:
+        info = zc.get_service_info(type_, name)
+        if info:
+            agent_name = name.replace(f".{SERVICE_TYPE}", "")
+            host = socket.inet_ntoa(info.addresses[0]) if info.addresses else "localhost"
+            port = info.port
+
+            self.discovered[name] = DiscoveredAgent(
+                name=agent_name,
+                host=host,
+                port=port,
+                url=f"http://{host}:{port}",
+                properties={
+                    k.decode() if isinstance(k, bytes) else k:
+                    v.decode() if isinstance(v, bytes) else str(v)
+                    for k, v in info.properties.items()
+                },
+            )
+
+    def remove_service(self, zc: Zeroconf, type_: str, name: str) -> None:
+        self.discovered.pop(name, None)
+
+    def update_service(self, zc: Zeroconf, type_: str, name: str) -> None:
+        self.add_service(zc, type_, name)

a2a_lite/executor.py

@@ -0,0 +1,317 @@
+"""
+Wrapper around A2A's AgentExecutor that dispatches to registered skill handlers.
+"""
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+from typing import Any, Callable, Dict, List, Optional
+
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.events import EventQueue
+
+from .decorators import SkillDefinition
+from .middleware import MiddlewareChain, MiddlewareContext
+from .streaming import is_generator_function, stream_generator
+
+
+class LiteAgentExecutor(AgentExecutor):
+    """
+    Simplified AgentExecutor with optional enterprise features.
+
+    Features (all optional):
+    - Middleware chain
+    - Streaming support
+    - Pydantic model conversion
+    - Task context injection
+    - Interaction context injection
+    - Authentication
+    - File part handling
+    """
+
+    def __init__(
+        self,
+        skills: Dict[str, SkillDefinition],
+        error_handler: Optional[Callable] = None,
+        middleware: Optional[MiddlewareChain] = None,
+        on_complete: Optional[List[Callable]] = None,
+        auth_provider: Optional[Any] = None,
+        task_store: Optional[Any] = None,
+    ):
+        self.skills = skills
+        self.error_handler = error_handler
+        self.middleware = middleware or MiddlewareChain()
+        self.on_complete = on_complete or []
+        self.auth_provider = auth_provider
+        self.task_store = task_store
+
+    async def execute(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+    ) -> None:
+        """Execute a skill based on the incoming request."""
+        from a2a.utils import new_agent_text_message
+
+        try:
+            # Extract message and parts
+            message, parts = self._extract_message_and_parts(context)
+
+            # Parse skill call
+            skill_name, params = self._parse_message(message)
+
+            # Build middleware context
+            ctx = MiddlewareContext(
+                skill=skill_name,
+                params=params,
+                message=message,
+            )
+
+            # Store parts in metadata for skill access
+            ctx.metadata["parts"] = parts
+            ctx.metadata["event_queue"] = event_queue
+
+            # Define final handler
+            async def final_handler(ctx: MiddlewareContext) -> Any:
+                return await self._execute_skill(
+                    ctx.skill,
+                    ctx.params,
+                    event_queue,
+                    ctx.metadata,
+                )
+
+            # Execute through middleware chain
+            result = await self.middleware.execute(ctx, final_handler)
+
+            # If result is not None and not already streamed, send it
+            if result is not None:
+                if isinstance(result, (dict, list)):
+                    response_text = json.dumps(result, indent=2, default=str)
+                else:
+                    response_text = str(result)
+                await event_queue.enqueue_event(new_agent_text_message(response_text))
+
+            # Call completion hooks
+            for hook in self.on_complete:
+                try:
+                    if asyncio.iscoroutinefunction(hook):
+                        await hook(skill_name, result, ctx)
+                    else:
+                        hook(skill_name, result, ctx)
+                except Exception:
+                    pass
+
+        except Exception as e:
+            await self._handle_error(e, event_queue)
+
+    async def _execute_skill(
+        self,
+        skill_name: Optional[str],
+        params: Dict[str, Any],
+        event_queue: EventQueue,
+        metadata: Dict[str, Any],
+    ) -> Any:
+        """Execute a skill with the given parameters."""
+        if skill_name is None:
+            if not self.skills:
+                return {"error": "No skills registered"}
+            skill_name = list(self.skills.keys())[0]
+
+        if skill_name not in self.skills:
+            return {
+                "error": f"Unknown skill: {skill_name}",
+                "available_skills": list(self.skills.keys()),
+            }
+
+        skill_def = self.skills[skill_name]
+
+        # Convert Pydantic models and file parts in params
+        params = self._convert_params(skill_def, params, metadata)
+
+        # Inject special contexts if needed
+        if skill_def.needs_task_context and self.task_store:
+            from .tasks import TaskContext, Task, TaskStatus, TaskState
+            task = self.task_store.create(skill_name, params)
+            # Only pass event_queue for streaming skills (status updates go via SSE)
+            eq = event_queue if skill_def.is_streaming else None
+            task_ctx = TaskContext(task, eq)
+            params["task"] = task_ctx
+
+        if skill_def.needs_interaction:
+            from .human_loop import InteractionContext
+            task_id = metadata.get("task_id", "unknown")
+            interaction_ctx = InteractionContext(task_id, event_queue)
+            params["ctx"] = interaction_ctx
+
+        # Call the handler
+        handler = skill_def.handler
+
+        if skill_def.is_streaming or is_generator_function(handler):
+            gen = handler(**params)
+            await stream_generator(gen, event_queue)
+            return None
+        else:
+            return await self._call_handler(handler, **params)
+
+    def _convert_params(
+        self,
+        skill_def: SkillDefinition,
+        params: Dict[str, Any],
+        metadata: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        """Convert parameters to Pydantic models and file parts if needed."""
+        handler = skill_def.handler
+        hints = getattr(handler, '__annotations__', {})
+
+        converted = {}
+        for param_name, value in params.items():
+            param_type = hints.get(param_name)
+
+            if param_type is None:
+                converted[param_name] = value
+                continue
+
+            type_name = str(param_type)
+
+            # Skip special context types
+            if "TaskContext" in type_name or "InteractionContext" in type_name:
+                continue
+
+            # Convert FilePart
+            if "FilePart" in type_name:
+                from .parts import FilePart
+                if isinstance(value, dict):
+                    # Handle both A2A format and simple dict format
+                    if "file" in value:
+                        converted[param_name] = FilePart.from_a2a(value)
+                    else:
+                        # Simple format: {name, data, mime_type}
+                        data = value.get("data")
+                        if isinstance(data, str):
+                            data = data.encode("utf-8")
+                        converted[param_name] = FilePart(
+                            name=value.get("name", "unknown"),
+                            mime_type=value.get("mime_type", "application/octet-stream"),
+                            data=data,
+                            uri=value.get("uri"),
+                        )
+                else:
+                    converted[param_name] = value
+                continue
+
+            # Convert DataPart
+            if "DataPart" in type_name:
+                from .parts import DataPart
+                if isinstance(value, dict):
+                    # Handle both A2A format and simple dict format
+                    if "type" in value and value.get("type") == "data":
+                        converted[param_name] = DataPart.from_a2a(value)
+                    else:
+                        # Simple format: pass the dict directly as data
+                        converted[param_name] = DataPart(data=value)
+                else:
+                    converted[param_name] = value
+                continue
+
+            # Convert Pydantic models
+            if hasattr(param_type, 'model_validate'):
+                if isinstance(value, dict):
+                    converted[param_name] = param_type.model_validate(value)
+                else:
+                    converted[param_name] = value
+                continue
+
+            # Default: keep as-is
+            converted[param_name] = value
+
+        return converted
+
+    def _parse_message(self, message: str) -> tuple[Optional[str], Dict[str, Any]]:
+        """Parse message to extract skill name and params."""
+        try:
+            data = json.loads(message)
+            if isinstance(data, dict) and 'skill' in data:
+                return data['skill'], data.get('params', {})
+        except json.JSONDecodeError:
+            pass
+
+        return None, {"message": message}
+
+    def _extract_message_and_parts(self, context: RequestContext) -> tuple[str, List[Any]]:
+        """Extract message text and any file/data parts."""
+        text = ""
+        parts = []
+
+        if hasattr(context, 'message') and context.message:
+            message = context.message
+            if hasattr(message, 'parts'):
+                raw_parts = message.parts
+            else:
+                raw_parts = message.get('parts', [])
+
+            for part in raw_parts:
+                if hasattr(part, 'root'):
+                    part = part.root
+
+                # Get part type
+                if hasattr(part, 'text'):
+                    text = part.text
+                elif isinstance(part, dict):
+                    part_type = part.get('type') or part.get('kind')
+                    if part_type == 'text':
+                        text = part.get('text', '')
+                    elif part_type in ('file', 'data'):
+                        parts.append(part)
+
+        return text, parts
+
+    async def _handle_error(self, e: Exception, event_queue: EventQueue) -> None:
+        """Handle execution errors."""
+        from a2a.utils import new_agent_text_message
+
+        if self.error_handler:
+            try:
+                result = await self._call_handler(self.error_handler, e)
+                await event_queue.enqueue_event(
+                    new_agent_text_message(json.dumps(result, default=str))
+                )
+                return
+            except Exception as handler_error:
+                await event_queue.enqueue_event(
+                    new_agent_text_message(json.dumps({
+                        "error": str(e),
+                        "handler_error": str(handler_error),
+                        "type": type(e).__name__,
+                    }))
+                )
+                return
+
+        await event_queue.enqueue_event(
+            new_agent_text_message(json.dumps({
+                "error": str(e),
+                "type": type(e).__name__,
+            }))
+        )
+
+    async def cancel(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+    ) -> None:
+        """Handle cancellation requests."""
+        from a2a.utils import new_agent_text_message
+        await event_queue.enqueue_event(
+            new_agent_text_message(json.dumps({"status": "cancelled"}))
+        )
+
+    async def _call_handler(self, handler: Callable, *args, **kwargs) -> Any:
+        """Call a handler, handling both sync and async functions."""
+        if asyncio.iscoroutinefunction(handler):
+            return await handler(*args, **kwargs)
+        else:
+            loop = asyncio.get_event_loop()
+            return await loop.run_in_executor(
+                None,
+                lambda: handler(*args, **kwargs)
+            )