jaf-py 2.5.10__py3-none-any.whl → 2.5.12__py3-none-any.whl

jaf/a2a/agent.py

@@ -27,7 +27,7 @@ def create_a2a_agent(
     description: str,
     instruction: str,
     tools: List[A2AAgentTool],
-    supported_content_types: Optional[List[str]] = None
+    supported_content_types: Optional[List[str]] = None,
 ) -> A2AAgent:
     """Pure function to create A2A compatible agent"""
     return A2AAgent(
@@ -35,7 +35,7 @@ def create_a2a_agent(
         description=description,
         supportedContentTypes=supported_content_types or ["text/plain", "application/json"],
         instruction=instruction,
-        tools=tools
+        tools=tools,
     )
 
 
@@ -43,14 +43,11 @@ def create_a2a_tool(
     name: str,
     description: str,
     parameters: Dict[str, Any],
-    execute_func: Callable[[Any, Optional[ToolContext]], Any]
+    execute_func: Callable[[Any, Optional[ToolContext]], Any],
 ) -> A2AAgentTool:
     """Pure function to create A2A tool"""
     return A2AAgentTool(
-        name=name,
-        description=description,
-        parameters=parameters,
-        execute=execute_func
+        name=name, description=description, parameters=parameters, execute=execute_func
     )
 
 
@@ -66,7 +63,7 @@ def create_initial_agent_state(session_id: str) -> AgentState:
         messages=[],
         context={},
         artifacts=[],
-        timestamp=datetime.now().isoformat()
+        timestamp=datetime.now().isoformat(),
     )
 
 
@@ -77,14 +74,14 @@ def add_message_to_state(state: AgentState, message: Any) -> AgentState:
         messages=[*state.messages, message],
         context=state.context,
         artifacts=state.artifacts,
-        timestamp=datetime.now().isoformat()
+        timestamp=datetime.now().isoformat(),
     )
 
 
 def update_state_from_run_result(state: AgentState, outcome: Any) -> AgentState:
     """Pure function to update state from run result"""
     new_artifacts = state.artifacts
-    if hasattr(outcome, 'artifacts') and outcome.artifacts:
+    if hasattr(outcome, "artifacts") and outcome.artifacts:
         new_artifacts = [*state.artifacts, *outcome.artifacts]
 
     return AgentState(
@@ -92,7 +89,7 @@ def update_state_from_run_result(state: AgentState, outcome: Any) -> AgentState:
         messages=state.messages,
         context=state.context,
         artifacts=new_artifacts,
-        timestamp=datetime.now().isoformat()
+        timestamp=datetime.now().isoformat(),
     )
 
 
@@ -110,24 +107,16 @@ def transform_a2a_agent_to_jaf(a2a_agent: A2AAgent) -> Agent:
     # Transform tools using functional approach
     jaf_tools = [transform_a2a_tool_to_jaf(a2a_tool) for a2a_tool in a2a_agent.tools]
 
-    return Agent(
-        name=a2a_agent.name,
-        instructions=instructions_func,
-        tools=jaf_tools
-    )
+    return Agent(name=a2a_agent.name, instructions=instructions_func, tools=jaf_tools)
 
 
-def transform_a2a_tool_to_jaf(a2a_tool: A2AAgentTool) -> 'JAFToolImplementation':
+def transform_a2a_tool_to_jaf(a2a_tool: A2AAgentTool) -> "JAFToolImplementation":
     """Pure function to transform A2A tool to JAF tool"""
 
     async def execute_wrapper(args: Any, context: Any) -> Any:
         tool_context = ToolContext(
-            actions={
-                "requiresInput": False,
-                "skipSummarization": False,
-                "escalate": False
-            },
-            metadata=context or {}
+            actions={"requiresInput": False, "skipSummarization": False, "escalate": False},
+            metadata=context or {},
         )
 
         result = await a2a_tool.execute(args, tool_context)
@@ -140,16 +129,12 @@ def transform_a2a_tool_to_jaf(a2a_tool: A2AAgentTool) -> 'JAFToolImplementation'
 
     # Create ToolSchema object
     from ..core.types import ToolSchema
+
     tool_schema = ToolSchema(
-        name=a2a_tool.name,
-        description=a2a_tool.description,
-        parameters=a2a_tool.parameters
+        name=a2a_tool.name, description=a2a_tool.description, parameters=a2a_tool.parameters
     )
 
-    return JAFToolImplementation(
-        schema=tool_schema,
-        execute=execute_wrapper
-    )
+    return JAFToolImplementation(schema=tool_schema, execute=execute_wrapper)
 
 
 class JAFToolImplementation:
@@ -170,21 +155,18 @@ class JAFToolImplementation:
         return await self._execute(args, context)
 
 
-def create_run_config_for_a2a_agent(
-    a2a_agent: A2AAgent,
-    model_provider: Any
-) -> RunConfig:
+def create_run_config_for_a2a_agent(a2a_agent: A2AAgent, model_provider: Any) -> RunConfig:
     """Pure function to create run configuration for A2A agent"""
     jaf_agent = transform_a2a_agent_to_jaf(a2a_agent)
 
     def event_handler(event: Any) -> None:
         # Handle different event types properly
-        if hasattr(event, 'data'):
+        if hasattr(event, "data"):
             event_data = event.data
             event_type = type(event).__name__
         elif isinstance(event, dict):
-            event_data = event.get('data', '')
-            event_type = event.get('type', 'unknown')
+            event_data = event.get("data", "")
+            event_type = event.get("type", "unknown")
         else:
             event_data = str(event)
             event_type = type(event).__name__
@@ -195,21 +177,19 @@ def create_run_config_for_a2a_agent(
         agent_registry={a2a_agent.name: jaf_agent},
         model_provider=model_provider,
         max_turns=10,
-        on_event=event_handler
+        on_event=event_handler,
     )
 
 
 def transform_to_run_state(
-    state: AgentState,
-    agent_name: str,
-    context: Optional[Dict[str, Any]] = None
+    state: AgentState, agent_name: str, context: Optional[Dict[str, Any]] = None
 ) -> RunState:
     """Pure function to transform agent state to JAF run state"""
     from ..core.types import generate_run_id, generate_trace_id
 
     # Convert A2A messages to JAF Message format
     jaf_messages = [
-        create_user_message(msg.content if hasattr(msg, 'content') else str(msg))
+        create_user_message(msg.content if hasattr(msg, "content") else str(msg))
         for msg in state.messages
         if msg is not None
     ]
@@ -220,15 +200,12 @@ def transform_to_run_state(
         messages=jaf_messages,
         current_agent_name=agent_name,
         context=context or {},
-        turn_count=0
+        turn_count=0,
     )
 
 
 async def process_agent_query(
-    agent: A2AAgent,
-    query: str,
-    state: AgentState,
-    model_provider: Any
+    agent: A2AAgent, query: str, state: AgentState, model_provider: Any
 ) -> AsyncGenerator[StreamEvent, None]:
     """Pure async generator function to process agent query with improved streaming"""
     start_time = time.time()
@@ -248,7 +225,7 @@ async def process_agent_query(
         new_state=new_state.model_dump(),
         timestamp=datetime.now().isoformat(),
         updates="Initializing agent state",
-        metrics={"start_time": start_time, "status": "starting"}
+        metrics={"start_time": start_time, "status": "starting"},
     )
 
     try:
@@ -266,19 +243,19 @@ async def process_agent_query(
             metrics={
                 "start_time": start_time,
                 "processing_time": processing_time,
-                "status": "processing"
-            }
+                "status": "processing",
+            },
         )
 
         # Execute JAF engine (pure function)
         result = await run(run_state, run_config)
         completion_time = time.time()
 
-        if hasattr(result.outcome, 'status') and result.outcome.status == 'completed':
+        if hasattr(result.outcome, "status") and result.outcome.status == "completed":
             final_state = update_state_from_run_result(new_state, result.outcome)
             yield StreamEvent(
                 isTaskComplete=True,
-                content=getattr(result.outcome, 'output', 'Task completed'),
+                content=getattr(result.outcome, "output", "Task completed"),
                 new_state=final_state.model_dump(),
                 timestamp=datetime.now().isoformat(),
                 updates="Task completed successfully",
@@ -286,12 +263,12 @@ async def process_agent_query(
                     "start_time": start_time,
                     "completion_time": completion_time,
                     "total_duration": completion_time - start_time,
-                    "status": "completed"
-                }
+                    "status": "completed",
+                },
             )
         else:
             final_state = update_state_from_run_result(new_state, result.outcome)
-            error_content = getattr(result.outcome, 'error', 'Unknown error')
+            error_content = getattr(result.outcome, "error", "Unknown error")
             yield StreamEvent(
                 isTaskComplete=True,
                 content=f"Error: {json.dumps(error_content) if isinstance(error_content, dict) else str(error_content)}",
@@ -302,14 +279,15 @@ async def process_agent_query(
                     "start_time": start_time,
                     "completion_time": completion_time,
                     "total_duration": completion_time - start_time,
-                    "status": "error"
-                }
+                    "status": "error",
+                },
             )
     except Exception as error:
         error_time = time.time()
 
         # Log the actual error for debugging but don't expose internal details
         import logging
+
         logging.error(f"Agent execution error for {agent.name}: {error!s}", exc_info=True)
 
         # Determine error type and provide safe error message
@@ -335,8 +313,8 @@ async def process_agent_query(
                 "error_time": error_time,
                 "total_duration": error_time - start_time,
                 "status": "failed",
-                "error_type": type(error).__name__
-            }
+                "error_type": type(error).__name__,
+            },
         )
 
 
@@ -346,19 +324,13 @@ def extract_text_from_a2a_message(message: Dict[str, Any]) -> str:
         return ""
 
     # Use functional approach instead of mutation
-    text_parts = [
-        part.get("text", "")
-        for part in message["parts"]
-        if part.get("kind") == "text"
-    ]
+    text_parts = [part.get("text", "") for part in message["parts"] if part.get("kind") == "text"]
 
     return "\n".join(text_parts)
 
 
 def create_a2a_text_message(
-    text: str,
-    context_id: str,
-    task_id: Optional[str] = None
+    text: str, context_id: str, task_id: Optional[str] = None
 ) -> Dict[str, Any]:
     """Pure function to create A2A text message"""
     return {
@@ -368,14 +340,12 @@ def create_a2a_text_message(
         "contextId": context_id,
         "taskId": task_id,
         "kind": "message",
-        "timestamp": datetime.now().isoformat()
+        "timestamp": datetime.now().isoformat(),
     }
 
 
 def create_a2a_data_message(
-    data: Dict[str, Any],
-    context_id: str,
-    task_id: Optional[str] = None
+    data: Dict[str, Any], context_id: str, task_id: Optional[str] = None
 ) -> Dict[str, Any]:
     """Pure function to create A2A data message"""
     return {
@@ -385,14 +355,11 @@ def create_a2a_data_message(
         "contextId": context_id,
         "taskId": task_id,
         "kind": "message",
-        "timestamp": datetime.now().isoformat()
+        "timestamp": datetime.now().isoformat(),
     }
 
 
-def create_a2a_task(
-    message: Dict[str, Any],
-    context_id: Optional[str] = None
-) -> Dict[str, Any]:
+def create_a2a_task(message: Dict[str, Any], context_id: Optional[str] = None) -> Dict[str, Any]:
     """Pure function to create A2A task"""
     task_id = f"task_{int(time.time() * 1000)}_{uuid.uuid4().hex[:8]}"
     current_time = datetime.now().isoformat()
@@ -400,42 +367,35 @@ def create_a2a_task(
     return {
         "id": task_id,
         "contextId": context_id or f"ctx_{int(time.time() * 1000)}",
-        "status": {
-            "state": "submitted",
-            "timestamp": current_time
-        },
+        "status": {"state": "submitted", "timestamp": current_time},
         "history": [message],
         "artifacts": [],
-        "kind": "task"
+        "kind": "task",
     }
 
 
 def update_a2a_task_status(
-    task: Dict[str, Any],
-    state: str,
-    message: Optional[Dict[str, Any]] = None
+    task: Dict[str, Any], state: str, message: Optional[Dict[str, Any]] = None
 ) -> Dict[str, Any]:
     """Pure function to update A2A task status"""
     updated_task = task.copy()
     updated_task["status"] = {
         "state": state,
         "message": message,
-        "timestamp": datetime.now().isoformat()
+        "timestamp": datetime.now().isoformat(),
     }
     return updated_task
 
 
 def add_artifact_to_a2a_task(
-    task: Dict[str, Any],
-    parts: List[Dict[str, Any]],
-    name: str
+    task: Dict[str, Any], parts: List[Dict[str, Any]], name: str
 ) -> Dict[str, Any]:
     """Pure function to add artifact to A2A task"""
     artifact = {
         "artifactId": f"artifact_{int(time.time() * 1000)}_{uuid.uuid4().hex[:8]}",
         "name": name,
         "parts": parts,
-        "timestamp": datetime.now().isoformat()
+        "timestamp": datetime.now().isoformat(),
     }
 
     updated_task = task.copy()
@@ -452,7 +412,7 @@ def complete_a2a_task(task: Dict[str, Any], result: Optional[Any] = None) -> Dic
             "artifactId": f"result_{int(time.time() * 1000)}_{uuid.uuid4().hex[:8]}",
             "name": "final_result",
             "parts": [{"kind": "text", "text": str(result)}],
-            "timestamp": datetime.now().isoformat()
+            "timestamp": datetime.now().isoformat(),
         }
 
         updated_task["artifacts"] = [*task.get("artifacts", []), result_artifact]
@@ -462,10 +422,9 @@ def complete_a2a_task(task: Dict[str, Any], result: Optional[Any] = None) -> Dic
 
 # Execution functions
 
+
 async def execute_a2a_agent(
-    context: Dict[str, Any],
-    agent: A2AAgent,
-    model_provider: Any
+    context: Dict[str, Any], agent: A2AAgent, model_provider: Any
 ) -> Dict[str, Any]:
     """Execute A2A agent and return result"""
     message = context.get("message", {})
@@ -495,40 +454,36 @@ async def execute_a2a_agent(
                     updated_task = update_a2a_task_status(
                         current_task,
                         "failed",
-                        create_a2a_text_message(error_message, session_id, current_task["id"])
+                        create_a2a_text_message(error_message, session_id, current_task["id"]),
                     )
 
                     return {
                         "events": [*events, *processing_events],
                         "final_task": updated_task,
-                        "error": error_message
+                        "error": error_message,
                     }
                 else:
                     # Success case
                     content = event.content
                     updated_task = add_artifact_to_a2a_task(
-                        current_task,
-                        [{"kind": "text", "text": str(content)}],
-                        "response"
+                        current_task, [{"kind": "text", "text": str(content)}], "response"
                     )
                     updated_task = complete_a2a_task(updated_task, content)
 
-                    return {
-                        "events": [*events, *processing_events],
-                        "final_task": updated_task
-                    }
+                    return {"events": [*events, *processing_events], "final_task": updated_task}
 
         # If we reach here without completion, mark as failed
         updated_task = update_a2a_task_status(current_task, "failed")
         return {
             "events": [*events, *processing_events],
             "final_task": updated_task,
-            "error": "Agent processing did not complete"
+            "error": "Agent processing did not complete",
         }
 
     except Exception as error:
         # Log the actual error for debugging but don't expose internal details
         import logging
+
         logging.error(f"Agent execution error for {agent.name}: {error!s}", exc_info=True)
 
         # Determine error type and provide safe error message
@@ -546,20 +501,14 @@ async def execute_a2a_agent(
         failed_task = update_a2a_task_status(
             current_task,
             "failed",
-            create_a2a_text_message(safe_error, session_id, current_task["id"])
+            create_a2a_text_message(safe_error, session_id, current_task["id"]),
         )
 
-        return {
-            "events": events,
-            "final_task": failed_task,
-            "error": safe_error
-        }
+        return {"events": events, "final_task": failed_task, "error": safe_error}
 
 
 async def execute_a2a_agent_with_streaming(
-    context: Dict[str, Any],
-    agent: A2AAgent,
-    model_provider: Any
+    context: Dict[str, Any], agent: A2AAgent, model_provider: Any
 ) -> AsyncGenerator[Dict[str, Any], None]:
     """Execute A2A agent with streaming"""
     message = context.get("message", {})
@@ -576,7 +525,7 @@ async def execute_a2a_agent_with_streaming(
             "taskId": current_task["id"],
             "contextId": current_task["contextId"],
             "status": {"state": "submitted", "timestamp": datetime.now().isoformat()},
-            "final": False
+            "final": False,
         }
 
     yield {
@@ -584,7 +533,7 @@ async def execute_a2a_agent_with_streaming(
         "taskId": current_task["id"],
         "contextId": current_task["contextId"],
         "status": {"state": "working", "timestamp": datetime.now().isoformat()},
-        "final": False
+        "final": False,
     }
 
     try:
@@ -601,11 +550,11 @@ async def execute_a2a_agent_with_streaming(
                         "message": create_a2a_text_message(
                             event.updates or "Processing...",
                             current_task["contextId"],
-                            current_task["id"]
+                            current_task["id"],
                         ),
-                        "timestamp": event.timestamp
+                        "timestamp": event.timestamp,
                     },
-                    "final": False
+                    "final": False,
                 }
             else:
                 # Handle final result
@@ -617,10 +566,12 @@ async def execute_a2a_agent_with_streaming(
                         "contextId": current_task["contextId"],
                         "status": {
                             "state": "failed",
-                            "message": create_a2a_text_message(error_message, current_task["contextId"], current_task["id"]),
-                            "timestamp": event.timestamp
+                            "message": create_a2a_text_message(
+                                error_message, current_task["contextId"], current_task["id"]
+                            ),
+                            "timestamp": event.timestamp,
                         },
-                        "final": True
+                        "final": True,
                     }
                 else:
                     yield {
@@ -630,8 +581,8 @@ async def execute_a2a_agent_with_streaming(
                         "artifact": {
                             "artifactId": f"result_{int(time.time() * 1000)}",
                             "name": "response",
-                            "parts": [{"kind": "text", "text": str(event.content)}]
-                        }
+                            "parts": [{"kind": "text", "text": str(event.content)}],
+                        },
                     }
 
                     yield {
@@ -639,7 +590,7 @@ async def execute_a2a_agent_with_streaming(
                         "taskId": current_task["id"],
                         "contextId": current_task["contextId"],
                         "status": {"state": "completed", "timestamp": event.timestamp},
-                        "final": True
+                        "final": True,
                     }
                 break
 
@@ -651,8 +602,10 @@ async def execute_a2a_agent_with_streaming(
             "contextId": current_task["contextId"],
             "status": {
                 "state": "failed",
-                "message": create_a2a_text_message(error_message, current_task["contextId"], current_task["id"]),
-                "timestamp": datetime.now().isoformat()
+                "message": create_a2a_text_message(
+                    error_message, current_task["contextId"], current_task["id"]
+                ),
+                "timestamp": datetime.now().isoformat(),
             },
-            "final": True
+            "final": True,
         }