praisonaiagents 0.0.97__tar.gz → 0.0.99__tar.gz

{praisonaiagents-0.0.97 → praisonaiagents-0.0.99}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: praisonaiagents
-Version: 0.0.97
+Version: 0.0.99
 Summary: Praison AI agents for completing complex tasks with Self Reflection Agents
 Author: Mervin Praison
 Requires-Python: >=3.10

{praisonaiagents-0.0.97 → praisonaiagents-0.0.99}/praisonaiagents/__init__.py

@@ -12,6 +12,7 @@ from .knowledge.knowledge import Knowledge
 from .knowledge.chunking import Chunking
 from .mcp.mcp import MCP
 from .session import Session
+from .guardrails import GuardrailResult, LLMGuardrail
 from .main import (
     TaskOutput,
     ReflectionOutput,
@@ -55,5 +56,7 @@ __all__ = [
     'async_display_callbacks',
     'Knowledge',
     'Chunking',
-    'MCP'
+    'MCP',
+    'GuardrailResult',
+    'LLMGuardrail'
 ] 

{praisonaiagents-0.0.97 → praisonaiagents-0.0.99}/praisonaiagents/agent/agent.py

@@ -3,7 +3,7 @@ import time
 import json
 import logging
 import asyncio
-from typing import List, Optional, Any, Dict, Union, Literal, TYPE_CHECKING
+from typing import List, Optional, Any, Dict, Union, Literal, TYPE_CHECKING, Callable, Tuple
 from rich.console import Console
 from rich.live import Live
 from openai import AsyncOpenAI
@@ -32,6 +32,7 @@ _shared_apps = {}  # Dict of port -> FastAPI app
 
 if TYPE_CHECKING:
     from ..task.task import Task
+    from ..main import TaskOutput
 
 @dataclass
 class ChatCompletionMessage:
@@ -368,7 +369,9 @@ class Agent:
         min_reflect: int = 1,
         reflect_llm: Optional[str] = None,
         user_id: Optional[str] = None,
-        reasoning_steps: bool = False
+        reasoning_steps: bool = False,
+        guardrail: Optional[Union[Callable[['TaskOutput'], Tuple[bool, Any]], str]] = None,
+        max_guardrail_retries: int = 3
     ):
         # Add check at start if memory is requested
         if memory is not None:
@@ -483,6 +486,12 @@ Your Goal: {self.goal}
         # Store user_id
         self.user_id = user_id or "praison"
         self.reasoning_steps = reasoning_steps
+        
+        # Initialize guardrail settings
+        self.guardrail = guardrail
+        self.max_guardrail_retries = max_guardrail_retries
+        self._guardrail_fn = None
+        self._setup_guardrail()
 
         # Check if knowledge parameter has any values
         if not knowledge:
@@ -512,6 +521,152 @@ Your Goal: {self.goal}
         except Exception as e:
             logging.error(f"Error processing knowledge item: {knowledge_item}, error: {e}")
 
+    def _setup_guardrail(self):
+        """Setup the guardrail function based on the provided guardrail parameter."""
+        if self.guardrail is None:
+            self._guardrail_fn = None
+            return
+            
+        if callable(self.guardrail):
+            # Validate function signature
+            sig = inspect.signature(self.guardrail)
+            positional_args = [
+                param for param in sig.parameters.values()
+                if param.default is inspect.Parameter.empty
+            ]
+            if len(positional_args) != 1:
+                raise ValueError("Agent guardrail function must accept exactly one parameter (TaskOutput)")
+            
+            # Check return annotation if present
+            from typing import get_args, get_origin
+            return_annotation = sig.return_annotation
+            if return_annotation != inspect.Signature.empty:
+                return_annotation_args = get_args(return_annotation)
+                if not (
+                    get_origin(return_annotation) is tuple
+                    and len(return_annotation_args) == 2
+                    and return_annotation_args[0] is bool
+                    and (
+                        return_annotation_args[1] is Any
+                        or return_annotation_args[1] is str
+                        or str(return_annotation_args[1]).endswith('TaskOutput')
+                        or str(return_annotation_args[1]).startswith('typing.Union')
+                    )
+                ):
+                    raise ValueError(
+                        "If return type is annotated, it must be Tuple[bool, Any] or Tuple[bool, Union[str, TaskOutput]]"
+                    )
+            
+            self._guardrail_fn = self.guardrail
+        elif isinstance(self.guardrail, str):
+            # Create LLM-based guardrail
+            from ..guardrails import LLMGuardrail
+            llm = getattr(self, 'llm', None) or getattr(self, 'llm_instance', None)
+            self._guardrail_fn = LLMGuardrail(description=self.guardrail, llm=llm)
+        else:
+            raise ValueError("Agent guardrail must be either a callable or a string description")
+
+    def _process_guardrail(self, task_output):
+        """Process the guardrail validation for a task output.
+        
+        Args:
+            task_output: The task output to validate
+            
+        Returns:
+            GuardrailResult: The result of the guardrail validation
+        """
+        from ..guardrails import GuardrailResult
+        
+        if not self._guardrail_fn:
+            return GuardrailResult(success=True, result=task_output)
+        
+        try:
+            # Call the guardrail function
+            result = self._guardrail_fn(task_output)
+            
+            # Convert the result to a GuardrailResult
+            return GuardrailResult.from_tuple(result)
+            
+        except Exception as e:
+            logging.error(f"Agent {self.name}: Error in guardrail validation: {e}")
+            # On error, return failure
+            return GuardrailResult(
+                success=False,
+                result=None,
+                error=f"Agent guardrail validation error: {str(e)}"
+            )
+
+    def _apply_guardrail_with_retry(self, response_text, prompt, temperature=0.2, tools=None):
+        """Apply guardrail validation with retry logic.
+        
+        Args:
+            response_text: The response to validate
+            prompt: Original prompt for regeneration if needed
+            temperature: Temperature for regeneration
+            tools: Tools for regeneration
+            
+        Returns:
+            str: The validated response text or None if validation fails after retries
+        """
+        if not self._guardrail_fn:
+            return response_text
+            
+        from ..main import TaskOutput
+        
+        retry_count = 0
+        current_response = response_text
+        
+        while retry_count <= self.max_guardrail_retries:
+            # Create TaskOutput object
+            task_output = TaskOutput(
+                raw=current_response,
+                output=current_response,
+                pydantic=None,
+                json_dict=None,
+                name=f"{self.name}_output",
+                description="Agent response output"
+            )
+            
+            # Process guardrail
+            guardrail_result = self._process_guardrail(task_output)
+            
+            if guardrail_result.success:
+                logging.info(f"Agent {self.name}: Guardrail validation passed")
+                # Return the potentially modified result
+                if guardrail_result.result and hasattr(guardrail_result.result, 'raw'):
+                    return guardrail_result.result.raw
+                elif guardrail_result.result:
+                    return str(guardrail_result.result)
+                else:
+                    return current_response
+            
+            # Guardrail failed
+            if retry_count >= self.max_guardrail_retries:
+                raise Exception(
+                    f"Agent {self.name} response failed guardrail validation after {self.max_guardrail_retries} retries. "
+                    f"Last error: {guardrail_result.error}"
+                )
+            
+            retry_count += 1
+            logging.warning(f"Agent {self.name}: Guardrail validation failed (retry {retry_count}/{self.max_guardrail_retries}): {guardrail_result.error}")
+            
+            # Regenerate response for retry
+            try:
+                retry_prompt = f"{prompt}\n\nNote: Previous response failed validation due to: {guardrail_result.error}. Please provide an improved response."
+                response = self._chat_completion([{"role": "user", "content": retry_prompt}], temperature, tools)
+                if response and response.choices:
+                    current_response = response.choices[0].message.content.strip()
+                else:
+                    raise Exception("Failed to generate retry response")
+            except Exception as e:
+                logging.error(f"Agent {self.name}: Error during guardrail retry: {e}")
+                # If we can't regenerate, fail the guardrail
+                raise Exception(
+                    f"Agent {self.name} guardrail retry failed: {e}"
+                )
+        
+        return current_response
+
     def generate_task(self) -> 'Task':
         """Generate a Task object from the agent's instructions"""
         from ..task.task import Task
@@ -788,69 +943,92 @@ Your Goal: {self.goal}
                     )
             else:
                 # Use the standard OpenAI client approach
-                if stream:
-                    # Process as streaming response with formatted tools
-                    final_response = self._process_stream_response(
-                        messages, 
-                        temperature, 
-                        start_time, 
-                        formatted_tools=formatted_tools if formatted_tools else None,
-                        reasoning_steps=reasoning_steps
-                    )
-                else:
-                    # Process as regular non-streaming response
-                    final_response = client.chat.completions.create(
-                        model=self.llm,
-                        messages=messages,
-                        temperature=temperature,
-                        tools=formatted_tools if formatted_tools else None,
-                        stream=False
-                    )
-
-            tool_calls = getattr(final_response.choices[0].message, 'tool_calls', None)
-
-            if tool_calls:
-                messages.append({
-                    "role": "assistant", 
-                    "content": final_response.choices[0].message.content,
-                    "tool_calls": tool_calls
-                })
+                # Continue tool execution loop until no more tool calls are needed
+                max_iterations = 10  # Prevent infinite loops
+                iteration_count = 0
+                
+                while iteration_count < max_iterations:
+                    if stream:
+                        # Process as streaming response with formatted tools
+                        final_response = self._process_stream_response(
+                            messages, 
+                            temperature, 
+                            start_time, 
+                            formatted_tools=formatted_tools if formatted_tools else None,
+                            reasoning_steps=reasoning_steps
+                        )
+                    else:
+                        # Process as regular non-streaming response
+                        final_response = client.chat.completions.create(
+                            model=self.llm,
+                            messages=messages,
+                            temperature=temperature,
+                            tools=formatted_tools if formatted_tools else None,
+                            stream=False
+                        )
 
-                for tool_call in tool_calls:
-                    function_name = tool_call.function.name
-                    arguments = json.loads(tool_call.function.arguments)
+                    tool_calls = getattr(final_response.choices[0].message, 'tool_calls', None)
 
-                    if self.verbose:
-                        display_tool_call(f"Agent {self.name} is calling function '{function_name}' with arguments: {arguments}")
+                    if tool_calls:
+                        messages.append({
+                            "role": "assistant", 
+                            "content": final_response.choices[0].message.content,
+                            "tool_calls": tool_calls
+                        })
 
-                    tool_result = self.execute_tool(function_name, arguments)
-                    results_str = json.dumps(tool_result) if tool_result else "Function returned an empty output"
+                        for tool_call in tool_calls:
+                            function_name = tool_call.function.name
+                            arguments = json.loads(tool_call.function.arguments)
 
-                    if self.verbose:
-                        display_tool_call(f"Function '{function_name}' returned: {results_str}")
+                            if self.verbose:
+                                display_tool_call(f"Agent {self.name} is calling function '{function_name}' with arguments: {arguments}")
 
-                    messages.append({
-                        "role": "tool",
-                        "tool_call_id": tool_call.id,
-                        "content": results_str
-                    })
+                            tool_result = self.execute_tool(function_name, arguments)
+                            results_str = json.dumps(tool_result) if tool_result else "Function returned an empty output"
 
-                # Get final response after tool calls
-                if stream:
-                    final_response = self._process_stream_response(
-                        messages, 
-                        temperature, 
-                        start_time,
-                        formatted_tools=formatted_tools if formatted_tools else None,
-                        reasoning_steps=reasoning_steps
-                    )
-                else:
-                    final_response = client.chat.completions.create(
-                        model=self.llm,
-                        messages=messages,
-                        temperature=temperature,
-                        stream=False
-                    )
+                            if self.verbose:
+                                display_tool_call(f"Function '{function_name}' returned: {results_str}")
+
+                            messages.append({
+                                "role": "tool",
+                                "tool_call_id": tool_call.id,
+                                "content": results_str
+                            })
+
+                        # Check if we should continue (for tools like sequential thinking)
+                        should_continue = False
+                        for tool_call in tool_calls:
+                            function_name = tool_call.function.name
+                            arguments = json.loads(tool_call.function.arguments)
+                            
+                            # For sequential thinking tool, check if nextThoughtNeeded is True
+                            if function_name == "sequentialthinking" and arguments.get("nextThoughtNeeded", False):
+                                should_continue = True
+                                break
+
+                        if not should_continue:
+                            # Get final response after tool calls
+                            if stream:
+                                final_response = self._process_stream_response(
+                                    messages, 
+                                    temperature, 
+                                    start_time,
+                                    formatted_tools=formatted_tools if formatted_tools else None,
+                                    reasoning_steps=reasoning_steps
+                                )
+                            else:
+                                final_response = client.chat.completions.create(
+                                    model=self.llm,
+                                    messages=messages,
+                                    temperature=temperature,
+                                    stream=False
+                                )
+                            break
+                        
+                        iteration_count += 1
+                    else:
+                        # No tool calls, we're done
+                        break
 
             return final_response
 
@@ -944,7 +1122,13 @@ Your Goal: {self.goal}
                     total_time = time.time() - start_time
                     logging.debug(f"Agent.chat completed in {total_time:.2f} seconds")
 
-                return response_text
+                # Apply guardrail validation for custom LLM response
+                try:
+                    validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools)
+                    return validated_response
+                except Exception as e:
+                    logging.error(f"Agent {self.name}: Guardrail validation failed for custom LLM: {e}")
+                    return None
             except Exception as e:
                 display_error(f"Error in LLM chat: {e}")
                 return None
@@ -1032,8 +1216,20 @@ Your Goal: {self.goal}
                         display_interaction(original_prompt, response_text, markdown=self.markdown, generation_time=time.time() - start_time, console=self.console)
                         # Return only reasoning content if reasoning_steps is True
                         if reasoning_steps and hasattr(response.choices[0].message, 'reasoning_content'):
-                            return response.choices[0].message.reasoning_content
-                        return response_text
+                            # Apply guardrail to reasoning content
+                            try:
+                                validated_reasoning = self._apply_guardrail_with_retry(response.choices[0].message.reasoning_content, original_prompt, temperature, tools)
+                                return validated_reasoning
+                            except Exception as e:
+                                logging.error(f"Agent {self.name}: Guardrail validation failed for reasoning content: {e}")
+                                return None
+                        # Apply guardrail to regular response
+                        try:
+                            validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools)
+                            return validated_response
+                        except Exception as e:
+                            logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
+                            return None
 
                     reflection_prompt = f"""
 Reflect on your previous response: '{response_text}'.
@@ -1066,7 +1262,13 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
                             self.chat_history.append({"role": "user", "content": prompt})
                             self.chat_history.append({"role": "assistant", "content": response_text})
                             display_interaction(prompt, response_text, markdown=self.markdown, generation_time=time.time() - start_time, console=self.console)
-                            return response_text
+                            # Apply guardrail validation after satisfactory reflection
+                            try:
+                                validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools)
+                                return validated_response
+                            except Exception as e:
+                                logging.error(f"Agent {self.name}: Guardrail validation failed after reflection: {e}")
+                                return None
 
                         # Check if we've hit max reflections
                         if reflection_count >= self.max_reflect - 1:
@@ -1075,7 +1277,13 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
                             self.chat_history.append({"role": "user", "content": prompt})
                             self.chat_history.append({"role": "assistant", "content": response_text})
                             display_interaction(prompt, response_text, markdown=self.markdown, generation_time=time.time() - start_time, console=self.console)
-                            return response_text
+                            # Apply guardrail validation after max reflections
+                            try:
+                                validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools)
+                                return validated_response
+                            except Exception as e:
+                                logging.error(f"Agent {self.name}: Guardrail validation failed after max reflections: {e}")
+                                return None
 
                         logging.debug(f"{self.name} reflection count {reflection_count + 1}, continuing reflection process")
                         messages.append({"role": "user", "content": "Now regenerate your response using the reflection you made"})
@@ -1099,8 +1307,16 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
         if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
             total_time = time.time() - start_time
             logging.debug(f"Agent.chat completed in {total_time:.2f} seconds")
-            
-        return response_text
+        
+        # Apply guardrail validation before returning    
+        try:
+            validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools)
+            return validated_response
+        except Exception as e:
+            logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
+            if self.verbose:
+                display_error(f"Guardrail validation failed: {e}", console=self.console)
+            return None
 
     def clean_json_output(self, output: str) -> str:
         """Clean and extract JSON from response text."""
@@ -1670,6 +1886,7 @@ Your Goal: {self.goal}
                 import threading
                 import time
                 import inspect
+                import asyncio  # Import asyncio in the MCP scope
                 # logging is already imported at the module level
                 
             except ImportError as e: