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

jaf/core/__init__.py

@@ -13,6 +13,14 @@ from .agent_tool import (
     get_current_run_config,
     set_current_run_config,
 )
+from .parallel_agents import (
+    ParallelAgentGroup,
+    ParallelExecutionConfig,
+    create_parallel_agents_tool,
+    create_simple_parallel_tool,
+    create_language_specialists_tool,
+    create_domain_experts_tool,
+)
 from .proxy import ProxyConfig, ProxyAuth, create_proxy_config, get_default_proxy_config
 
 __all__ = [
@@ -23,6 +31,8 @@ __all__ = [
     "Message",
     "ModelConfig",
     "ModelProvider",
+    "ParallelAgentGroup",
+    "ParallelExecutionConfig",
     "ProxyAuth",
     "ProxyConfig",
     "RunConfig",
@@ -41,9 +51,13 @@ __all__ = [
     "create_agent_tool",
     "create_conditional_enabler",
     "create_default_output_extractor",
+    "create_domain_experts_tool",
     "create_json_output_extractor",
+    "create_language_specialists_tool",
+    "create_parallel_agents_tool",
     "create_proxy_config",
     "create_run_id",
+    "create_simple_parallel_tool",
     "create_trace_id",
     "get_current_run_config",
     "get_default_proxy_config",

jaf/core/engine.py

@@ -32,6 +32,8 @@ from .types import (
     Interruption,
     GuardrailEvent,
     GuardrailEventData,
+    GuardrailViolationEvent,
+    GuardrailViolationEventData,
     MemoryEvent,
     MemoryEventData,
     OutputParseEvent,
@@ -61,6 +63,15 @@ from .types import (
     ToolCallFunction,
     ToolCallStartEvent,
     ToolCallStartEventData,
+    Guardrail,
+    ValidValidationResult,
+    InvalidValidationResult,
+)
+from .guardrails import (
+    build_effective_guardrails,
+    execute_input_guardrails_sequential,
+    execute_input_guardrails_parallel,
+    execute_output_guardrails,
 )
 
 
@@ -399,36 +410,6 @@ async def _run_internal(
     if resumed:
         return resumed
     
-    # Check initial input guardrails on first turn
-    if state.turn_count == 0:
-        first_user_message = next((m for m in state.messages if m.role == ContentRole.USER or m.role == 'user'), None)
-        if first_user_message and config.initial_input_guardrails:
-            for guardrail in config.initial_input_guardrails:
-                if config.on_event:
-                    config.on_event(GuardrailEvent(data=GuardrailEventData(
-                        guardrail_name=getattr(guardrail, '__name__', 'unknown_guardrail'),
-                        content=get_text_content(first_user_message.content)
-                    )))
-                if asyncio.iscoroutinefunction(guardrail):
-                    result = await guardrail(get_text_content(first_user_message.content))
-                else:
-                    result = guardrail(get_text_content(first_user_message.content))
-
-                if not result.is_valid:
-                    if config.on_event:
-                        config.on_event(GuardrailEvent(data=GuardrailEventData(
-                            guardrail_name=getattr(guardrail, '__name__', 'unknown_guardrail'),
-                            content=get_text_content(first_user_message.content),
-                            is_valid=False,
-                            error_message=result.error_message
-                        )))
-                    return RunResult(
-                        final_state=state,
-                        outcome=ErrorOutcome(error=InputGuardrailTripwire(
-                            reason=result.error_message or "Input guardrail failed"
-                        ))
-                    )
-
     # Check max turns
     max_turns = config.max_turns or 50
     if state.turn_count >= max_turns:
@@ -445,6 +426,105 @@ async def _run_internal(
             outcome=ErrorOutcome(error=AgentNotFound(agent_name=state.current_agent_name))
         )
 
+    # Determine if agent has advanced guardrails configuration
+    has_advanced_guardrails = bool(
+        current_agent.advanced_config and 
+        current_agent.advanced_config.guardrails and
+        (current_agent.advanced_config.guardrails.input_prompt or 
+         current_agent.advanced_config.guardrails.output_prompt or 
+         current_agent.advanced_config.guardrails.require_citations)
+    )
+    
+    print('[JAF:ENGINE] Debug guardrails setup:', {
+        'agent_name': current_agent.name,
+        'has_advanced_config': bool(current_agent.advanced_config),
+        'has_advanced_guardrails': has_advanced_guardrails,
+        'initial_input_guardrails': len(config.initial_input_guardrails or []),
+        'final_output_guardrails': len(config.final_output_guardrails or [])
+    })
+
+    # Build effective guardrails
+    effective_input_guardrails: List[Guardrail] = []
+    effective_output_guardrails: List[Guardrail] = []
+    
+    if has_advanced_guardrails:
+        result = await build_effective_guardrails(current_agent, config)
+        effective_input_guardrails, effective_output_guardrails = result
+    else:
+        effective_input_guardrails = list(config.initial_input_guardrails or [])
+        effective_output_guardrails = list(config.final_output_guardrails or [])
+
+    # Execute input guardrails on first turn
+    input_guardrails_to_run = (effective_input_guardrails 
+                              if state.turn_count == 0 and effective_input_guardrails 
+                              else [])
+    
+    print('[JAF:ENGINE] Input guardrails to run:', {
+        'turn_count': state.turn_count,
+        'effective_input_length': len(effective_input_guardrails),
+        'input_guardrails_to_run_length': len(input_guardrails_to_run),
+        'has_advanced_guardrails': has_advanced_guardrails
+    })
+
+    if input_guardrails_to_run and state.turn_count == 0:
+        first_user_message = next((m for m in state.messages if m.role == ContentRole.USER or m.role == 'user'), None)
+        if first_user_message:
+            if has_advanced_guardrails:
+                execution_mode = (current_agent.advanced_config.guardrails.execution_mode 
+                                if current_agent.advanced_config and current_agent.advanced_config.guardrails 
+                                else 'parallel')
+            
+                if execution_mode == 'sequential':
+                    guardrail_result = await execute_input_guardrails_sequential(
+                        input_guardrails_to_run, first_user_message, config
+                    )
+                    if not guardrail_result.is_valid:
+                        return RunResult(
+                            final_state=state,
+                            outcome=ErrorOutcome(error=InputGuardrailTripwire(
+                                reason=getattr(guardrail_result, 'error_message', 'Input guardrail violation')
+                            ))
+                        )
+                else:
+                    # Parallel execution with LLM call overlap
+                    guardrail_result = await execute_input_guardrails_parallel(
+                        input_guardrails_to_run, first_user_message, config
+                    )
+                    if not guardrail_result.is_valid:
+                        print(f"🚨 Input guardrail violation: {getattr(guardrail_result, 'error_message', 'Unknown violation')}")
+                        return RunResult(
+                            final_state=state,
+                            outcome=ErrorOutcome(error=InputGuardrailTripwire(
+                                reason=getattr(guardrail_result, 'error_message', 'Input guardrail violation')
+                            ))
+                        )
+            else:
+                # Legacy guardrails path
+                print('[JAF:ENGINE] Using LEGACY guardrails path with', len(input_guardrails_to_run), 'guardrails')
+                for guardrail in input_guardrails_to_run:
+                    if config.on_event:
+                        config.on_event(GuardrailEvent(data=GuardrailEventData(
+                            guardrail_name=getattr(guardrail, '__name__', 'unknown_guardrail'),
+                            content=get_text_content(first_user_message.content)
+                        )))
+                    if asyncio.iscoroutinefunction(guardrail):
+                        result = await guardrail(get_text_content(first_user_message.content))
+                    else:
+                        result = guardrail(get_text_content(first_user_message.content))
+
+                    if not result.is_valid:
+                        if config.on_event:
+                            config.on_event(GuardrailViolationEvent(data=GuardrailViolationEventData(
+                                stage='input',
+                                reason=getattr(result, 'error_message', 'Input guardrail failed')
+                            )))
+                        return RunResult(
+                            final_state=state,
+                            outcome=ErrorOutcome(error=InputGuardrailTripwire(
+                                reason=getattr(result, 'error_message', 'Input guardrail failed')
+                            ))
+                        )
+
     # Agent debugging logs removed for performance
 
     # Get model name
@@ -516,12 +596,15 @@ async def _run_internal(
                     if len(partial_tool_calls) > 0:
                         message_tool_calls = []
                         for i, tc in enumerate(partial_tool_calls):
+                            arguments = tc["function"]["arguments"]
+                            if isinstance(arguments, str):
+                                arguments = _normalize_tool_call_arguments(arguments)
                             message_tool_calls.append({
                                 "id": tc["id"] or f"call_{i}",
                                 "type": "function",
                                 "function": {
                                     "name": tc["function"]["name"] or "",
-                                    "arguments": tc["function"]["arguments"]
+                                    "arguments": arguments
                                 }
                             })
 
@@ -534,7 +617,7 @@ async def _run_internal(
                                 type="function",
                                 function=ToolCallFunction(
                                     name=mc["function"]["name"],
-                                    arguments=mc["function"]["arguments"],
+                                    arguments=_normalize_tool_call_arguments(mc["function"]["arguments"])
                                 ),
                             ) for mc in message_tool_calls
                         ],
@@ -553,12 +636,15 @@ async def _run_internal(
             if len(partial_tool_calls) > 0:
                 final_tool_calls = []
                 for i, tc in enumerate(partial_tool_calls):
+                    arguments = tc["function"]["arguments"]
+                    if isinstance(arguments, str):
+                        arguments = _normalize_tool_call_arguments(arguments)
                     final_tool_calls.append({
                         "id": tc["id"] or f"call_{i}",
                         "type": "function",
                         "function": {
                             "name": tc["function"]["name"] or "",
-                            "arguments": tc["function"]["arguments"]
+                            "arguments": arguments
                         }
                     })
 
@@ -746,13 +832,27 @@ async def _run_internal(
                     )))
 
                 # Check final output guardrails
-                if config.final_output_guardrails:
-                    for guardrail in config.final_output_guardrails:
-                        if config.on_event:
-                            config.on_event(GuardrailEvent(data=GuardrailEventData(
-                                guardrail_name=getattr(guardrail, '__name__', 'unknown_guardrail'),
-                                content=output_data
-                            )))
+                if has_advanced_guardrails:
+                    # Use new advanced system
+                    output_guardrail_result = await execute_output_guardrails(
+                        effective_output_guardrails, output_data, config
+                    )
+                    if not output_guardrail_result.is_valid:
+                        return RunResult(
+                            final_state=replace(state, messages=new_messages),
+                            outcome=ErrorOutcome(error=OutputGuardrailTripwire(
+                                reason=getattr(output_guardrail_result, 'error_message', 'Output guardrail violation')
+                            ))
+                        )
+                else:
+                    # Legacy system
+                    if effective_output_guardrails:
+                        for guardrail in effective_output_guardrails:
+                            if config.on_event:
+                                config.on_event(GuardrailEvent(data=GuardrailEventData(
+                                    guardrail_name=getattr(guardrail, '__name__', 'unknown_guardrail'),
+                                    content=output_data
+                                )))
                         if asyncio.iscoroutinefunction(guardrail):
                             result = await guardrail(output_data)
                         else:
@@ -760,16 +860,14 @@ async def _run_internal(
 
                         if not result.is_valid:
                             if config.on_event:
-                                config.on_event(GuardrailEvent(data=GuardrailEventData(
-                                    guardrail_name=getattr(guardrail, '__name__', 'unknown_guardrail'),
-                                    content=output_data,
-                                    is_valid=False,
-                                    error_message=result.error_message
+                                config.on_event(GuardrailViolationEvent(data=GuardrailViolationEventData(
+                                    stage='output',
+                                    reason=getattr(result, 'error_message', 'Output guardrail failed')
                                 )))
                             return RunResult(
                                 final_state=replace(state, messages=new_messages, approvals=state.approvals),
                                 outcome=ErrorOutcome(error=OutputGuardrailTripwire(
-                                    reason=result.error_message or "Output guardrail failed"
+                                    reason=getattr(result, 'error_message', 'Output guardrail failed')
                                 ))
                             )
 
@@ -793,32 +891,44 @@ async def _run_internal(
                 )
         else:
             # No output codec, return content as string
-            if config.final_output_guardrails:
-                for guardrail in config.final_output_guardrails:
-                    if config.on_event:
-                        config.on_event(GuardrailEvent(data=GuardrailEventData(
-                            guardrail_name=getattr(guardrail, '__name__', 'unknown_guardrail'),
-                            content=get_text_content(assistant_message.content)
-                        )))
-                    if asyncio.iscoroutinefunction(guardrail):
-                        result = await guardrail(get_text_content(assistant_message.content))
-                    else:
-                        result = guardrail(get_text_content(assistant_message.content))
-
-                    if not result.is_valid:
+            if has_advanced_guardrails:
+                # Use new advanced system
+                output_guardrail_result = await execute_output_guardrails(
+                    effective_output_guardrails, get_text_content(assistant_message.content), config
+                )
+                if not output_guardrail_result.is_valid:
+                    return RunResult(
+                        final_state=replace(state, messages=new_messages),
+                        outcome=ErrorOutcome(error=OutputGuardrailTripwire(
+                            reason=getattr(output_guardrail_result, 'error_message', 'Output guardrail violation')
+                        ))
+                    )
+            else:
+                # Legacy system
+                if effective_output_guardrails:
+                    for guardrail in effective_output_guardrails:
                         if config.on_event:
                             config.on_event(GuardrailEvent(data=GuardrailEventData(
                                 guardrail_name=getattr(guardrail, '__name__', 'unknown_guardrail'),
-                                content=get_text_content(assistant_message.content),
-                                is_valid=False,
-                                error_message=result.error_message
+                                content=get_text_content(assistant_message.content)
                             )))
-                        return RunResult(
-                            final_state=replace(state, messages=new_messages, approvals=state.approvals),
-                            outcome=ErrorOutcome(error=OutputGuardrailTripwire(
-                                reason=result.error_message or "Output guardrail failed"
-                            ))
-                        )
+                        if asyncio.iscoroutinefunction(guardrail):
+                            result = await guardrail(get_text_content(assistant_message.content))
+                        else:
+                            result = guardrail(get_text_content(assistant_message.content))
+
+                        if not result.is_valid:
+                            if config.on_event:
+                                config.on_event(GuardrailViolationEvent(data=GuardrailViolationEventData(
+                                    stage='output',
+                                    reason=getattr(result, 'error_message', 'Output guardrail failed')
+                                )))
+                            return RunResult(
+                                final_state=replace(state, messages=new_messages, approvals=state.approvals),
+                                outcome=ErrorOutcome(error=OutputGuardrailTripwire(
+                                    reason=getattr(result, 'error_message', 'Output guardrail failed')
+                                ))
+                            )
 
             return RunResult(
                 final_state=replace(state, messages=new_messages, turn_count=state.turn_count + 1, approvals=state.approvals),
@@ -844,12 +954,33 @@ def _convert_tool_calls(tool_calls: Optional[List[Dict[str, Any]]]) -> Optional[
             type='function',
             function=ToolCallFunction(
                 name=tc['function']['name'],
-                arguments=tc['function']['arguments']
+                arguments=_normalize_tool_call_arguments(tc['function']['arguments'])
             )
         )
         for tc in tool_calls
     ]
 
+
+def _normalize_tool_call_arguments(arguments: Any) -> Any:
+    """Strip trailing streaming artifacts so arguments remain valid JSON strings."""
+    if not arguments or not isinstance(arguments, str):
+        return arguments
+
+    decoder = json.JSONDecoder()
+    try:
+        obj, end = decoder.raw_decode(arguments)
+    except json.JSONDecodeError:
+        return arguments
+
+    remainder = arguments[end:].strip()
+    if remainder:
+        try:
+            return json.dumps(obj)
+        except (TypeError, ValueError):
+            return arguments
+
+    return arguments
+
 async def _execute_tool_calls(
     tool_calls: List[ToolCall],
     agent: Agent[Ctx, Any],
@@ -865,7 +996,8 @@ async def _execute_tool_calls(
                 tool_name=tool_call.function.name,
                 args=_try_parse_json(tool_call.function.arguments),
                 trace_id=state.trace_id,
-                run_id=state.run_id
+                run_id=state.run_id,
+                call_id=tool_call.id
             ))))
 
         try:
@@ -891,7 +1023,8 @@ async def _execute_tool_calls(
                         trace_id=state.trace_id,
                         run_id=state.run_id,
                         status='error',
-                        tool_result={'error': 'tool_not_found'}
+                        tool_result={'error': 'tool_not_found'},
+                        call_id=tool_call.id
                     ))))
 
                 return {
@@ -925,7 +1058,8 @@ async def _execute_tool_calls(
                         trace_id=state.trace_id,
                         run_id=state.run_id,
                         status='error',
-                        tool_result={'error': 'validation_error', 'details': e.errors()}
+                        tool_result={'error': 'validation_error', 'details': e.errors()},
+                        call_id=tool_call.id
                     ))))
 
                 return {
@@ -1019,7 +1153,7 @@ async def _execute_tool_calls(
             else:
                 timeout = None
             if timeout is None:
-                timeout = config.default_tool_timeout if config.default_tool_timeout is not None else 30.0
+                timeout = config.default_tool_timeout if config.default_tool_timeout is not None else 300.0
 
             # Merge additional context if provided through approval
             additional_context = approval_status.additional_context if approval_status else None
@@ -1063,7 +1197,8 @@ async def _execute_tool_calls(
                         trace_id=state.trace_id,
                         run_id=state.run_id,
                         status='timeout',
-                        tool_result={'error': 'timeout_error'}
+                        tool_result={'error': 'timeout_error'},
+                        call_id=tool_call.id
                     ))))
 
                 return {
@@ -1115,7 +1250,8 @@ async def _execute_tool_calls(
                     trace_id=state.trace_id,
                     run_id=state.run_id,
                     tool_result=tool_result,
-                    status='success'
+                    status='success',
+                    call_id=tool_call.id
                 ))))
 
             # Check for handoff
@@ -1153,7 +1289,8 @@ async def _execute_tool_calls(
                     trace_id=state.trace_id,
                     run_id=state.run_id,
                     status='error',
-                    tool_result={'error': 'execution_error', 'detail': str(error)}
+                    tool_result={'error': 'execution_error', 'detail': str(error)},
+                    call_id=tool_call.id
                 ))))
 
             return {