a2a-lite 0.1.0__py3-none-any.whl

a2a_lite/__init__.py

@@ -0,0 +1,151 @@
+"""
+A2A Lite - Simple by default, powerful when needed.
+
+SIMPLE (8 lines):
+    from a2a_lite import Agent
+
+    agent = Agent(name="Bot", description="A bot")
+
+    @agent.skill("greet")
+    async def greet(name: str) -> str:
+        return f"Hello, {name}!"
+
+    agent.run()
+
+TEST IT (3 lines):
+    from a2a_lite import TestClient
+    client = TestClient(agent)
+    assert client.call("greet", name="World") == "Hello, World!"
+
+WITH PYDANTIC:
+    class User(BaseModel):
+        name: str
+
+    @agent.skill("create")
+    async def create(user: User) -> dict:
+        return {"created": user.name}
+
+WITH STREAMING:
+    @agent.skill("chat", streaming=True)
+    async def chat(msg: str):
+        for word in msg.split():
+            yield word
+
+WITH AUTH (opt-in):
+    from a2a_lite.auth import APIKeyAuth
+    agent = Agent(name="Bot", auth=APIKeyAuth(keys=["secret"]))
+
+WITH TASK TRACKING (opt-in):
+    from a2a_lite.tasks import TaskContext
+
+    @agent.skill("process")
+    async def process(data: str, task: TaskContext) -> str:
+        await task.update("working", progress=0.5)
+        return "done"
+
+WITH HUMAN-IN-THE-LOOP (opt-in):
+    from a2a_lite.human_loop import InteractionContext
+
+    @agent.skill("wizard")
+    async def wizard(ctx: InteractionContext) -> str:
+        name = await ctx.ask("What's your name?")
+        return f"Hello, {name}!"
+
+WITH FILES (opt-in):
+    from a2a_lite.parts import FilePart
+
+    @agent.skill("summarize")
+    async def summarize(doc: FilePart) -> str:
+        text = await doc.read_text()
+        return summarize(text)
+"""
+
+# Core - always available
+from .agent import Agent
+from .decorators import SkillDefinition
+from .discovery import AgentDiscovery, DiscoveredAgent
+from .testing import TestClient, AsyncTestClient
+
+# Middleware - always available
+from .middleware import (
+    MiddlewareContext,
+    MiddlewareChain,
+    logging_middleware,
+    timing_middleware,
+    retry_middleware,
+    rate_limit_middleware,
+)
+
+# Webhooks - always available
+from .webhooks import WebhookClient, WebhookConfig, NotificationManager
+
+# Streaming - always available
+from .streaming import StreamingResponse
+
+# Parts - opt-in for multi-modal
+from .parts import TextPart, FilePart, DataPart, Artifact
+
+# Tasks - opt-in for task lifecycle
+from .tasks import TaskContext, TaskState, TaskStatus, Task, TaskStore
+
+# Human-in-the-loop - opt-in
+from .human_loop import InteractionContext, ConversationMemory
+
+# Auth - opt-in
+from .auth import (
+    AuthProvider,
+    AuthResult,
+    NoAuth,
+    APIKeyAuth,
+    BearerAuth,
+    OAuth2Auth,
+    require_auth,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Core
+    "Agent",
+    "SkillDefinition",
+    "AgentDiscovery",
+    "DiscoveredAgent",
+    # Testing
+    "TestClient",
+    "AsyncTestClient",
+    # Middleware
+    "MiddlewareContext",
+    "MiddlewareChain",
+    "logging_middleware",
+    "timing_middleware",
+    "retry_middleware",
+    "rate_limit_middleware",
+    # Webhooks
+    "WebhookClient",
+    "WebhookConfig",
+    "NotificationManager",
+    # Streaming
+    "StreamingResponse",
+    # Parts (multi-modal)
+    "TextPart",
+    "FilePart",
+    "DataPart",
+    "Artifact",
+    # Tasks
+    "TaskContext",
+    "TaskState",
+    "TaskStatus",
+    "Task",
+    "TaskStore",
+    # Human-in-the-loop
+    "InteractionContext",
+    "ConversationMemory",
+    # Auth
+    "AuthProvider",
+    "AuthResult",
+    "NoAuth",
+    "APIKeyAuth",
+    "BearerAuth",
+    "OAuth2Auth",
+    "require_auth",
+]

a2a_lite/agent.py

@@ -0,0 +1,453 @@
+"""
+Core Agent class that wraps the A2A SDK complexity.
+
+Simple by default, powerful when needed.
+"""
+from __future__ import annotations
+
+import asyncio
+import inspect
+from typing import Any, Callable, Optional, Dict, List, Type, Union, get_origin, get_args
+from dataclasses import dataclass, field
+
+import uvicorn
+
+from a2a.server.apps import A2AStarletteApplication
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.tasks import InMemoryTaskStore
+from a2a.types import (
+    AgentCard,
+    AgentSkill,
+    AgentCapabilities,
+)
+
+from .executor import LiteAgentExecutor
+from .decorators import SkillDefinition
+from .utils import type_to_json_schema, extract_function_schemas
+from .middleware import MiddlewareChain, MiddlewareContext
+from .streaming import is_generator_function
+from .webhooks import NotificationManager, WebhookClient
+
+
+@dataclass
+class Agent:
+    """
+    Simplified A2A Agent - simple by default, powerful when needed.
+
+    SIMPLE (8 lines):
+        agent = Agent(name="Bot", description="A bot")
+
+        @agent.skill("greet")
+        async def greet(name: str) -> str:
+            return f"Hello, {name}!"
+
+        agent.run()
+
+    WITH PYDANTIC:
+        class User(BaseModel):
+            name: str
+
+        @agent.skill("create")
+        async def create(user: User) -> dict:
+            return {"created": user.name}
+
+    WITH STREAMING:
+        @agent.skill("chat", streaming=True)
+        async def chat(msg: str):
+            for word in msg.split():
+                yield word
+
+    WITH AUTH (optional):
+        from a2a_lite.auth import APIKeyAuth
+
+        agent = Agent(
+            name="SecureBot",
+            auth=APIKeyAuth(keys=["secret"]),
+        )
+
+    WITH TASK TRACKING (optional):
+        @agent.skill("process")
+        async def process(data: str, task: TaskContext) -> str:
+            await task.update("working", progress=0.5)
+            return "done"
+
+    WITH HUMAN-IN-THE-LOOP (optional):
+        @agent.skill("wizard")
+        async def wizard(ctx: InteractionContext) -> str:
+            name = await ctx.ask("What's your name?")
+            return f"Hello, {name}!"
+    """
+    name: str
+    description: str
+    version: str = "1.0.0"
+    url: Optional[str] = None
+
+    # Optional enterprise features
+    auth: Optional[Any] = None  # AuthProvider
+    task_store: Optional[Any] = None  # TaskStore or "memory"
+
+    def __post_init__(self):
+        # Internal state
+        self._skills: Dict[str, SkillDefinition] = {}
+        self._error_handler: Optional[Callable] = None
+        self._on_startup: List[Callable] = []
+        self._on_shutdown: List[Callable] = []
+        self._on_complete: List[Callable] = []
+        self._discovery = None
+        self._middleware = MiddlewareChain()
+        self._notifications = NotificationManager()
+        self._webhook = WebhookClient()
+        self._has_streaming = False
+
+        # Setup optional task store
+        if self.task_store == "memory":
+            from .tasks import TaskStore
+            self._task_store = TaskStore()
+        elif self.task_store:
+            self._task_store = self.task_store
+        else:
+            self._task_store = None
+
+        # Setup optional auth
+        if self.auth is None:
+            from .auth import NoAuth
+            self._auth = NoAuth()
+        else:
+            self._auth = self.auth
+
+    def skill(
+        self,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        streaming: bool = False,
+    ) -> Callable:
+        """
+        Decorator to register a function as an agent skill.
+
+        Simple:
+            @agent.skill("greet")
+            async def greet(name: str) -> str:
+                return f"Hello, {name}!"
+
+        With streaming:
+            @agent.skill("chat", streaming=True)
+            async def chat(message: str):
+                for word in message.split():
+                    yield word
+
+        With task context (opt-in):
+            @agent.skill("process")
+            async def process(data: str, task: TaskContext) -> str:
+                await task.update("working", progress=0.5)
+                return "done"
+
+        With human-in-the-loop (opt-in):
+            @agent.skill("wizard")
+            async def wizard(ctx: InteractionContext) -> str:
+                name = await ctx.ask("What's your name?")
+                return f"Hello, {name}!"
+        """
+        def decorator(func: Callable) -> Callable:
+            skill_name = name or func.__name__
+            skill_desc = description or func.__doc__ or f"Skill: {skill_name}"
+
+            # Clean up docstring
+            if skill_desc:
+                skill_desc = " ".join(skill_desc.split())
+
+            # Detect streaming
+            is_streaming = streaming or is_generator_function(func)
+            if is_streaming:
+                self._has_streaming = True
+
+            # Detect special parameter types
+            hints = getattr(func, '__annotations__', {})
+            needs_task_context = any(
+                str(h).endswith("TaskContext") or "TaskContext" in str(h)
+                for h in hints.values()
+            )
+            needs_interaction = any(
+                str(h).endswith("InteractionContext") or "InteractionContext" in str(h)
+                for h in hints.values()
+            )
+
+            # Extract schemas
+            input_schema, output_schema = extract_function_schemas(func)
+
+            skill_def = SkillDefinition(
+                name=skill_name,
+                description=skill_desc,
+                tags=tags or [],
+                handler=func,
+                input_schema=input_schema,
+                output_schema=output_schema,
+                is_async=asyncio.iscoroutinefunction(func) or is_streaming,
+                is_streaming=is_streaming,
+                needs_task_context=needs_task_context,
+                needs_interaction=needs_interaction,
+            )
+
+            self._skills[skill_name] = skill_def
+            return func
+
+        return decorator
+
+    def middleware(self, func: Callable) -> Callable:
+        """
+        Decorator to register middleware.
+
+        Example:
+            @agent.middleware
+            async def log_requests(ctx, next):
+                print(f"Calling: {ctx.skill}")
+                return await next()
+        """
+        self._middleware.add(func)
+        return func
+
+    def add_middleware(self, middleware: Callable) -> None:
+        """Add a middleware function (non-decorator version)."""
+        self._middleware.add(middleware)
+
+    def on_error(self, func: Callable) -> Callable:
+        """Decorator to register a global error handler."""
+        self._error_handler = func
+        return func
+
+    def on_startup(self, func: Callable) -> Callable:
+        """Decorator to register a startup hook."""
+        self._on_startup.append(func)
+        return func
+
+    def on_shutdown(self, func: Callable) -> Callable:
+        """Decorator to register a shutdown hook."""
+        self._on_shutdown.append(func)
+        return func
+
+    def on_complete(self, func: Callable) -> Callable:
+        """Decorator to register a task completion handler."""
+        self._on_complete.append(func)
+        return func
+
+    @property
+    def webhook(self) -> WebhookClient:
+        """Get the webhook client for sending notifications."""
+        return self._webhook
+
+    @property
+    def notifications(self) -> NotificationManager:
+        """Get the notification manager for push notifications."""
+        return self._notifications
+
+    def build_agent_card(self, host: str = "localhost", port: int = 8787) -> AgentCard:
+        """Generate A2A-compliant Agent Card from registered skills."""
+        skills = []
+
+        for skill_def in self._skills.values():
+            skill = AgentSkill(
+                id=skill_def.name,
+                name=skill_def.name,
+                description=skill_def.description,
+                tags=skill_def.tags,
+                inputModes=["application/json"],
+                outputModes=["application/json"],
+            )
+            skills.append(skill)
+
+        url = self.url or f"http://{host}:{port}"
+
+        # Check if any skills need human-in-the-loop
+        has_input_required = any(
+            getattr(s, 'needs_interaction', False)
+            for s in self._skills.values()
+        )
+
+        return AgentCard(
+            name=self.name,
+            description=self.description,
+            version=self.version,
+            url=url,
+            capabilities=AgentCapabilities(
+                streaming=self._has_streaming,
+                pushNotifications=bool(self._on_complete),
+            ),
+            defaultInputModes=["application/json"],
+            defaultOutputModes=["application/json"],
+            skills=skills,
+        )
+
+    def run(
+        self,
+        host: str = "0.0.0.0",
+        port: int = 8787,
+        reload: bool = False,
+        log_level: str = "info",
+        enable_discovery: bool = False,
+    ) -> None:
+        """
+        Start the A2A server.
+
+        Simple:
+            agent.run()
+
+        With options:
+            agent.run(port=9000, enable_discovery=True)
+        """
+        from rich.console import Console
+        from rich.panel import Panel
+
+        console = Console()
+
+        # Build components
+        display_host = "localhost" if host == "0.0.0.0" else host
+        agent_card = self.build_agent_card(display_host, port)
+        executor = LiteAgentExecutor(
+            skills=self._skills,
+            error_handler=self._error_handler,
+            middleware=self._middleware,
+            on_complete=self._on_complete,
+            auth_provider=self._auth,
+            task_store=self._task_store,
+        )
+
+        request_handler = DefaultRequestHandler(
+            agent_executor=executor,
+            task_store=InMemoryTaskStore(),
+        )
+
+        app_builder = A2AStarletteApplication(
+            agent_card=agent_card,
+            http_handler=request_handler,
+        )
+
+        # Build display info
+        skills_list = "\n".join([
+            f"  • {s.name}: {s.description}" +
+            (" [streaming]" if getattr(s, 'is_streaming', False) else "") +
+            (" [interactive]" if getattr(s, 'needs_interaction', False) else "")
+            for s in self._skills.values()
+        ])
+        if not skills_list:
+            skills_list = "  (no skills registered)"
+
+        # Collect enabled features
+        features = []
+        if len(self._middleware._middlewares):
+            features.append(f"{len(self._middleware._middlewares)} middleware")
+        if self._has_streaming:
+            features.append("streaming")
+        if self._on_complete:
+            features.append("webhooks")
+        if self.auth:
+            features.append("auth")
+        if self._task_store:
+            features.append("task-tracking")
+
+        features_str = f"\n\n[bold]Features:[/] {', '.join(features)}" if features else ""
+
+        console.print(Panel(
+            f"[bold green]{self.name}[/] v{self.version}\n\n"
+            f"[dim]{self.description}[/]\n\n"
+            f"[bold]Skills:[/]\n{skills_list}{features_str}\n\n"
+            f"[bold]Endpoints:[/]\n"
+            f"  • Agent Card: http://{display_host}:{port}/.well-known/agent.json\n"
+            f"  • API: http://{display_host}:{port}/",
+            title="🚀 A2A Lite Agent Started",
+            border_style="green",
+        ))
+
+        # Run startup hooks
+        for hook in self._on_startup:
+            if asyncio.iscoroutinefunction(hook):
+                asyncio.get_event_loop().run_until_complete(hook())
+            else:
+                hook()
+
+        # Enable discovery if requested
+        if enable_discovery:
+            from .discovery import AgentDiscovery
+            self._discovery = AgentDiscovery()
+            self._discovery.register(
+                name=self.name,
+                port=port,
+                properties={"version": self.version},
+            )
+            console.print(f"[dim]mDNS discovery enabled for {self.name}[/]")
+
+        # Start server
+        try:
+            uvicorn.run(
+                app_builder.build(),
+                host=host,
+                port=port,
+                log_level=log_level,
+            )
+        finally:
+            # Run shutdown hooks
+            for hook in self._on_shutdown:
+                if asyncio.iscoroutinefunction(hook):
+                    asyncio.get_event_loop().run_until_complete(hook())
+                else:
+                    hook()
+
+            # Unregister discovery
+            if self._discovery:
+                self._discovery.unregister()
+
+    async def call_remote(
+        self,
+        agent_url: str,
+        message: str,
+        timeout: float = 30.0,
+    ) -> Dict[str, Any]:
+        """Call a remote A2A agent."""
+        import httpx
+        from a2a.client import A2AClient
+        from a2a.types import MessageSendParams, SendMessageRequest
+        from uuid import uuid4
+
+        async with httpx.AsyncClient(timeout=timeout) as http_client:
+            card_url = f"{agent_url.rstrip('/')}/.well-known/agent.json"
+
+            client = await A2AClient.get_client_from_agent_card_url(
+                http_client, card_url
+            )
+
+            request = SendMessageRequest(
+                id=uuid4().hex,
+                params=MessageSendParams(
+                    message={
+                        "role": "user",
+                        "parts": [{"type": "text", "text": message}],
+                        "messageId": uuid4().hex,
+                    }
+                )
+            )
+
+            response = await client.send_message(request)
+            return response.model_dump()
+
+    def get_app(self):
+        """Get the Starlette application without running it."""
+        agent_card = self.build_agent_card()
+        executor = LiteAgentExecutor(
+            skills=self._skills,
+            error_handler=self._error_handler,
+            middleware=self._middleware,
+            on_complete=self._on_complete,
+            auth_provider=self._auth,
+            task_store=self._task_store,
+        )
+
+        request_handler = DefaultRequestHandler(
+            agent_executor=executor,
+            task_store=InMemoryTaskStore(),
+        )
+
+        app_builder = A2AStarletteApplication(
+            agent_card=agent_card,
+            http_handler=request_handler,
+        )
+
+        return app_builder.build()