smarta2a 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

smarta2a/server/state_manager.py

@@ -1,34 +1,188 @@
 # Library imports
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, List
 from uuid import uuid4
+from nats.aio.client import Client as NATS
+import asyncio
 
 # Local imports
 from smarta2a.state_stores.base_state_store import BaseStateStore
 from smarta2a.history_update_strategies.history_update_strategy import HistoryUpdateStrategy
-from smarta2a.utils.types import Message, StateData
+from smarta2a.utils.types import Message, StateData, Task, TaskStatus, TaskState, PushNotificationConfig, Part
+from smarta2a.server.nats_client import NATSClient
 
 class StateManager:
-    def __init__(self, store: Optional[BaseStateStore], history_strategy: HistoryUpdateStrategy):
-        self.store = store
+    def __init__(self, state_store: BaseStateStore, history_strategy: HistoryUpdateStrategy):
+        self.state_store = state_store
         self.strategy = history_strategy
+        self.nats_client = NATSClient(server_url="nats://localhost:4222")
+    
+    async def load(self):
+        await self.nats_client.connect()
+    
+
+    async def unload(self):
+        await self.nats_client.close()
+    
+    def _initialize_empty_state(
+        self,
+        task_id: str,
+        session_id: str,
+        push_notification_config: Optional[PushNotificationConfig] = None
+    ) -> StateData:
+        """
+        Build a fresh StateData and persist it.
+        """
+        initial_task = Task(
+            id=task_id,
+            sessionId=session_id,
+            status=TaskStatus(state=TaskState.WORKING),
+            artifacts=[],
+            history=[],
+            metadata={}
+        )
+        state = StateData(
+            task_id=task_id,
+            task=initial_task,
+            context_history=[],
+            push_notification_config=push_notification_config
+        )
+        self.state_store.initialize_state(state)
+        return state
+
+    async def get_or_create_and_update_state(
+        self,
+        task_id: str,
+        session_id: str,
+        message: Message,
+        metadata: Optional[Dict[str, Any]] = None,
+        push_notification_config: Optional[PushNotificationConfig] = None
+    ) -> StateData:
+        """
+        Fetch existing StateData, or initialize & persist a new one.
+        """
+        existing_state = self.state_store.get_state(task_id)
+        if not existing_state:
+            latest_state = self._initialize_empty_state(
+                task_id, session_id, push_notification_config
+            )
+        else:
+            latest_state = existing_state.copy()
+
+        latest_state.task.history.append(message)
+        latest_state.context_history = self.strategy.update_history(
+            existing_history=latest_state.context_history,
+            new_messages=[message]
+        )
+        latest_state.task.metadata = metadata
 
-    def init_or_get(self, session_id: Optional[str], message: Message, metadata: Dict[str, Any]) -> StateData:
-        sid = session_id or str(uuid4())
-        if not self.store:
-            return StateData(sessionId=sid, history=[message], metadata=metadata or {})
-        existing = self.store.get_state(sid) or StateData(sessionId=sid, history=[], metadata={})
-        existing.history.append(message)
-        existing.metadata = {**(existing.metadata or {}), **(metadata or {})}
-        self.store.update_state(sid, existing)
-        return existing
+        await self.update_state(latest_state)
 
-    def update(self, state: StateData):
-        if self.store:
-            self.store.update_state(state.sessionId, state)
+        return latest_state
+    
+    def get_and_update_state_from_webhook(self, task_id: str, result: Task) -> StateData:
+        """
+        Update existing state with webhook result data, including:
+        - Merges task history from result
+        - Extracts messages from artifacts' parts
+        - Updates context history using strategy
+        - Merges artifacts and metadata
+        
+        Raises ValueError if no existing state is found
+        """
+        existing_state = self.state_store.get_state(task_id)
+        if not existing_state:
+            raise ValueError(f"No existing state found for task_id: {task_id}")
+
+        updated_state = existing_state.copy()
+        new_messages = []
+
+        # Add messages from result's history
+        if result.history:
+            new_messages.extend(result.history)
+
+        # Extract messages from result's artifacts
+        for artifact in result.artifacts or []:
+            if artifact.parts:
+                artifact_message = Message(
+                    role="tool",
+                    parts=artifact.parts,
+                    metadata=artifact.metadata
+                )
+                new_messages.append(artifact_message)
+
+        # Update task history (merge with existing)
+        if new_messages:
+            if updated_state.task.history is None:
+                updated_state.task.history = []
+            updated_state.task.history.extend(new_messages)
+
+        # Update context history using strategy
+        updated_state.context_history = self.strategy.update_history(
+            existing_history=updated_state.context_history,
+            new_messages=new_messages
+        )
+
+        # Merge artifacts
+        if result.artifacts:
+            if updated_state.task.artifacts is None:
+                updated_state.task.artifacts = []
+            updated_state.task.artifacts.extend(result.artifacts)
+
+        # Merge metadata
+        if result.metadata:
+            updated_state.task.metadata = {
+                **(updated_state.task.metadata or {}),
+                **(result.metadata or {})
+            }
+
+        # Update task status if provided
+        if result.status:
+            updated_state.task.status = result.status
+
+        return updated_state
+    
+    def get_state(self, task_id: str) -> Optional[StateData]:
+        return self.state_store.get_state(task_id)
+
+    async def update_state(self, state_data: StateData):
+        self.state_store.update_state(state_data.task_id, state_data)
+
+        # Publish update through NATS client
+        payload = self._prepare_update_payload(state_data)
+        await self.nats_client.publish("state.updates", payload)
     
     def get_store(self) -> Optional[BaseStateStore]:
-        return self.store
+        return self.state_store
     
-    def get_strategy(self) -> HistoryUpdateStrategy:
+    def get_history_strategy(self) -> HistoryUpdateStrategy:
         return self.strategy
-    
+    
+    # Private methods
+
+    def _serialize_part(self, part: Part) -> dict:
+        """Serialize a Part to frontend-compatible format"""
+        part_data = part.model_dump()
+        if part.type == "file" and part.file:
+            if not part.file.bytes and not part.file.uri:
+                raise ValueError("FilePart must have either bytes or uri")
+        return part_data
+
+    def _prepare_update_payload(self, state: StateData) -> Dict[str, Any]:
+        """Prepare NATS message payload from state data"""
+        return {
+            "taskId": state.task_id,
+            "parts": self._extract_artifact_parts(state.task),
+            "complete": state.task.status.state == TaskState.COMPLETED
+        }
+
+    def _extract_artifact_parts(self, task: Task) -> List[dict]:
+        """Extract and serialize parts from all artifacts"""
+        parts = []
+        if task.artifacts:
+            for artifact in task.artifacts:
+                for part in artifact.parts:
+                    try:
+                        parts.append(self._serialize_part(part))
+                    except ValueError as e:
+                        print(f"Invalid part in artifact: {e}")
+        return parts

smarta2a/server/webhook_request_processor.py

@@ -0,0 +1,112 @@
+# Library imports
+from typing import Callable, Any, Optional
+
+# Local imports
+from smarta2a.server.state_manager import StateManager
+from smarta2a.utils.types import WebhookRequest, WebhookResponse, StateData, Message
+from smarta2a.client.a2a_client import A2AClient
+
+class WebhookRequestProcessor:
+    def __init__(self, webhook_fn: Callable[[WebhookRequest], Any], state_manager: Optional[StateManager] = None):
+        self.webhook_fn = webhook_fn
+        self.state_manager = state_manager
+        self.a2a_aclient = A2AClient()
+
+    async def process_request(self, request: WebhookRequest) -> WebhookResponse:
+        if self.state_manager:
+            state_data = self.state_manager.get_and_update_state_from_webhook(request.id, request.result)
+            return await self._handle_webhook(request, state_data)
+        else:
+            return await self._handle_webhook(request)
+
+
+    async def _webhook_handler(self, request: WebhookRequest, state_data: Optional[StateData] = None) -> WebhookResponse:
+        try:
+            # --- Step 1: Process Incoming Task ---
+            if request.result:
+                incoming_task = request.result
+                
+                # Initialize state_data if missing
+                if not state_data:
+                    state_data = StateData(
+                        task_id=incoming_task.id,
+                        task=incoming_task.copy(update={"artifacts": incoming_task.artifacts}),
+                        context_history=[],
+                        push_notification_config=None
+                    )
+                else:
+                    existing_task = state_data.task
+                    
+                    # Overwrite artifacts
+                    existing_task.artifacts = incoming_task.artifacts.copy() if incoming_task.artifacts else []
+                    
+                    # Merge metadata
+                    existing_task.metadata = {**(existing_task.metadata or {}), **(incoming_task.metadata or {})}
+                    
+                    # Build messages from artifact parts (role="agent" as in handle_send_task)
+                    all_parts = [part for artifact in incoming_task.artifacts for part in artifact.parts] if incoming_task.artifacts else []
+                    new_messages = [Message(role="agent", parts=all_parts, metadata=incoming_task.metadata)]
+                    
+                    # Update context history using strategy
+                    history_strategy = self.state_manager.get_history_strategy()
+                    state_data.context_history = history_strategy.update_history(
+                        existing_history=state_data.context_history,
+                        new_messages=new_messages
+                    )
+                    
+                    # Persist state
+                    await self.state_manager.update_state(state_data)
+
+            # --- Step 2: Call Webhook Function ---
+            webhook_response = await self.webhook_fn(request, state_data) if state_data else await self.webhook_fn(request)
+            
+            # --- Step 3: Process Webhook Response ---
+            if webhook_response.result:
+                updated_task = webhook_response.result
+                existing_task = state_data.task
+                
+                # Overwrite artifacts from response
+                existing_task.artifacts = updated_task.artifacts.copy() if updated_task.artifacts else []
+                
+                # Merge metadata
+                existing_task.metadata = {**(existing_task.metadata or {}), **(updated_task.metadata or {})}
+                
+                # Build messages from updated artifacts
+                updated_parts = [part for artifact in updated_task.artifacts for part in artifact.parts] if updated_task.artifacts else []
+                updated_messages = [Message(role="agent", parts=updated_parts, metadata=updated_task.metadata)]
+                
+                # Update context history again
+                state_data.context_history = history_strategy.update_history(
+                    existing_history=state_data.context_history,
+                    new_messages=updated_messages
+                )
+                
+                await self.state_manager.update_state(state_data)
+
+            # --- Step 4: Push Notification ---
+            push_url = (
+                state_data.push_notification_config.url 
+                if state_data and state_data.push_notification_config 
+                else None
+            )
+            if push_url:
+                try:
+                    self.a2a_aclient.send_to_webhook(
+                        webhook_url=push_url,
+                        id=state_data.task_id,
+                        task=state_data.task
+                    )
+                except Exception as e:
+                    return WebhookResponse(
+                        id=request.id,
+                        error=f"Push notification failed: {str(e)}"
+                    )
+
+            # --- Step 5: Return Final Response ---
+            return WebhookResponse(
+                id=request.id,
+                result=state_data.task if state_data else None
+            )
+        
+        except Exception as e:
+            return WebhookResponse(id=request.id, error=f"Internal error: {str(e)}")

smarta2a/state_stores/base_state_store.py

@@ -8,13 +8,13 @@ from smarta2a.utils.types import StateData, Message
 class BaseStateStore(ABC):
     
     @abstractmethod
-    async def get_state(self, session_id: str) -> Optional[StateData]:
+    async def get_state(self, task_id: str) -> Optional[StateData]:
         pass
     
     @abstractmethod
-    async def update_state(self, session_id: str, state_data: StateData) -> None:
+    async def update_state(self, task_id: str, state_data: StateData) -> None:
         pass
     
     @abstractmethod
-    async def delete_state(self, session_id: str) -> None:
+    async def delete_state(self, task_id: str) -> None:
         pass

smarta2a/state_stores/inmemory_state_store.py

@@ -10,12 +10,26 @@ class InMemoryStateStore(BaseStateStore):
     def __init__(self):
         self.states: Dict[str, StateData] = {}
     
-    def get_state(self, session_id: str) -> Optional[StateData]:
-        return self.states.get(session_id)
+    def initialize_state(self, state_data: StateData) -> None:
+        self.states[state_data.task_id] = state_data
     
-    def update_state(self, session_id: str, state_data: StateData):
-        self.states[session_id] = state_data
+    def get_state(self, task_id: str) -> Optional[StateData]:
+        return self.states.get(task_id)
     
-    def delete_state(self, session_id: str):
-        if session_id in self.states:
-            del self.states[session_id]
+    def update_state(self, task_id: str, state_data: StateData):
+        self.states[task_id] = state_data
+    
+    def delete_state(self, task_id: str):
+        if task_id in self.states:
+            del self.states[task_id]
+    
+    def get_all_tasks(self, fields: Optional[str] = None) -> List[Dict[str, Any]]:
+        all_tasks = [state_data.task.model_dump() for state_data in self.states.values()]
+        if fields:
+            requested_fields = fields.split(",")
+            fields_filtered_tasks = [
+                {field: task[field] for field in requested_fields if field in task}
+                for task in all_tasks
+            ]
+            return fields_filtered_tasks
+        return all_tasks

smarta2a/utils/agent_discovery_manager.py

@@ -0,0 +1,121 @@
+# Library imports
+import httpx
+from typing import List, Optional
+from pydantic import HttpUrl, ValidationError
+
+# Local imports
+from smarta2a.utils.types import AgentCard
+
+class AgentDiscoveryManager:
+    """Centralized service for discovering agents through multiple methods"""
+    
+    def __init__(
+        self,
+        agent_cards: Optional[List[AgentCard]] = None,
+        agent_base_urls: Optional[List[HttpUrl]] = None,
+        discovery_endpoint: Optional[HttpUrl] = None,
+        timeout: float = 5.0,
+        retries: int = 2
+    ):
+        self.explicit_cards = agent_cards or []
+        self.agent_base_urls = agent_base_urls or []
+        self.discovery_endpoint = discovery_endpoint
+        self.timeout = timeout
+        self.retries = retries
+        self.discovered_cards: List[AgentCard] = []
+
+    async def discover_agents(self) -> List[AgentCard]:
+        """Discover agents through all configured methods"""
+        self.discovered_cards = []
+        
+        # 1. Add explicit cards first
+        self.discovered_cards.extend(self.explicit_cards)
+        
+        # 2. Discover via base URLs
+        if self.agent_base_urls:
+            base_url_cards = await self._discover_via_base_urls()
+            self.discovered_cards.extend(base_url_cards)
+        
+        # 3. Discover via central endpoint
+        if self.discovery_endpoint:
+            endpoint_cards = await self._discover_via_endpoint()
+            self.discovered_cards.extend(endpoint_cards)
+        
+        return self.discovered_cards
+
+    async def _discover_via_base_urls(self) -> List[AgentCard]:
+        """Discover agents from provided base URLs"""
+        cards = []
+        client = httpx.AsyncClient()
+        
+        try:
+            for base_url in self.agent_base_urls:
+                try:
+                    agent_url = f"{base_url}/.well-known/agent.json"
+                    card = await self._fetch_agent_card(client, agent_url)
+                    cards.append(card)
+                except Exception:
+                    pass
+        finally:
+            await client.aclose()
+            
+        return cards
+
+    async def _discover_via_endpoint(self) -> List[AgentCard]:
+        """Discover agents through a centralized endpoint"""
+        client = httpx.AsyncClient()
+        cards = []
+        
+        try:
+            # Fetch service registry
+            try:
+                response = await client.get(
+                    str(self.discovery_endpoint),
+                    timeout=self.timeout
+                )
+                response.raise_for_status()
+                services = response.json()["services"]
+            except Exception:
+                return []
+
+            # Fetch all discovered agent cards
+            for service in services:
+                try:
+                    agent_url = f"{service['base_url']}/.well-known/agent.json"
+                    card = await self._fetch_agent_card(client, agent_url)
+                    cards.append(card)
+                except Exception:
+                    pass
+        finally:
+            await client.aclose()
+            
+        return cards
+
+    async def _fetch_agent_card(self, client: httpx.AsyncClient, url: str) -> AgentCard:
+        """Fetch and validate a single agent.json"""
+        try:
+            response = await client.get(
+                url,
+                timeout=self.timeout,
+                follow_redirects=True,
+                headers={"User-Agent": "AgentDiscovery/1.0"}
+            )
+            response.raise_for_status()
+            
+            # Validate content type
+            content_type = response.headers.get('Content-Type', '')
+            if 'application/json' not in content_type:
+                raise ValueError(f"Unexpected Content-Type: {content_type}")
+                
+            data = response.json()
+
+            # Enforce required 'url' field
+            if "url" not in data:
+                raise ValueError("AgentCard requires 'url' field in agent.json")
+                
+            return AgentCard(**data)
+            
+        except ValidationError as e:
+            raise
+        except Exception as e:
+            raise

smarta2a/utils/prompt_helpers.py

@@ -2,7 +2,7 @@
 from typing import Optional, List
 
 # Local imports
-from smarta2a.client.tools_manager import ToolsManager
+from smarta2a.utils.tools_manager import ToolsManager
 from smarta2a.utils.types import AgentCard
 
 def build_system_prompt(

smarta2a/utils/tools_manager.py

@@ -0,0 +1,108 @@
+# Library imports
+import json
+from typing import List, Dict, Any, Union, Literal, Optional
+
+# Local imports
+from smarta2a.client.mcp_client import MCPClient
+from smarta2a.client.a2a_client import A2AClient
+from smarta2a.utils.types import AgentCard, Tool
+
+class ToolsManager:
+    """
+    Manages loading, describing, and invoking tools from various providers.
+    Acts as a wrapper around the MCP and A2A clients.
+    """
+    def __init__(self):
+        self.tools_list: List[Any] = []
+        self.clients: Dict[str, Union[MCPClient, A2AClient]] = {}
+
+    async def load_mcp_tools(self, urls_or_paths: List[str]) -> None:
+        for url in urls_or_paths:
+            mcp_client = await MCPClient.create(url)
+            tools = await mcp_client.list_tools()
+            for tool in tools:
+                # Generate key and ensure Tool type with key
+                key = f"mcp---{tool.name}"
+                validated_tool = Tool(
+                    key=key,
+                    **tool.model_dump()  # Pydantic 2.x syntax (use .dict() for Pydantic 1.x)
+                )
+                self.tools_list.append(validated_tool)
+                self.clients[key] = mcp_client
+
+    async def load_a2a_tools(self, agent_cards: List[AgentCard]) -> None:
+        for agent_card in agent_cards:
+            a2a_client = A2AClient(agent_card)
+            tools_list = await a2a_client.list_tools()
+            for tool_dict in tools_list:
+                # Generate key from agent URL and tool name
+                key = f"{agent_card.name}---{tool_dict['name']}"
+
+                # Build new description
+                components = []
+                original_desc = tool_dict['description']
+                if original_desc:
+                    components.append(original_desc)
+                if agent_card.description:
+                    components.append(f"Agent Description: {agent_card.description}")
+                
+                # Collect skill descriptions
+                skill_descriptions = []
+                for skill in agent_card.skills:
+                    if skill.description:
+                        skill_descriptions.append(skill.description)
+                if skill_descriptions:
+                    components.append(f"Agent's skills: {', '.join(skill_descriptions)}")
+            
+                # Update tool_dict with new description
+                tool_dict['description'] = ". ".join(components)
+                
+                validated_tool = Tool(
+                    key=key,
+                    **tool_dict
+                )
+                self.tools_list.append(validated_tool)
+                self.clients[key] = a2a_client
+
+    def get_tools(self) -> List[Any]:
+        return self.tools_list
+
+
+    def describe_tools(self, client_type: Literal["mcp", "a2a"]) -> str:
+        lines = []
+        for tool in self.tools_list:
+            schema = json.dumps(tool.inputSchema, indent=2)  # Fix: use inputSchema
+            if client_type == "mcp":
+                lines.append(
+                    f"- **{tool.key}**: {tool.description}\n  Parameters schema:\n  ```json\n{schema}\n```"
+                )
+            elif client_type == "a2a":
+                lines.append(
+                    f"- **{tool.key}**: {tool.description}\n  Parameters schema:\n  ```json\n{schema}\n```"
+                )
+
+        return "\n".join(lines)
+
+    def get_client(self, tool_key: str) -> Any:
+        return self.clients.get(tool_key)
+    
+    async def call_tool(self, tool_key: str, args: Dict[str, Any], override_args: Optional[Dict[str, Any]] = None) -> Any:
+        try:
+            client = self.get_client(tool_key)
+            tool_name = self._get_tool_name(tool_key)
+            new_args = self._replace_with_override_args(args, override_args)
+            result = await client.call_tool(tool_name, new_args)
+            return result
+
+        except Exception as e:
+            # This will catch ANY error in the body above
+            raise
+    
+    def _get_tool_name(self, tool_key: str) -> str:
+        return tool_key.split("---")[1]
+    
+    def _replace_with_override_args(self, args: Dict[str, Any], override_args: Optional[Dict[str, Any]] = None):
+        new_args = args.copy()
+        if override_args:
+            new_args.update(override_args)
+        return new_args

smarta2a/utils/types.py

@@ -479,11 +479,26 @@ class A2AStreamResponse(BaseModel):
     metadata: dict[str, Any] | None = None
 
 class StateData(BaseModel):
-    sessionId: str
-    history: List[Message]
-    metadata: Dict[str, Any]
+    task_id: str
+    task: Task
+    context_history: List[Message]
+    push_notification_config: PushNotificationConfig | None = None
 
 class Tool(BaseModel):
+    key: str
     name: str
     description: str
     inputSchema: Dict[str, Any]
+
+'''
+The callback request may simply be a message without a result - basically acknowledging the task was completed.
+It can also include a result, which is the task that was completed along with the full artifact.
+'''
+class WebhookRequest(BaseModel):
+    id: str
+    result: Task | None = None
+
+class WebhookResponse(BaseModel):
+    id: str
+    result: Task | None = None
+    error: str | None = None