smarta2a 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

smarta2a/server/server.py

@@ -1,3 +1,4 @@
+# Library imports
 from typing import Callable, Any, Optional, Dict, Union, List, AsyncGenerator
 import json
 from datetime import datetime
@@ -10,7 +11,15 @@ import uvicorn
 from fastapi.responses import StreamingResponse
 from uuid import uuid4
 
-from smarta2a.common.types import (
+# Local imports
+from smarta2a.server.handler_registry import HandlerRegistry
+from smarta2a.server.state_manager import StateManager
+from smarta2a.state_stores.base_state_store import BaseStateStore
+from smarta2a.history_update_strategies.history_update_strategy import HistoryUpdateStrategy
+from smarta2a.history_update_strategies.append_strategy import AppendStrategy
+from smarta2a.utils.task_builder import TaskBuilder
+
+from smarta2a.utils.types import (
     JSONRPCResponse,
     Task,
     Artifact,
@@ -19,6 +28,7 @@ from smarta2a.common.types import (
     FileContent,
     DataPart,
     Part,
+    Message,
     TaskStatus,
     TaskState,
     JSONRPCError,
@@ -50,41 +60,78 @@ from smarta2a.common.types import (
     SetTaskPushNotificationResponse,
     GetTaskPushNotificationResponse,
     TaskPushNotificationConfig,
+    StateData
 )
 
 class SmartA2A:
-    def __init__(self, name: str, **fastapi_kwargs):
+    def __init__(self, name: str, state_store: Optional[BaseStateStore] = None, history_strategy: HistoryUpdateStrategy = AppendStrategy(), **fastapi_kwargs):
         self.name = name
-        self.handlers: Dict[str, Callable] = {}
-        self.subscriptions: Dict[str, Callable] = {}
+        self.registry = HandlerRegistry()
+        self.state_mgr = StateManager(state_store, history_strategy)
         self.app = FastAPI(title=name, **fastapi_kwargs)
         self.router = APIRouter()
-        self._registered_decorators = set()
+        self.state_store = state_store
+        self.history_strategy = history_strategy
         self._setup_routes()
         self.server_config = {
             "host": "0.0.0.0",
             "port": 8000,
             "reload": False
         }
+        self.task_builder = TaskBuilder(default_status=TaskState.COMPLETED)
+
+    # Add this method to delegate ASGI calls
+    async def __call__(self, scope, receive, send):
+        return await self.app(scope, receive, send)
         
+    def on_send_task(self):
+        def decorator(func: Callable[[SendTaskRequest, Optional[StateData]], Any]) -> Callable:
+            self.registry.register("tasks/send", func)
+            return func
+        return decorator
+    
+    def on_send_subscribe_task(self):
+        def decorator(fn: Callable[[SendTaskStreamingRequest, Optional[StateData]], Any]):
+            self.registry.register("tasks/sendSubscribe", fn, subscription=True)
+            return fn
+        return decorator
+    
+    def task_get(self):
+        def decorator(fn: Callable[[GetTaskRequest], Any]):
+            self.registry.register("tasks/get", fn)
+            return fn
+        return decorator
+
+    def task_cancel(self):
+        def decorator(fn: Callable[[CancelTaskRequest], Any]):
+            self.registry.register("tasks/cancel", fn)
+            return fn
+        return decorator
+
+    def set_notification(self):
+        def decorator(fn: Callable[[SetTaskPushNotificationRequest], Any]):
+            self.registry.register("tasks/pushNotification/set", fn)
+            return fn
+        return decorator
+
+    def get_notification(self):
+        def decorator(fn: Callable[[GetTaskPushNotificationRequest], Any]):
+            self.registry.register("tasks/pushNotification/get", fn)
+            return fn
+        return decorator
+    
 
     def _setup_routes(self):
         @self.app.post("/")    
         async def handle_request(request: Request):
             try:
                 data = await request.json()
-                request_obj = JSONRPCRequest(**data)
+                req = JSONRPCRequest.model_validate(data)
+                #request_obj = JSONRPCRequest(**data)
             except Exception as e:
-                return JSONRPCResponse(
-                    id=None,
-                    error=JSONRPCError(
-                        code=-32700,
-                        message="Parse error",
-                        data=str(e)
-                    )
-                ).model_dump()
-            
-            response = await self.process_request(request_obj.model_dump())
+                return JSONRPCResponse(id=None, error=JSONRPCError(code=-32700, message="Parse error", data=str(e))).model_dump()
+                
+            response = await self.process_request(req)
 
             # <-- Accept both SSE‐style responses:
             if isinstance(response, (EventSourceResponse, StreamingResponse)):
@@ -92,114 +139,117 @@ class SmartA2A:
 
             # <-- Everything else is a normal pydantic JSONRPCResponse
             return response.model_dump()
-
-    def _register_handler(self, method: str, func: Callable, handler_name: str, handler_type: str = "handler"):
-        """Shared registration logic with duplicate checking"""
-        if method in self._registered_decorators:
-            raise RuntimeError(
-                f"@{handler_name} decorator for method '{method}' "
-                f"can only be used once per SmartA2A instance"
-            )
-        
-        if handler_type == "handler":
-            self.handlers[method] = func
-        else:
-            self.subscriptions[method] = func
-            
-        self._registered_decorators.add(method)
-
-    def on_send_task(self) -> Callable:
-        def decorator(func: Callable[[SendTaskRequest], Any]) -> Callable:
-            self._register_handler("tasks/send", func, "on_send_task", "handler")
-            return func
-        return decorator
-
-    def on_send_subscribe_task(self) -> Callable:
-        def decorator(func: Callable) -> Callable:
-            self._register_handler("tasks/sendSubscribe", func, "on_send_subscribe_task", "subscription")
-            return func
-        return decorator
-    
-    def task_get(self):
-        def decorator(func: Callable[[GetTaskRequest], Task]):
-            self._register_handler("tasks/get", func, "task_get", "handler")
-            return func
-        return decorator
     
-    def task_cancel(self):
-        def decorator(func: Callable[[CancelTaskRequest], Task]):
-            self._register_handler("tasks/cancel", func, "task_cancel", "handler")
-            return func
-        return decorator
-    
-    def set_notification(self):
-        def decorator(func: Callable[[SetTaskPushNotificationRequest], None]) -> Callable:
-            self._register_handler("tasks/pushNotification/set", func, "set_notification", "handler")
-            return func
-        return decorator
-    
-    def get_notification(self):
-        def decorator(func: Callable[[GetTaskPushNotificationRequest], Union[TaskPushNotificationConfig, GetTaskPushNotificationResponse]]):
-            self._register_handler("tasks/pushNotification/get", func, "get_notification", "handler")
-            return func
-        return decorator
 
-    async def process_request(self, request_data: dict) -> JSONRPCResponse:
+    async def process_request(self, request: JSONRPCRequest) -> JSONRPCResponse:
+        
         try:
-            method = request_data.get("method")
+            method = request.method
+            params = request.params
+            state_store = self.state_mgr.get_store()
             if method == "tasks/send":
-                return self._handle_send_task(request_data)
+                state_data = self.state_mgr.init_or_get(params.get("sessionId"), params.get("message"), params.get("metadata") or {})
+                if state_store:
+                    return self._handle_send_task(request, state_data)
+                else:
+                    return self._handle_send_task(request)
             elif method == "tasks/sendSubscribe":
-                return await self._handle_subscribe_task(request_data)
+                state_data = self.state_mgr.init_or_get(params.get("sessionId"), params.get("message"), params.get("metadata") or {})
+                if state_store:
+                    return await self._handle_subscribe_task(request, state_data)
+                else:
+                    return await self._handle_subscribe_task(request)
             elif method == "tasks/get":
-                return self._handle_get_task(request_data)
+                return self._handle_get_task(request)
             elif method == "tasks/cancel":
-                return self._handle_cancel_task(request_data)
+                return self._handle_cancel_task(request)
             elif method == "tasks/pushNotification/set":
-                return self._handle_set_notification(request_data)
+                return self._handle_set_notification(request)
             elif method == "tasks/pushNotification/get":
-                return self._handle_get_notification(request_data)
+                return self._handle_get_notification(request)
             else:
-                return self._error_response(
-                    request_data.get("id"),
-                    -32601,
-                    "Method not found"
-                )
+                return JSONRPCResponse(id=request.id, error=MethodNotFoundError()).model_dump() 
         except ValidationError as e:
-            return self._error_response(
-                request_data.get("id"),
-                -32600,
-                "Invalid params",
-                e.errors()
-            )
+                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 _handle_send_task(self, request_data: dict) -> SendTaskResponse:
+    def _handle_send_task(self, request_data: JSONRPCRequest, state_data: Optional[StateData] = None) -> SendTaskResponse:
         try:
             # Validate request format
-            request = SendTaskRequest.model_validate(request_data)
-            handler = self.handlers.get("tasks/send")
+            request = SendTaskRequest.model_validate(request_data.model_dump())
+            handler = self.registry.get_handler("tasks/send")
             
             if not handler:
                 return SendTaskResponse(
                     id=request.id,
                     error=MethodNotFoundError()
                 )
+            
+            user_message = request.params.message
+            request_metadata = request.params.metadata or {}
+            if state_data:
+                session_id = state_data.sessionId
+                existing_history = state_data.history.copy() or []
+                metadata = state_data.metadata or {} # Request metadata has already been merged so need to do it here
+            else:
+                session_id = request.params.sessionId or str(uuid4())
+                existing_history = [user_message]
+                metadata = request_metadata
+
 
             try:
-                raw_result = handler(request)
-            
+
+                if state_data:
+                    raw_result = handler(request, state_data)
+                else:
+                    raw_result = handler(request)
+
+                # Handle direct SendTaskResponse returns
                 if isinstance(raw_result, SendTaskResponse):
                     return raw_result
 
-                # Use unified task builder
-                task = self._build_task(
+                # Build task with updated history (before agent response)
+                task = self.task_builder.build(
                     content=raw_result,
                     task_id=request.params.id,
-                    session_id=request.params.sessionId,
-                    default_status=TaskState.COMPLETED,
-                    metadata=request.params.metadata or {}
+                    session_id=session_id,  # Always use generated session ID
+                    metadata=metadata,  # Use merged metadata
+                    history=existing_history  # History
                 )
 
+                # Process messages through strategy
+                messages = []
+                if task.artifacts:
+                    agent_parts = [p for a in task.artifacts for p in a.parts]
+                    agent_message = Message(
+                        role="agent",
+                        parts=agent_parts,
+                        metadata=task.metadata
+                    )
+                    messages.append(agent_message)
+
+                final_history = self.history_strategy.update_history(
+                    existing_history=existing_history,
+                    new_messages=messages
+                )
+
+                # Update task with final state
+                task.history = final_history
+
+                # State store update (if enabled)
+                if self.state_store:
+                    self.state_store.update_state(
+                        session_id=session_id,
+                        state_data=StateData(
+                            sessionId=session_id,
+                            history=final_history,
+                            metadata=metadata  # Use merged metadata
+                        )
+                    )
+
                 return SendTaskResponse(
                     id=request.id,
                     result=task
@@ -229,10 +279,11 @@ class SmartA2A:
             )
 
         
-    async def _handle_subscribe_task(self, request_data: dict) -> Union[EventSourceResponse, SendTaskStreamingResponse]:
+    async def _handle_subscribe_task(self, request_data: JSONRPCRequest, state_data: Optional[StateData] = None) -> Union[EventSourceResponse, SendTaskStreamingResponse]:
         try:
-            request = SendTaskStreamingRequest.model_validate(request_data)
-            handler = self.subscriptions.get("tasks/sendSubscribe")
+            request = SendTaskStreamingRequest.model_validate(request_data.model_dump())
+            #handler = self.subscriptions.get("tasks/sendSubscribe")
+            handler = self.registry.get_subscription("tasks/sendSubscribe")
             
             if not handler:
                 return SendTaskStreamingResponse(
@@ -240,15 +291,74 @@ class SmartA2A:
                     id=request.id,
                     error=MethodNotFoundError()
                 )
+            
+            user_message = request.params.message
+            request_metadata = request.params.metadata or {}
+            if state_data:
+                session_id = state_data.sessionId
+                existing_history = state_data.history.copy() or []
+                metadata = state_data.metadata or {} # Request metadata has already been merged so need to do it here
+            else:
+                session_id = request.params.sessionId or str(uuid4())
+                existing_history = [user_message]
+                metadata = request_metadata
+
 
             async def event_generator():
                 
                 try:
-                    raw_events = handler(request)
+                    
+                    if state_data:
+                        raw_events = handler(request, state_data)
+                    else:
+                        raw_events = handler(request)
+
                     normalized_events = self._normalize_subscription_events(request.params, raw_events)
 
+                    # Initialize streaming state
+                    stream_history = existing_history.copy()
+                    stream_metadata = metadata.copy()
+
                     async for item in normalized_events:
                         try:
+
+                            # Process artifact updates
+                            if isinstance(item, TaskArtifactUpdateEvent):
+                                # Create agent message from artifact parts
+                                agent_message = Message(
+                                    role="agent",
+                                    parts=[p for p in item.artifact.parts],
+                                    metadata=item.artifact.metadata
+                                )
+                                
+                                # Update history using strategy
+                                new_history = self.history_strategy.update_history(
+                                    existing_history=stream_history,
+                                    new_messages=[agent_message]
+                                )
+                                
+                                # Merge metadata
+                                new_metadata = {
+                                    **stream_metadata,
+                                    **(item.artifact.metadata or {})
+                                }
+
+                                # Update state store if configured
+                                if self.state_store:
+                                    self.state_store.update_state(
+                                        session_id=session_id,
+                                        state_data=StateData(
+                                            sessionId=session_id,
+                                            history=new_history,
+                                            metadata=new_metadata
+                                        )
+                                    )
+
+                                # Update streaming state
+                                stream_history = new_history
+                                stream_metadata = new_metadata
+
+                                
                             if isinstance(item, SendTaskStreamingResponse):
                                 yield item.model_dump_json()
                                 continue
@@ -318,11 +428,11 @@ class SmartA2A:
             )
 
 
-    def _handle_get_task(self, request_data: dict) -> GetTaskResponse:
+    def _handle_get_task(self, request_data: JSONRPCRequest) -> GetTaskResponse:
         try:
             # Validate request structure
-            request = GetTaskRequest.model_validate(request_data)
-            handler = self.handlers.get("tasks/get")
+            request = GetTaskRequest.model_validate(request_data.model_dump())
+            handler = self.registry.get_handler("tasks/get")
             
             if not handler:
                 return GetTaskResponse(
@@ -337,11 +447,10 @@ class SmartA2A:
                     return self._validate_response_id(raw_result, request)
 
                 # Use unified task builder with different defaults
-                task = self._build_task(
+                task = self.task_builder.build(
                     content=raw_result,
                     task_id=request.params.id,
-                    default_status=TaskState.COMPLETED,
-                    metadata=request.params.metadata or {}
+                    metadata=getattr(raw_result, "metadata", {}) or {}
                 )
 
                 return self._finalize_task_response(request, task)
@@ -370,11 +479,11 @@ class SmartA2A:
             )
         
 
-    def _handle_cancel_task(self, request_data: dict) -> CancelTaskResponse:
+    def _handle_cancel_task(self, request_data: JSONRPCRequest) -> CancelTaskResponse:
         try:
             # Validate request structure
-            request = CancelTaskRequest.model_validate(request_data)
-            handler = self.handlers.get("tasks/cancel")
+            request = CancelTaskRequest.model_validate(request_data.model_dump())
+            handler = self.registry.get_handler("tasks/cancel")
             
             if not handler:
                 return CancelTaskResponse(
@@ -385,24 +494,26 @@ class SmartA2A:
             try:
                 raw_result = handler(request)
             
+                cancel_task_builder = TaskBuilder(default_status=TaskState.CANCELED)
                 # Handle direct CancelTaskResponse returns
                 if isinstance(raw_result, CancelTaskResponse):
                     return self._validate_response_id(raw_result, request)
 
                 # Handle A2AStatus returns
                 if isinstance(raw_result, A2AStatus):
-                    task = self._build_task_from_status(
-                        status=raw_result,
+                    task = cancel_task_builder.normalize_from_status(
+                        status=raw_result.status,
                         task_id=request.params.id,
-                        metadata=raw_result.metadata or {}
+                        metadata=getattr(raw_result, "metadata", {}) or {}
                     )
                 else:
                     # Existing processing for other return types
-                    task = self._build_task(
+                    task = cancel_task_builder.build(
                         content=raw_result,
                         task_id=request.params.id,
-                        metadata=raw_result.metadata or {}
+                        metadata=getattr(raw_result, "metadata", {}) or {}
                     )
+                    print(task)
 
                 # Final validation and packaging
                 return self._finalize_cancel_response(request, task)
@@ -440,10 +551,10 @@ class SmartA2A:
                 error=InternalError(data=str(e))
             )
 
-    def _handle_set_notification(self, request_data: dict) -> SetTaskPushNotificationResponse:
+    def _handle_set_notification(self, request_data: JSONRPCRequest) -> SetTaskPushNotificationResponse:
         try:
-            request = SetTaskPushNotificationRequest.model_validate(request_data)
-            handler = self.handlers.get("tasks/pushNotification/set")
+            request = SetTaskPushNotificationRequest.model_validate(request_data.model_dump())
+            handler = self.registry.get_handler("tasks/pushNotification/set")
             
             if not handler:
                 return SetTaskPushNotificationResponse(
@@ -485,10 +596,10 @@ class SmartA2A:
             )
                       
 
-    def _handle_get_notification(self, request_data: dict) -> GetTaskPushNotificationResponse:
+    def _handle_get_notification(self, request_data: JSONRPCRequest) -> GetTaskPushNotificationResponse:
         try:
-            request = GetTaskPushNotificationRequest.model_validate(request_data)
-            handler = self.handlers.get("tasks/pushNotification/get")
+            request = GetTaskPushNotificationRequest.model_validate(request_data.model_dump())
+            handler = self.registry.get_handler("tasks/pushNotification/get")
             
             if not handler:
                 return GetTaskPushNotificationResponse(
@@ -535,139 +646,6 @@ class SmartA2A:
                 error=JSONParseError(data=str(e))
             )
 
-
-    def _normalize_artifacts(self, content: Any) -> List[Artifact]:
-        """Handle both A2AResponse content and regular returns"""
-        if isinstance(content, Artifact):
-            return [content]
-        
-        if isinstance(content, list):
-            # Handle list of artifacts
-            if all(isinstance(item, Artifact) for item in content):
-                return content
-            
-            # Handle mixed parts in list
-            parts = []
-            for item in content:
-                if isinstance(item, Artifact):
-                    parts.extend(item.parts)
-                else:
-                    parts.append(self._create_part(item))
-            return [Artifact(parts=parts)]
-        
-        # Handle single part returns
-        if isinstance(content, (str, Part, dict)):
-            return [Artifact(parts=[self._create_part(content)])]
-        
-        # Handle raw artifact dicts
-        try:
-            return [Artifact.model_validate(content)]
-        except ValidationError:
-            return [Artifact(parts=[TextPart(text=str(content))])]
-
-
-    def _build_task(
-    self,
-    content: Any,
-    task_id: str,
-    session_id: Optional[str] = None,
-    default_status: TaskState = TaskState.COMPLETED,
-    metadata: Optional[dict] = None
-) -> Task:
-        """Universal task construction from various return types."""
-        if isinstance(content, Task):
-            return content
-        
-        # Handle A2AResponse for sendTask case
-        if isinstance(content, A2AResponse):
-            status = content.status if isinstance(content.status, TaskStatus) \
-                else TaskStatus(state=content.status)
-            artifacts = self._normalize_content(content.content)
-            return Task(
-                id=task_id,
-                sessionId=session_id or str(uuid4()),  # Generate if missing
-                status=status,
-                artifacts=artifacts,
-                metadata=metadata or {}
-            )
-
-        try:  # Attempt direct validation for dicts
-            return Task.model_validate(content)
-        except ValidationError:
-            pass
-
-        # Fallback to content normalization
-        artifacts = self._normalize_content(content)
-        return Task(
-            id=task_id,
-            sessionId=session_id,
-            status=TaskStatus(state=default_status),
-            artifacts=artifacts,
-            metadata=metadata or {}
-        )
-    
-    def _build_task_from_status(self, status: A2AStatus, task_id: str, metadata: dict) -> Task:
-        """Convert A2AStatus to a Task with proper cancellation state."""
-        return Task(
-            id=task_id,
-            status=TaskStatus(
-                state=TaskState(status.status),
-                timestamp=datetime.now()
-            ),
-            metadata=metadata,
-            # Include empty/default values for required fields
-            sessionId="",  
-            artifacts=[],
-            history=[]
-        )
-
-
-    def _normalize_content(self, content: Any) -> List[Artifact]:
-        """Handle all content types for both sendTask and getTask cases."""
-        if isinstance(content, Artifact):
-            return [content]
-        
-        if isinstance(content, list):
-            if all(isinstance(item, Artifact) for item in content):
-                return content
-            return [Artifact(parts=self._parts_from_mixed(content))]
-        
-        if isinstance(content, (str, Part, dict)):
-            return [Artifact(parts=[self._create_part(content)])]
-        
-        try:  # Handle raw artifact dicts
-            return [Artifact.model_validate(content)]
-        except ValidationError:
-            return [Artifact(parts=[TextPart(text=str(content))])]
-
-    def _parts_from_mixed(self, items: List[Any]) -> List[Part]:
-        """Extract parts from mixed content lists."""
-        parts = []
-        for item in items:
-            if isinstance(item, Artifact):
-                parts.extend(item.parts)
-            else:
-                parts.append(self._create_part(item))
-        return parts
-
-
-    def _create_part(self, item: Any) -> Part:
-        """Convert primitive types to proper Part models"""
-        if isinstance(item, (TextPart, FilePart, DataPart)):
-            return item
-        
-        if isinstance(item, str):
-            return TextPart(text=item)
-        
-        if isinstance(item, dict):
-            try:
-                return Part.model_validate(item)
-            except ValidationError:
-                return TextPart(text=str(item))
-        
-        return TextPart(text=str(item))
-    
-
     # Response validation helper
     def _validate_response_id(self, response: Union[SendTaskResponse, GetTaskResponse], request) -> Union[SendTaskResponse, GetTaskResponse]:
         if response.result and response.result.id != request.params.id:
@@ -678,7 +656,7 @@ class SmartA2A:
                 )
             )
         return response
-    
+
     # Might refactor this later 
     def _finalize_task_response(self, request: GetTaskRequest, task: Task) -> GetTaskResponse:
         """Final validation and processing for getTask responses."""
@@ -721,8 +699,8 @@ class SmartA2A:
             id=request.id,
             result=task
         )
-    
-    
+
+
     async def _normalize_subscription_events(self, params: TaskSendParams, events: AsyncGenerator) -> AsyncGenerator[Union[SendTaskStreamingResponse, TaskStatusUpdateEvent, TaskArtifactUpdateEvent], None]:
         artifact_state = defaultdict(lambda: {"index": 0, "last_chunk": False})