smarta2a 0.3.1__py3-none-any.whl → 0.4.1__py3-none-any.whl

smarta2a/model_providers/openai_provider.py

@@ -1,13 +1,16 @@
 # Library imports
 import json
+import httpx
 from typing import AsyncGenerator, List, Dict, Optional, Union, Any
 from openai import AsyncOpenAI
+from pydantic import HttpUrl, ValidationError
 
 # Local imports
-from smarta2a.utils.types import Message, TextPart, FilePart, DataPart, Part, AgentCard
+from smarta2a.utils.types import Message, TextPart, FilePart, DataPart, Part, AgentCard, StateData
 from smarta2a.model_providers.base_llm_provider import BaseLLMProvider
-from smarta2a.client.tools_manager import ToolsManager
+from smarta2a.utils.tools_manager import ToolsManager
 from smarta2a.utils.prompt_helpers import build_system_prompt
+from smarta2a.utils.agent_discovery_manager import AgentDiscoveryManager
 
 class OpenAIProvider(BaseLLMProvider):
     def __init__(
@@ -17,22 +20,39 @@ class OpenAIProvider(BaseLLMProvider):
         base_system_prompt: Optional[str] = None,
         mcp_server_urls_or_paths: Optional[List[str]] = None,
         agent_cards: Optional[List[AgentCard]] = None,
-        # enable_discovery: bool = False
+        agent_base_urls: Optional[List[HttpUrl]] = None,
+        discovery_endpoint: Optional[HttpUrl] = None,
+        timeout: float = 5.0,
+        retries: int = 2
     ):
         self.client = AsyncOpenAI(api_key=api_key)
         self.model = model
         self.mcp_server_urls_or_paths = mcp_server_urls_or_paths
-        self.agent_cards = agent_cards
         # Store the base system prompt; will be enriched by tool descriptions
         self.base_system_prompt = base_system_prompt
         self.supported_media_types = [
             "image/png", "image/jpeg", "image/gif", "image/webp"
         ]
+
+        # Initialize discovery manager
+        self.agent_discovery = AgentDiscoveryManager(
+            agent_cards=agent_cards,
+            agent_base_urls=agent_base_urls,
+            discovery_endpoint=discovery_endpoint,
+            timeout=timeout,
+            retries=retries
+        )
+
+        self.agent_cards: List[AgentCard] = []
         # Initialize ToolsManager 
         self.tools_manager = ToolsManager()
         
     
     async def load(self):
+        """Async initialization of resources"""
+        # Discover agents first
+        self.agent_cards = await self.agent_discovery.discover_agents()
+
         if self.mcp_server_urls_or_paths:
             await self.tools_manager.load_mcp_tools(self.mcp_server_urls_or_paths)
         
@@ -116,102 +136,108 @@ class OpenAIProvider(BaseLLMProvider):
         return openai_messages
     
 
-    def _format_openai_tools(self) -> List[dict]:
+    def _format_openai_functions(self) -> List[dict]:
         """
         Convert internal tools metadata to OpenAI's function-call schema.
         """
-        openai_tools = []
+        functions = []
         for tool in self.tools_manager.get_tools():
-            openai_tools.append({
-                "type": "function",
-                "function": {
-                    "name": tool.key,
-                    "description": tool.description,
-                    "parameters": tool.inputSchema
-                }
+            functions.append({
+                "name": tool.key,
+                "description": tool.description,
+                "parameters": tool.inputSchema,
             })
-        return openai_tools
-
+        return functions
+        
 
-    async def generate(self, messages: List[Dict[str, Any]], **kwargs) -> str:
+    async def generate(self, state: StateData, **kwargs) -> str:
         """
         Generate a complete response, invoking tools as needed.
         """
-        # Ensure messages are Message objects
-        messages = [msg if isinstance(msg, Message) else Message(**msg) for msg in messages]
-        converted_messages = self._convert_messages(messages)
+        
+        # Prepare history
+        messages = [msg if isinstance(msg, Message) else Message(**msg) for msg in state.context_history]
+        openai_messages = self._convert_messages(messages)
         max_iterations = 30
 
-        for iteration in range(max_iterations):
+        for _ in range(max_iterations):
+            # Call ChatCompletion with functions (not 'tools')
             response = await self.client.chat.completions.create(
                 model=self.model,
-                messages=converted_messages,
-                tools=self._format_openai_tools(),
+                messages=openai_messages,
+                functions=self._format_openai_functions(),
                 **kwargs
             )
-            message = response.choices[0].message
-
-            # Detect and extract the tool/function call
-            if getattr(message, 'function_call', None):
-                name = message.function_call.name
-                args_raw = message.function_call.arguments
-            elif getattr(message, 'tool_calls', None):
-                tc = message.tool_calls[0]
-                name = tc.function.name
-                args_raw = tc.function.arguments
-            else:
-                return message.content
-            print(message)
-            # Append the assistant's intent
-            converted_messages.append({
+            msg = response.choices[0].message
+
+            # If no function call, return content
+            if not msg.function_call:
+                return msg.content
+
+            # Extract function call details
+            fn_name = msg.function_call.name
+            fn_args = json.loads(msg.function_call.arguments or '{}')
+
+            # Append assistant function_call message
+            openai_messages.append({
                 "role": "assistant",
                 "content": None,
-                "function_call": {"name": name, "arguments": args_raw}
+                "function_call": {
+                    "name": fn_name,
+                    "arguments": msg.function_call.arguments,
+                }
             })
 
-            # Parse arguments safely
+            # Call the actual tool
             try:
-                args = json.loads(args_raw or '{}')
-            except json.JSONDecodeError:
-                args = {}
+                override_args = {
+                    'id': state.task_id,
+                    'sessionId': state.task.sessionId
+                }
 
-            # Call the tool manager with name and parsed args
-            try:
-                tool_result = await self.tools_manager.call_tool(name, args)
+                tool_result = await self.tools_manager.call_tool(fn_name, fn_args, override_args)
             except Exception as e:
-                tool_result = {"content": f"Error calling {name}: {e}"}
+                tool_result = {"content": f"Error calling {fn_name}: {e}"}
+
+            # Case 1: Handle list of TextContent objects
+            if isinstance(tool_result, list) and len(tool_result) > 0:
+                # Extract text from the first TextContent item in the list
+                result = tool_result[0].text  # Access the `text` attribute
 
-            # Extract content
-            if hasattr(tool_result, 'content'):
-                result_content = tool_result.content
-            elif isinstance(tool_result, dict) and 'content' in tool_result:
-                result_content = tool_result['content']
+            # Case 2: Handle error case (dict with 'content' key)
+            elif isinstance(tool_result, dict):
+                result = tool_result.get('content', str(tool_result))
+
+            # Fallback for unexpected types
             else:
-                result_content = str(tool_result)
+                result = str(tool_result)
 
-            # Append the function/tool's response
-            converted_messages.append({
+            # Append function response
+            openai_messages.append({
                 "role": "function",
-                "name": name,
-                "content": result_content
+                "name": fn_name,
+                "content": result,
             })
 
         raise RuntimeError("Max tool iteration depth reached in generate().")
 
 
-
-    async def generate_stream(
-        self, messages: List[Message], **kwargs
-    ) -> AsyncGenerator[str, None]:
+    async def generate_stream(self, state: StateData, **kwargs) -> AsyncGenerator[str, None]:
         """
-        Stream response chunks, handling tool calls when complete.
+        Stream response chunks, invoking tools as needed.
         """
-        # Prepare messages including dynamic system prompt
-        converted_messages = self._convert_messages(messages)
-        max_iterations = 10
+        context_history = state.context_history
+        # Normalize incoming messages to your Message model
+        msgs = [
+            msg if isinstance(msg, Message) else Message(**msg)
+            for msg in context_history
+        ]
+        # Convert to OpenAI schema, including any prior tool results
+        converted_messages = self._convert_messages(msgs)
+        max_iterations = 30
 
         for _ in range(max_iterations):
-            # Start streaming completion with function-call support
+            # Kick off the streaming completion
             stream = await self.client.chat.completions.create(
                 model=self.model,
                 messages=converted_messages,
@@ -221,25 +247,28 @@ class OpenAIProvider(BaseLLMProvider):
                 **kwargs
             )
 
-            full_content = []
+            full_content = ""
             tool_calls: List[Dict[str, Any]] = []
 
-            # Collect streamed tokens and tool call deltas
+            # As chunks arrive, yield them and collect any tool_call deltas
             async for chunk in stream:
                 delta = chunk.choices[0].delta
-                # Yield text content immediately
-                if hasattr(delta, 'content') and delta.content:
-                    full_content.append(delta.content)
+
+                # 1) Stream content immediately
+                if hasattr(delta, "content") and delta.content:
                     yield delta.content
+                    full_content += delta.content
 
-                # Accumulate tool call metadata
-                if hasattr(delta, 'tool_calls') and delta.tool_calls:
+                # 2) Buffer up any function/tool calls for after the stream
+                if hasattr(delta, "tool_calls") and delta.tool_calls:
                     for d in delta.tool_calls:
                         idx = d.index
-                        # Ensure sufficient list length
+                        # Ensure list is long enough
                         while len(tool_calls) <= idx:
-                            tool_calls.append({"id": "", "function": {"name": "", "arguments": ""}})
-                        # Assign fields if present
+                            tool_calls.append({
+                                "id": "",
+                                "function": {"name": "", "arguments": ""}
+                            })
                         if d.id:
                             tool_calls[idx]["id"] = d.id
                         if d.function.name:
@@ -247,35 +276,39 @@ class OpenAIProvider(BaseLLMProvider):
                         if d.function.arguments:
                             tool_calls[idx]["function"]["arguments"] += d.function.arguments
 
-            # If no tool calls were invoked, stream is complete
+            # If the assistant didn't invoke any tools, we're done
             if not tool_calls:
                 return
 
-            # Append completed assistant message with tool calls
+            # Otherwise, append the assistant's outgoing call and loop for tool execution
             converted_messages.append({
                 "role": "assistant",
-                "content": "".join(full_content),
+                "content": full_content,
                 "tool_calls": [
-                    {"id": tc["id"],
-                     "type": "function",
-                     "function": {"name": tc["function"]["name"],
-                                   "arguments": tc["function"]["arguments"]}
+                    {
+                        "id": tc["id"],
+                        "type": "function",
+                        "function": {
+                            "name": tc["function"]["name"],
+                            "arguments": tc["function"]["arguments"]
+                        }
                     }
                     for tc in tool_calls
                 ]
             })
 
-            # Execute each tool call and append responses
+            # Execute each tool in turn and append its result
             for tc in tool_calls:
                 name = tc["function"]["name"]
                 try:
-                    args = json.loads(tc["function"]["arguments"])
+                    args = json.loads(tc["function"]["arguments"] or "{}")
                 except json.JSONDecodeError:
                     args = {}
-
                 try:
-                    result = await self.tools_manager.call_tool(name, args)
-                    result_content = result.content
+                    tool_res = await self.tools_manager.call_tool(name, args)
+                    result_content = getattr(tool_res, "content", None) or (
+                        tool_res.get("content") if isinstance(tool_res, dict) else str(tool_res)
+                    )
                 except Exception as e:
                     result_content = f"Error executing {name}: {e}"
 
@@ -284,5 +317,9 @@ class OpenAIProvider(BaseLLMProvider):
                     "content": result_content,
                     "tool_call_id": tc["id"]
                 })
-        # If iterations exhausted without final completion
-        raise RuntimeError("Max tool iteration depth reached in generate_stream().")
+
+        raise RuntimeError("Max tool iteration depth reached in generate_stream().")
+
+
+
+    

smarta2a/server/json_rpc_request_processor.py

@@ -0,0 +1,130 @@
+# Library imports
+from typing import Optional, Any
+from pydantic import BaseModel, ValidationError
+from fastapi import HTTPException
+
+# Local imports
+from smarta2a.utils.types import (
+    JSONRPCRequest, 
+    JSONRPCResponse,
+    SendTaskRequest,
+    SendTaskStreamingRequest,
+    GetTaskRequest,
+    CancelTaskRequest,
+    SetTaskPushNotificationRequest,
+    GetTaskPushNotificationRequest,
+    GetTaskResponse,
+    CancelTaskResponse,
+    SetTaskPushNotificationResponse,
+    GetTaskPushNotificationResponse,
+    TaskStatus,
+    TaskState
+)
+from smarta2a.server.state_manager import StateManager
+from smarta2a.server.request_handler import RequestHandler
+from smarta2a.server.handler_registry import HandlerRegistry
+from smarta2a.utils.types import (
+    TaskNotFoundError,
+    MethodNotFoundError,
+    InvalidParamsError,
+    UnsupportedOperationError,
+    InternalError,
+    InvalidRequestError
+)
+
+class JSONRPCRequestProcessor:
+    def __init__(self, registry: HandlerRegistry, state_manager: Optional[StateManager] = None):
+        self.request_handler = RequestHandler(registry, state_manager)
+        self.state_manager = state_manager
+
+    async def process_request(self, request: JSONRPCRequest) -> JSONRPCResponse:
+        
+        try:
+            method = request.method
+            params = request.params
+            
+
+            match method:
+                case "tasks/send":
+                    send_task_request = self._validate_request(request, SendTaskRequest)
+                    
+                    if self.state_manager:
+                        state_data = await self.state_manager.get_or_create_and_update_state(send_task_request.params.id, send_task_request.params.sessionId, send_task_request.params.message, send_task_request.params.metadata, send_task_request.params.pushNotification)
+                        return await self.request_handler.handle_send_task(send_task_request, state_data)
+                    else:
+                        return await self.request_handler.handle_send_task(send_task_request)
+                    
+                case "tasks/sendSubscribe":
+                    send_subscribe_request = self._validate_request(request, SendTaskStreamingRequest)
+                    if self.state_manager:
+                        state_data = await self.state_manager.get_or_create_and_update_state(send_subscribe_request.params.id, send_subscribe_request.params.sessionId, send_subscribe_request.params.message, send_subscribe_request.params.metadata, send_subscribe_request.params.pushNotification)
+                        return await self.request_handler.handle_subscribe_task(send_subscribe_request, state_data)
+                    else:
+                        return await self.request_handler.handle_subscribe_task(send_subscribe_request)
+                    
+                case "tasks/get":
+                    get_task_request = self._validate_request(request, GetTaskRequest)
+                    if self.state_manager:
+                        state_data = self.state_manager.get_state(get_task_request.id)
+                        if state_data:
+                            return GetTaskResponse(
+                                id=get_task_request.id,
+                                result=state_data.task
+                            )
+                        else:
+                            return JSONRPCResponse(id=request.id, error=TaskNotFoundError())
+                    else:
+                        return self.request_handler.handle_get_task(get_task_request)
+                    
+                case "tasks/cancel":
+                    cancel_task_request = self._validate_request(request, CancelTaskRequest)
+                    if self.state_manager:
+                        state_data = self.state_manager.get_state(cancel_task_request.id)
+                        if state_data:
+                            state_data.task.status = TaskStatus(state=TaskState.CANCELLED)
+                            self.state_manager.update_state(cancel_task_request.id, state_data)
+                            return CancelTaskResponse(id=cancel_task_request.id)
+                        else:
+                            return JSONRPCResponse(id=request.id, error=TaskNotFoundError())
+                    else:
+                        return self.request_handler.handle_cancel_task(cancel_task_request)
+                    
+                case "tasks/pushNotification/set":
+                    set_push_notification_request = self._validate_request(request, SetTaskPushNotificationRequest)
+                    if self.state_manager:
+                        state_data = self.state_manager.get_state(set_push_notification_request.id)
+                        if state_data:
+                            state_data.push_notification_config = set_push_notification_request.pushNotificationConfig
+                            self.state_manager.update_state(set_push_notification_request.id, state_data)
+                            return SetTaskPushNotificationResponse(id=set_push_notification_request.id, result=state_data.push_notification_config)
+                        else:
+                            return JSONRPCResponse(id=request.id, error=TaskNotFoundError())
+                    else:
+                        return self.request_handler.handle_set_notification(set_push_notification_request)
+                    
+                case "tasks/pushNotification/get":
+                    get_push_notification_request = self._validate_request(request, GetTaskPushNotificationRequest)
+                    if self.state_manager:
+                        state_data = self.state_manager.get_state(get_push_notification_request.id)
+                        if state_data:
+                            return GetTaskPushNotificationResponse(id=get_push_notification_request.id, result=state_data.push_notification_config)
+                        else:
+                            return JSONRPCResponse(id=request.id, error=TaskNotFoundError())
+                    else:
+                        return self.request_handler.handle_get_notification(get_push_notification_request)
+                    
+                case _:
+                    return JSONRPCResponse(id=request.id, error=MethodNotFoundError()).model_dump()
+                
+        except ValidationError as e:
+            return JSONRPCResponse(id=request.id, error=InvalidParamsError(data=e.errors())).model_dump()
+        except HTTPException as e:
+            err = UnsupportedOperationError() if e.status_code == 405 else InternalError(data=str(e))
+            return JSONRPCResponse(id=request.id, error=err).model_dump()
+
+        
+    def _validate_request(self, request: JSONRPCRequest, validation_schema: BaseModel) -> Any:
+        try:
+            return validation_schema.model_validate(request.model_dump())
+        except ValidationError as e:
+            return JSONRPCResponse(id=request.id, error=InvalidRequestError(data=e.errors())).model_dump()

smarta2a/server/nats_client.py

@@ -0,0 +1,49 @@
+# Library imports
+import asyncio
+import json
+from typing import Any, Dict
+from nats.aio.client import Client as NATS
+
+# Local imports
+
+class NATSClient:
+    def __init__(self, server_url: str = "nats://localhost:4222"):
+        self.server_url = server_url
+        self.nats = NATS()
+        self._connected = False
+
+    async def connect(self) -> None:
+        """Establishes an asynchronous connection to the NATS server."""
+        if not self._connected:
+            try:
+                # Use the current running loop by default
+                await self.nats.connect(self.server_url)
+                self._connected = True
+                print(f"Connected to NATS at {self.server_url}")
+            except Exception as e:
+                print(f"Failed to connect to NATS: {e}")
+                raise
+
+    async def publish(self, subject: str, payload: Dict[str, Any]) -> None:
+        """Publishes a JSON-encoded message to a NATS subject, auto-connecting if needed."""
+        if not self._connected:
+            # Ensure connection before publishing
+            await self.connect()
+
+        try:
+            data = json.dumps(payload).encode()
+            await self.nats.publish(subject, data)
+        except Exception as e:
+            print(f"Failed to publish message: {e}")
+            raise
+
+    async def close(self) -> None:
+        """Close NATS connection gracefully"""
+        if self._connected:
+            await self.nats.close()
+            self._connected = False
+            print("NATS connection closed")
+
+    @property
+    def is_connected(self) -> bool:
+        return self._connected