jaf-py 2.4.4__py3-none-any.whl → 2.4.6__py3-none-any.whl

jaf/core/parallel_agents.py

@@ -0,0 +1,339 @@
+"""
+Parallel Agent Execution for JAF Framework.
+
+This module provides functionality to execute multiple sub-agents in parallel groups,
+allowing for coordinated parallel execution with configurable grouping and result aggregation.
+"""
+
+import asyncio
+import json
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union, Callable, TypeVar
+
+from .types import (
+    Agent,
+    Tool,
+    ToolSchema,
+    ToolSource,
+    RunConfig,
+    RunState,
+    RunResult,
+    Message,
+    ContentRole,
+    generate_run_id,
+    generate_trace_id,
+)
+from .agent_tool import create_agent_tool, AgentToolInput
+
+Ctx = TypeVar('Ctx')
+Out = TypeVar('Out')
+
+
+@dataclass
+class ParallelAgentGroup:
+    """Configuration for a group of agents to be executed in parallel."""
+    name: str
+    agents: List[Agent[Ctx, Out]]
+    shared_input: bool = True  # Whether all agents receive the same input
+    result_aggregation: str = "combine"  # "combine", "first", "majority", "custom"
+    custom_aggregator: Optional[Callable[[List[str]], str]] = None
+    timeout: Optional[float] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+@dataclass
+class ParallelExecutionConfig:
+    """Configuration for parallel agent execution."""
+    groups: List[ParallelAgentGroup]
+    inter_group_execution: str = "sequential"  # "sequential" or "parallel"
+    global_timeout: Optional[float] = None
+    preserve_session: bool = False
+
+
+class ParallelAgentsTool:
+    """Tool that executes multiple agent groups in parallel."""
+    
+    def __init__(
+        self,
+        config: ParallelExecutionConfig,
+        tool_name: str = "execute_parallel_agents",
+        tool_description: str = "Execute multiple agents in parallel groups"
+    ):
+        self.config = config
+        self.tool_name = tool_name
+        self.tool_description = tool_description
+        
+        # Create tool schema
+        self.schema = ToolSchema(
+            name=tool_name,
+            description=tool_description,
+            parameters=AgentToolInput,
+            timeout=config.global_timeout
+        )
+        self.source = ToolSource.NATIVE
+        self.metadata = {"source": "parallel_agents", "groups": len(config.groups)}
+    
+    async def execute(self, args: AgentToolInput, context: Ctx) -> str:
+        """Execute all configured agent groups."""
+        try:
+            if self.config.inter_group_execution == "parallel":
+                # Execute all groups in parallel
+                group_results = await asyncio.gather(*[
+                    self._execute_group(group, args.input, context) 
+                    for group in self.config.groups
+                ])
+            else:
+                # Execute groups sequentially
+                group_results = []
+                for group in self.config.groups:
+                    result = await self._execute_group(group, args.input, context)
+                    group_results.append(result)
+            
+            # Combine results from all groups
+            final_result = {
+                "parallel_execution_results": {
+                    group.name: result for group, result in zip(self.config.groups, group_results)
+                },
+                "execution_mode": self.config.inter_group_execution,
+                "total_groups": len(self.config.groups)
+            }
+            
+            return json.dumps(final_result, indent=2)
+            
+        except Exception as e:
+            return json.dumps({
+                "error": "parallel_execution_failed",
+                "message": f"Failed to execute parallel agents: {str(e)}",
+                "groups_attempted": len(self.config.groups)
+            })
+    
+    async def _execute_group(
+        self, 
+        group: ParallelAgentGroup, 
+        input_text: str, 
+        context: Ctx
+    ) -> Dict[str, Any]:
+        """Execute a single group of agents in parallel."""
+        try:
+            # Create agent tools for all agents in the group
+            agent_tools = []
+            for agent in group.agents:
+                tool = create_agent_tool(
+                    agent=agent,
+                    tool_name=f"run_{agent.name.lower().replace(' ', '_')}",
+                    tool_description=f"Execute the {agent.name} agent",
+                    timeout=group.timeout,
+                    preserve_session=self.config.preserve_session
+                )
+                agent_tools.append((agent.name, tool))
+            
+            # Execute all agents in the group in parallel
+            if group.shared_input:
+                # All agents get the same input
+                tasks = [
+                    tool.execute(AgentToolInput(input=input_text), context)
+                    for _, tool in agent_tools
+                ]
+            else:
+                # This could be extended to support different inputs per agent
+                tasks = [
+                    tool.execute(AgentToolInput(input=input_text), context)
+                    for _, tool in agent_tools
+                ]
+            
+            # Execute with timeout if specified
+            if group.timeout:
+                results = await asyncio.wait_for(
+                    asyncio.gather(*tasks, return_exceptions=True),
+                    timeout=group.timeout
+                )
+            else:
+                results = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            # Process results
+            agent_results = {}
+            for (agent_name, _), result in zip(agent_tools, results):
+                if isinstance(result, Exception):
+                    agent_results[agent_name] = {
+                        "error": True,
+                        "message": str(result),
+                        "type": type(result).__name__
+                    }
+                else:
+                    agent_results[agent_name] = {
+                        "success": True,
+                        "result": result
+                    }
+            
+            # Apply result aggregation
+            aggregated_result = self._aggregate_results(group, agent_results)
+            
+            return {
+                "group_name": group.name,
+                "agent_count": len(group.agents),
+                "individual_results": agent_results,
+                "aggregated_result": aggregated_result,
+                "execution_time_ms": None  # Could be added with timing
+            }
+            
+        except asyncio.TimeoutError:
+            return {
+                "group_name": group.name,
+                "error": "timeout",
+                "message": f"Group {group.name} execution timed out after {group.timeout} seconds",
+                "agent_count": len(group.agents)
+            }
+        except Exception as e:
+            return {
+                "group_name": group.name,
+                "error": "execution_failed",
+                "message": str(e),
+                "agent_count": len(group.agents)
+            }
+    
+    def _aggregate_results(
+        self, 
+        group: ParallelAgentGroup, 
+        agent_results: Dict[str, Any]
+    ) -> Union[str, Dict[str, Any]]:
+        """Aggregate results from parallel agent execution."""
+        successful_results = [
+            result["result"] for result in agent_results.values() 
+            if result.get("success") and "result" in result
+        ]
+        
+        if not successful_results:
+            return {"error": "no_successful_results", "message": "All agents failed"}
+        
+        if group.result_aggregation == "first":
+            return successful_results[0]
+        elif group.result_aggregation == "combine":
+            return {
+                "combined_results": successful_results,
+                "result_count": len(successful_results)
+            }
+        elif group.result_aggregation == "majority":
+            # Simple majority logic - could be enhanced
+            if len(successful_results) >= len(group.agents) // 2 + 1:
+                return successful_results[0]  # Return first as majority representative
+            else:
+                return {"error": "no_majority", "results": successful_results}
+        elif group.result_aggregation == "custom" and group.custom_aggregator:
+            try:
+                return group.custom_aggregator(successful_results)
+            except Exception as e:
+                return {"error": "custom_aggregation_failed", "message": str(e)}
+        else:
+            return {"combined_results": successful_results}
+
+
+def create_parallel_agents_tool(
+    groups: List[ParallelAgentGroup],
+    tool_name: str = "execute_parallel_agents",
+    tool_description: str = "Execute multiple agents in parallel groups",
+    inter_group_execution: str = "sequential",
+    global_timeout: Optional[float] = None,
+    preserve_session: bool = False
+) -> Tool:
+    """
+    Create a tool that executes multiple agent groups in parallel.
+    
+    Args:
+        groups: List of parallel agent groups to execute
+        tool_name: Name of the tool
+        tool_description: Description of the tool
+        inter_group_execution: How to execute groups ("sequential" or "parallel")
+        global_timeout: Global timeout for all executions
+        preserve_session: Whether to preserve session across agent calls
+        
+    Returns:
+        A Tool that can execute parallel agent groups
+    """
+    config = ParallelExecutionConfig(
+        groups=groups,
+        inter_group_execution=inter_group_execution,
+        global_timeout=global_timeout,
+        preserve_session=preserve_session
+    )
+    
+    return ParallelAgentsTool(config, tool_name, tool_description)
+
+
+def create_simple_parallel_tool(
+    agents: List[Agent],
+    group_name: str = "parallel_group",
+    tool_name: str = "execute_parallel_agents",
+    shared_input: bool = True,
+    result_aggregation: str = "combine",
+    timeout: Optional[float] = None
+) -> Tool:
+    """
+    Create a simple parallel agents tool from a list of agents.
+    
+    Args:
+        agents: List of agents to execute in parallel
+        group_name: Name for the parallel group
+        tool_name: Name of the tool
+        shared_input: Whether all agents receive the same input
+        result_aggregation: How to aggregate results ("combine", "first", "majority")
+        timeout: Timeout for parallel execution
+        
+    Returns:
+        A Tool that executes all agents in parallel
+    """
+    group = ParallelAgentGroup(
+        name=group_name,
+        agents=agents,
+        shared_input=shared_input,
+        result_aggregation=result_aggregation,
+        timeout=timeout
+    )
+    
+    return create_parallel_agents_tool([group], tool_name=tool_name)
+
+
+# Convenience functions for common parallel execution patterns
+
+def create_language_specialists_tool(
+    language_agents: Dict[str, Agent],
+    tool_name: str = "consult_language_specialists",
+    timeout: Optional[float] = 300.0
+) -> Tool:
+    """Create a tool that consults multiple language specialists in parallel."""
+    group = ParallelAgentGroup(
+        name="language_specialists",
+        agents=list(language_agents.values()),
+        shared_input=True,
+        result_aggregation="combine",
+        timeout=timeout,
+        metadata={"languages": list(language_agents.keys())}
+    )
+    
+    return create_parallel_agents_tool(
+        [group], 
+        tool_name=tool_name,
+        tool_description="Consult multiple language specialists in parallel"
+    )
+
+
+def create_domain_experts_tool(
+    expert_agents: Dict[str, Agent],
+    tool_name: str = "consult_domain_experts",
+    result_aggregation: str = "combine",
+    timeout: Optional[float] = 60.0
+) -> Tool:
+    """Create a tool that consults multiple domain experts in parallel."""
+    group = ParallelAgentGroup(
+        name="domain_experts",
+        agents=list(expert_agents.values()),
+        shared_input=True,
+        result_aggregation=result_aggregation,
+        timeout=timeout,
+        metadata={"domains": list(expert_agents.keys())}
+    )
+    
+    return create_parallel_agents_tool(
+        [group],
+        tool_name=tool_name,
+        tool_description="Consult multiple domain experts in parallel"
+    )

jaf/core/streaming.py

@@ -209,20 +209,37 @@ async def run_streaming(
         trace_id=initial_state.trace_id
     )
 
-    tool_call_ids = {} # To map tool calls to their IDs
+    tool_call_ids: Dict[str, str] = {}  # Map call_id -> tool_name for in-flight tool calls
 
     def event_handler(event: TraceEvent) -> None:
         """Handle trace events and put them into the queue."""
         nonlocal tool_call_ids
         streaming_event = None
+        payload = event.data
+
+        def _get_event_value(keys: List[str]) -> Any:
+            for key in keys:
+                if isinstance(payload, dict) and key in payload:
+                    return payload[key]
+                if hasattr(payload, key):
+                    return getattr(payload, key)
+            return None
+
         if event.type == 'tool_call_start':
-            # Generate a unique ID for the tool call
-            call_id = f"call_{uuid.uuid4().hex[:8]}"
-            tool_call_ids[event.data.tool_name] = call_id
-            
+            tool_name = _get_event_value(['tool_name', 'toolName']) or 'unknown'
+            args = _get_event_value(['args', 'arguments'])
+            call_id = _get_event_value(['call_id', 'tool_call_id', 'toolCallId'])
+
+            if not call_id:
+                call_id = f"call_{uuid.uuid4().hex[:8]}"
+                if isinstance(payload, dict):
+                    payload['call_id'] = call_id
+
+            tool_call_ids[call_id] = tool_name
+
             tool_call = StreamingToolCall(
-                tool_name=event.data.tool_name,
-                arguments=event.data.args,
+                tool_name=tool_name,
+                arguments=args,
                 call_id=call_id,
                 status='started'
             )
@@ -233,18 +250,26 @@ async def run_streaming(
                 trace_id=initial_state.trace_id
             )
         elif event.type == 'tool_call_end':
-            if event.data.tool_name not in tool_call_ids:
-                raise RuntimeError(
-                    f"Tool call end event received for unknown tool '{event.data.tool_name}'. "
-                    f"Known tool calls: {list(tool_call_ids.keys())}. "
-                    f"This may indicate a missing tool_call_start event or a bug in the streaming implementation."
-                )
-            call_id = tool_call_ids[event.data.tool_name]
+            tool_name = _get_event_value(['tool_name', 'toolName']) or 'unknown'
+            call_id = _get_event_value(['call_id', 'tool_call_id', 'toolCallId'])
+
+            if not call_id:
+                # Fallback to locate a pending tool call with the same tool name
+                matching_call_id = next((cid for cid, name in tool_call_ids.items() if name == tool_name), None)
+                if matching_call_id:
+                    call_id = matching_call_id
+                else:
+                    raise RuntimeError(
+                        f"Tool call end event received for unknown tool '{tool_name}'. "
+                        f"Pending call IDs: {list(tool_call_ids.keys())}."
+                    )
+
+            tool_call_ids.pop(call_id, None)
             tool_result = StreamingToolResult(
-                tool_name=event.data.tool_name,
+                tool_name=tool_name,
                 call_id=call_id,
-                result=event.data.result,
-                status=event.data.status or 'completed'
+                result=_get_event_value(['result']),
+                status=_get_event_value(['status']) or 'completed'
             )
             streaming_event = StreamingEvent(
                 type=StreamingEventType.TOOL_RESULT,

jaf/core/tracing.py

@@ -10,6 +10,7 @@ import json
 import time
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Protocol
+import uuid
 
 from opentelemetry import trace
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
@@ -652,28 +653,36 @@ class LangfuseTraceCollector:
                 # Start a span for tool calls with detailed input information
                 tool_name = event.data.get('tool_name', 'unknown')
                 tool_args = event.data.get("args", {})
+                call_id = event.data.get("call_id")
+                if not call_id:
+                    call_id = f"{tool_name}-{uuid.uuid4().hex[:8]}"
+                    try:
+                        event.data["call_id"] = call_id
+                    except TypeError:
+                        # event.data may be immutable; log and rely on synthetic ID tracking downstream
+                        print(f"[LANGFUSE] Generated synthetic call_id for tool start: {call_id}")
                 
-                print(f"[LANGFUSE] Starting span for tool call: {tool_name}")
+                print(f"[LANGFUSE] Starting span for tool call: {tool_name} ({call_id})")
                 
                 # Track this tool call for the trace
                 tool_call_data = {
                     "tool_name": tool_name,
                     "arguments": tool_args,
-                    "call_id": event.data.get("call_id"),
+                    "call_id": call_id,
                     "timestamp": datetime.now().isoformat()
                 }
                 
                 # Ensure trace_id exists in tracking
                 if trace_id not in self.trace_tool_calls:
                     self.trace_tool_calls[trace_id] = []
-                
+
                 self.trace_tool_calls[trace_id].append(tool_call_data)
                 
                 # Create comprehensive input data for the tool call
                 tool_input = {
                     "tool_name": tool_name,
                     "arguments": tool_args,
-                    "call_id": event.data.get("call_id"),
+                    "call_id": call_id,
                     "timestamp": datetime.now().isoformat()
                 }
                 
@@ -682,7 +691,7 @@ class LangfuseTraceCollector:
                     input=tool_input,
                     metadata={
                         "tool_name": tool_name,
-                        "call_id": event.data.get("call_id"),
+                        "call_id": call_id,
                         "framework": "jaf",
                         "event_type": "tool_call"
                     }
@@ -696,14 +705,15 @@ class LangfuseTraceCollector:
                 if span_id in self.active_spans:
                     tool_name = event.data.get('tool_name', 'unknown')
                     tool_result = event.data.get("result")
+                    call_id = event.data.get("call_id")
                     
-                    print(f"[LANGFUSE] Ending span for tool call: {tool_name}")
+                    print(f"[LANGFUSE] Ending span for tool call: {tool_name} ({call_id})")
                     
                     # Track this tool result for the trace
                     tool_result_data = {
                         "tool_name": tool_name,
                         "result": tool_result,
-                        "call_id": event.data.get("call_id"),
+                        "call_id": call_id,
                         "timestamp": datetime.now().isoformat(),
                         "status": event.data.get("status", "completed"),
                         "tool_result": event.data.get("tool_result")
@@ -718,7 +728,7 @@ class LangfuseTraceCollector:
                     tool_output = {
                         "tool_name": tool_name,
                         "result": tool_result,
-                        "call_id": event.data.get("call_id"),
+                        "call_id": call_id,
                         "timestamp": datetime.now().isoformat(),
                         "status": event.data.get("status", "completed")
                     }
@@ -729,7 +739,7 @@ class LangfuseTraceCollector:
                         output=tool_output,
                         metadata={
                             "tool_name": tool_name,
-                            "call_id": event.data.get("call_id"),
+                            "call_id": call_id,
                             "result_length": len(str(tool_result)) if tool_result else 0,
                             "framework": "jaf",
                             "event_type": "tool_call_end"
@@ -791,6 +801,9 @@ class LangfuseTraceCollector:
         
         # Use consistent identifiers that don't depend on timestamp
         if event.type.startswith('tool_call'):
+            call_id = event.data.get('call_id') or event.data.get('tool_call_id')
+            if call_id:
+                return f"tool-{trace_id}-{call_id}"
             tool_name = event.data.get('tool_name') or event.data.get('toolName', 'unknown')
             return f"tool-{tool_name}-{trace_id}"
         elif event.type.startswith('llm_call'):

jaf/core/types.py

@@ -288,6 +288,7 @@ class Agent(Generic[Ctx, Out]):
     output_codec: Optional[Any] = None  # Type that can validate Out (like Pydantic model or Zod equivalent)
     handoffs: Optional[List[str]] = None
     model_config: Optional[ModelConfig] = None
+    advanced_config: Optional['AdvancedConfig'] = None
 
     def as_tool(
         self,
@@ -331,6 +332,74 @@ class Agent(Generic[Ctx, Out]):
 # Guardrail type
 Guardrail = Callable[[Any], Union[ValidationResult, Awaitable[ValidationResult]]]
 
+@dataclass(frozen=True)
+class AdvancedGuardrailsConfig:
+    """Configuration for advanced guardrails with LLM-based validation."""
+    input_prompt: Optional[str] = None
+    output_prompt: Optional[str] = None
+    require_citations: bool = False
+    fast_model: Optional[str] = None
+    fail_safe: Literal['allow', 'block'] = 'allow'
+    execution_mode: Literal['parallel', 'sequential'] = 'parallel'
+    timeout_ms: int = 30000
+
+    def __post_init__(self):
+        """Validate configuration."""
+        if self.timeout_ms < 1000:
+            object.__setattr__(self, 'timeout_ms', 1000)
+
+@dataclass(frozen=True)
+class AdvancedConfig:
+    """Advanced agent configuration including guardrails."""
+    guardrails: Optional[AdvancedGuardrailsConfig] = None
+
+def validate_guardrails_config(config: Optional[AdvancedGuardrailsConfig]) -> AdvancedGuardrailsConfig:
+    """Validate and provide defaults for guardrails configuration."""
+    if config is None:
+        return AdvancedGuardrailsConfig()
+    
+    return AdvancedGuardrailsConfig(
+        input_prompt=config.input_prompt.strip() if isinstance(config.input_prompt, str) and config.input_prompt else None,
+        output_prompt=config.output_prompt.strip() if isinstance(config.output_prompt, str) and config.output_prompt else None,
+        require_citations=config.require_citations,
+        fast_model=config.fast_model.strip() if isinstance(config.fast_model, str) and config.fast_model else None,
+        fail_safe=config.fail_safe,
+        execution_mode=config.execution_mode,
+        timeout_ms=max(1000, config.timeout_ms)
+    )
+
+def json_parse_llm_output(text: str) -> Optional[Dict[str, Any]]:
+    """Parse JSON from LLM output, handling common formatting issues."""
+    import json
+    import re
+    
+    if not text:
+        return None
+    
+    # Try direct parsing first
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        pass
+    
+    # Try to extract JSON from markdown code blocks
+    json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', text, re.DOTALL)
+    if json_match:
+        try:
+            return json.loads(json_match.group(1))
+        except json.JSONDecodeError:
+            pass
+    
+    # Try to find the first JSON object in the text
+    json_match = re.search(r'\{.*?\}', text, re.DOTALL)
+    if json_match:
+        try:
+            return json.loads(json_match.group(0))
+        except json.JSONDecodeError:
+            pass
+    
+    return None
+
 @dataclass(frozen=True)
 class ApprovalValue:
     """Represents an approval decision with context."""
@@ -541,11 +610,12 @@ class ToolCallStartEventData:
     args: Any
     trace_id: TraceId
     run_id: RunId
+    call_id: Optional[str] = None
 
 @dataclass(frozen=True)
 class ToolCallStartEvent:
     type: Literal['tool_call_start'] = 'tool_call_start'
-    data: ToolCallStartEventData = field(default_factory=lambda: ToolCallStartEventData("", None, TraceId(""), RunId("")))
+    data: ToolCallStartEventData = field(default_factory=lambda: ToolCallStartEventData("", None, TraceId(""), RunId(""), None))
 
 @dataclass(frozen=True)
 class ToolCallEndEventData:
@@ -556,11 +626,12 @@ class ToolCallEndEventData:
     run_id: RunId
     tool_result: Optional[Any] = None
     status: Optional[str] = None
+    call_id: Optional[str] = None
 
 @dataclass(frozen=True)
 class ToolCallEndEvent:
     type: Literal['tool_call_end'] = 'tool_call_end'
-    data: ToolCallEndEventData = field(default_factory=lambda: ToolCallEndEventData("", "", TraceId(""), RunId("")))
+    data: ToolCallEndEventData = field(default_factory=lambda: ToolCallEndEventData("", "", TraceId(""), RunId(""), None, None))
 
 @dataclass(frozen=True)
 class HandoffEventData:
@@ -598,6 +669,17 @@ class GuardrailEvent:
     type: Literal['guardrail_check'] = 'guardrail_check'
     data: GuardrailEventData = field(default_factory=lambda: GuardrailEventData(""))
 
+@dataclass(frozen=True)
+class GuardrailViolationEventData:
+    """Data for guardrail violation events."""
+    stage: Literal['input', 'output']
+    reason: str
+
+@dataclass(frozen=True)
+class GuardrailViolationEvent:
+    type: Literal['guardrail_violation'] = 'guardrail_violation'
+    data: GuardrailViolationEventData = field(default_factory=lambda: GuardrailViolationEventData("input", ""))
+
 @dataclass(frozen=True)
 class MemoryEventData:
     """Data for memory operation events."""
@@ -630,6 +712,7 @@ class OutputParseEvent:
 TraceEvent = Union[
     RunStartEvent,
     GuardrailEvent,
+    GuardrailViolationEvent,
     MemoryEvent,
     OutputParseEvent,
     LLMCallStartEvent,
@@ -708,7 +791,8 @@ class RunConfig(Generic[Ctx]):
     initial_input_guardrails: Optional[List[Guardrail]] = None
     final_output_guardrails: Optional[List[Guardrail]] = None
     on_event: Optional[Callable[[TraceEvent], None]] = None
-    memory: Optional['MemoryConfig'] = None
+    memory: Optional[Any] = None  # MemoryConfig - avoiding circular import
     conversation_id: Optional[str] = None
-    default_tool_timeout: Optional[float] = 30.0  # Default timeout for tool execution in seconds
-    approval_storage: Optional['ApprovalStorage'] = None  # Storage for approval decisions
+    default_fast_model: Optional[str] = None  # Default model for fast operations like guardrails
+    default_tool_timeout: Optional[float] = 300.0  # Default timeout for tool execution in seconds
+    approval_storage: Optional['ApprovalStorage'] = None  # Storage for approval decisions