praisonaiagents 0.0.29__py3-none-any.whl → 0.0.53__py3-none-any.whl

praisonaiagents/agent/agent.py

@@ -16,13 +16,142 @@ from ..main import (
     display_self_reflection,
     ReflectionOutput,
     client,
-    error_logs
+    error_logs,
+    adisplay_instruction
 )
 import inspect
+import uuid
+from dataclasses import dataclass
 
 if TYPE_CHECKING:
     from ..task.task import Task
 
+@dataclass
+class ChatCompletionMessage:
+    content: str
+    role: str = "assistant"
+    refusal: Optional[str] = None
+    audio: Optional[str] = None
+    function_call: Optional[dict] = None
+    tool_calls: Optional[List] = None
+    reasoning_content: Optional[str] = None
+
+@dataclass
+class Choice:
+    finish_reason: Optional[str]
+    index: int
+    message: ChatCompletionMessage
+    logprobs: Optional[dict] = None
+
+@dataclass
+class CompletionTokensDetails:
+    accepted_prediction_tokens: Optional[int] = None
+    audio_tokens: Optional[int] = None
+    reasoning_tokens: Optional[int] = None
+    rejected_prediction_tokens: Optional[int] = None
+
+@dataclass
+class PromptTokensDetails:
+    audio_tokens: Optional[int] = None
+    cached_tokens: int = 0
+
+@dataclass
+class CompletionUsage:
+    completion_tokens: int = 0
+    prompt_tokens: int = 0
+    total_tokens: int = 0
+    completion_tokens_details: Optional[CompletionTokensDetails] = None
+    prompt_tokens_details: Optional[PromptTokensDetails] = None
+    prompt_cache_hit_tokens: int = 0
+    prompt_cache_miss_tokens: int = 0
+
+@dataclass
+class ChatCompletion:
+    id: str
+    choices: List[Choice]
+    created: int
+    model: str
+    object: str = "chat.completion"
+    system_fingerprint: Optional[str] = None
+    service_tier: Optional[str] = None
+    usage: Optional[CompletionUsage] = None
+
+def process_stream_chunks(chunks):
+    """Process streaming chunks into combined response"""
+    if not chunks:
+        return None
+    
+    try:
+        first_chunk = chunks[0]
+        last_chunk = chunks[-1]
+        
+        # Basic metadata
+        id = getattr(first_chunk, "id", None) 
+        created = getattr(first_chunk, "created", None)
+        model = getattr(first_chunk, "model", None)
+        system_fingerprint = getattr(first_chunk, "system_fingerprint", None)
+        
+        # Track usage
+        completion_tokens = 0
+        prompt_tokens = 0
+        
+        content_list = []
+        reasoning_list = []
+
+        for chunk in chunks:
+            if not hasattr(chunk, "choices") or not chunk.choices:
+                continue
+            
+            # Track usage from each chunk
+            if hasattr(chunk, "usage"):
+                completion_tokens += getattr(chunk.usage, "completion_tokens", 0)
+                prompt_tokens += getattr(chunk.usage, "prompt_tokens", 0)
+                
+            delta = getattr(chunk.choices[0], "delta", None)
+            if not delta:
+                continue
+                
+            if hasattr(delta, "content") and delta.content:
+                content_list.append(delta.content)
+            if hasattr(delta, "reasoning_content") and delta.reasoning_content:
+                reasoning_list.append(delta.reasoning_content)
+
+        combined_content = "".join(content_list) if content_list else ""
+        combined_reasoning = "".join(reasoning_list) if reasoning_list else None
+        finish_reason = getattr(last_chunk.choices[0], "finish_reason", None) if hasattr(last_chunk, "choices") and last_chunk.choices else None
+
+        message = ChatCompletionMessage(
+            content=combined_content,
+            reasoning_content=combined_reasoning
+        )
+        
+        choice = Choice(
+            finish_reason=finish_reason,
+            index=0,
+            message=message
+        )
+
+        usage = CompletionUsage(
+            completion_tokens=completion_tokens,
+            prompt_tokens=prompt_tokens,
+            total_tokens=completion_tokens + prompt_tokens,
+            completion_tokens_details=CompletionTokensDetails(),
+            prompt_tokens_details=PromptTokensDetails()
+        )
+        
+        return ChatCompletion(
+            id=id,
+            choices=[choice],
+            created=created,
+            model=model,
+            system_fingerprint=system_fingerprint,
+            usage=usage
+        )
+        
+    except Exception as e:
+        print(f"Error processing chunks: {e}")
+        return None
+
 class Agent:
     def _generate_tool_definition(self, function_name):
         """
@@ -69,10 +198,15 @@ class Agent:
 
         import inspect
         # Langchain tools
-        if inspect.isclass(func) and hasattr(func, 'run'):
+        if inspect.isclass(func) and hasattr(func, 'run') and not hasattr(func, '_run'):
             original_func = func
             func = func.run
             function_name = original_func.__name__
+        # CrewAI tools
+        elif inspect.isclass(func) and hasattr(func, '_run'):
+            original_func = func
+            func = func._run
+            function_name = original_func.__name__
 
         sig = inspect.signature(func)
         logging.debug(f"Function signature: {sig}")
@@ -175,18 +309,37 @@ class Agent:
         respect_context_window: bool = True,
         code_execution_mode: Literal["safe", "unsafe"] = "safe",
         embedder_config: Optional[Dict[str, Any]] = None,
-        knowledge_sources: Optional[List[Any]] = None,
+        knowledge: Optional[List[str]] = None,
+        knowledge_config: Optional[Dict[str, Any]] = None,
         use_system_prompt: Optional[bool] = True,
         markdown: bool = True,
         self_reflect: bool = False,
         max_reflect: int = 3,
         min_reflect: int = 1,
-        reflect_llm: Optional[str] = None
+        reflect_llm: Optional[str] = None,
+        user_id: Optional[str] = None,
+        reasoning_steps: bool = False
     ):
+        # Add check at start if memory is requested
+        if memory is not None:
+            try:
+                from ..memory.memory import Memory
+                MEMORY_AVAILABLE = True
+            except ImportError:
+                raise ImportError(
+                    "Memory features requested in Agent but memory dependencies not installed. "
+                    "Please install with: pip install \"praisonaiagents[memory]\""
+                )
+
         # Handle backward compatibility for required fields
         if all(x is None for x in [name, role, goal, backstory, instructions]):
             raise ValueError("At least one of name, role, goal, backstory, or instructions must be provided")
 
+        # Configure logging to suppress unwanted outputs
+        logging.getLogger("litellm").setLevel(logging.WARNING)
+        logging.getLogger("httpx").setLevel(logging.WARNING)
+        logging.getLogger("httpcore").setLevel(logging.WARNING)
+
         # If instructions are provided, use them to set role, goal, and backstory
         if instructions:
             self.name = name or "Agent"
@@ -206,7 +359,34 @@ class Agent:
         
         self.instructions = instructions
         # Check for model name in environment variable if not provided
-        self.llm = llm or os.getenv('OPENAI_MODEL_NAME', 'gpt-4o')
+        self._using_custom_llm = False
+
+        # If the user passes a dictionary (for advanced configuration)
+        if isinstance(llm, dict) and "model" in llm:
+            try:
+                from ..llm.llm import LLM
+                self.llm_instance = LLM(**llm)  # Pass all dict items as kwargs
+                self._using_custom_llm = True
+            except ImportError as e:
+                raise ImportError(
+                    "LLM features requested but dependencies not installed. "
+                    "Please install with: pip install \"praisonaiagents[llm]\""
+                ) from e
+        # If the user passes a string with a slash (provider/model)
+        elif isinstance(llm, str) and "/" in llm:
+            try:
+                from ..llm.llm import LLM
+                # Pass the entire string so LiteLLM can parse provider/model
+                self.llm_instance = LLM(model=llm)
+                self._using_custom_llm = True
+            except ImportError as e:
+                raise ImportError(
+                    "LLM features requested but dependencies not installed. "
+                    "Please install with: pip install \"praisonaiagents[llm]\""
+                ) from e
+        # Otherwise, fall back to OpenAI environment/name
+        else:
+            self.llm = llm or os.getenv('OPENAI_MODEL_NAME', 'gpt-4o')
         self.tools = tools if tools else []  # Store original tools
         self.function_calling_llm = function_calling_llm
         self.max_iter = max_iter
@@ -225,7 +405,7 @@ class Agent:
         self.respect_context_window = respect_context_window
         self.code_execution_mode = code_execution_mode
         self.embedder_config = embedder_config
-        self.knowledge_sources = knowledge_sources
+        self.knowledge = knowledge
         self.use_system_prompt = use_system_prompt
         self.chat_history = []
         self.markdown = markdown
@@ -241,6 +421,41 @@ Your Role: {self.role}\n
 Your Goal: {self.goal}
         """
 
+        # Generate unique IDs
+        self.agent_id = str(uuid.uuid4())
+
+        # Store user_id
+        self.user_id = user_id or "praison"
+        self.reasoning_steps = reasoning_steps
+
+        # Check if knowledge parameter has any values
+        if not knowledge:
+            self.knowledge = None
+        else:
+            # Initialize Knowledge with provided or default config
+            from praisonaiagents.knowledge import Knowledge
+            self.knowledge = Knowledge(knowledge_config or None)
+            
+            # Handle knowledge
+            if knowledge:
+                for source in knowledge:
+                    self._process_knowledge(source)
+
+    def _process_knowledge(self, knowledge_item):
+        """Process and store knowledge from a file path, URL, or string."""
+        try:
+            if os.path.exists(knowledge_item):
+                # It's a file path
+                self.knowledge.add(knowledge_item, user_id=self.user_id, agent_id=self.agent_id)
+            elif knowledge_item.startswith("http://") or knowledge_item.startswith("https://"):
+                # It's a URL
+                pass
+            else:
+                # It's a string content
+                self.knowledge.store(knowledge_item, user_id=self.user_id, agent_id=self.agent_id)
+        except Exception as e:
+            logging.error(f"Error processing knowledge item: {knowledge_item}, error: {e}")
+
     def generate_task(self) -> 'Task':
         """Generate a Task object from the agent's instructions"""
         from ..task.task import Task
@@ -279,14 +494,22 @@ Your Goal: {self.goal}
 
         if func:
             try:
-                # If it's a class with run method, instantiate and call run
-                if inspect.isclass(func) and hasattr(func, 'run'):
+                # Langchain: If it's a class with run but not _run, instantiate and call run
+                if inspect.isclass(func) and hasattr(func, 'run') and not hasattr(func, '_run'):
                     instance = func()
-                    # Extract only the parameters that run() expects
                     run_params = {k: v for k, v in arguments.items() 
-                                if k in inspect.signature(instance.run).parameters 
-                                and k != 'self'}
+                                  if k in inspect.signature(instance.run).parameters 
+                                  and k != 'self'}
                     return instance.run(**run_params)
+
+                # CrewAI: If it's a class with an _run method, instantiate and call _run
+                elif inspect.isclass(func) and hasattr(func, '_run'):
+                    instance = func()
+                    run_params = {k: v for k, v in arguments.items() 
+                                  if k in inspect.signature(instance._run).parameters 
+                                  and k != 'self'}
+                    return instance._run(**run_params)
+
                 # Otherwise treat as regular function
                 elif callable(func):
                     return func(**arguments)
@@ -305,7 +528,7 @@ Your Goal: {self.goal}
     def __str__(self):
         return f"Agent(name='{self.name}', role='{self.role}', goal='{self.goal}')"
 
-    def _chat_completion(self, messages, temperature=0.2, tools=None, stream=True):
+    def _chat_completion(self, messages, temperature=0.2, tools=None, stream=True, reasoning_steps=False):
         start_time = time.time()
         logging.debug(f"{self.name} sending messages to LLM: {messages}")
 
@@ -375,30 +598,35 @@ Your Goal: {self.goal}
                     stream=True
                 )
                 full_response_text = ""
+                reasoning_content = ""
+                chunks = []
                 
                 # Create Live display with proper configuration
                 with Live(
                     display_generating("", start_time),
                     console=self.console,
                     refresh_per_second=4,
-                    transient=False,  # Changed to False to preserve output
+                    transient=True,
                     vertical_overflow="ellipsis",
                     auto_refresh=True
                 ) as live:
                     for chunk in response_stream:
+                        chunks.append(chunk)
                         if chunk.choices[0].delta.content:
                             full_response_text += chunk.choices[0].delta.content
                             live.update(display_generating(full_response_text, start_time))
+                        
+                        # Update live display with reasoning content if enabled
+                        if reasoning_steps and hasattr(chunk.choices[0].delta, "reasoning_content"):
+                            rc = chunk.choices[0].delta.reasoning_content
+                            if rc:
+                                reasoning_content += rc
+                                live.update(display_generating(f"{full_response_text}\n[Reasoning: {reasoning_content}]", start_time))
                 
                 # Clear the last generating display with a blank line
                 self.console.print()
                 
-                final_response = client.chat.completions.create(
-                    model=self.llm,
-                    messages=messages,
-                    temperature=temperature,
-                    stream=False
-                )
+                final_response = process_stream_chunks(chunks)
                 return final_response
             else:
                 if tool_calls:
@@ -416,178 +644,237 @@ Your Goal: {self.goal}
             display_error(f"Error in chat completion: {e}")
             return None
 
-    def chat(self, prompt, temperature=0.2, tools=None, output_json=None, output_pydantic=None):
-        if self.use_system_prompt:
-            system_prompt = f"""{self.backstory}\n
-Your Role: {self.role}\n
-Your Goal: {self.goal}
-            """
-            if output_json:
-                system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {output_json.schema_json()}"
-            elif output_pydantic:
-                system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {output_pydantic.schema_json()}"
-        else:
-            system_prompt = None
-
-        messages = []
-        if system_prompt:
-            messages.append({"role": "system", "content": system_prompt})
-        messages.extend(self.chat_history)
-
-        # Modify prompt if output_json or output_pydantic is specified
-        original_prompt = prompt
-        if output_json or output_pydantic:
-            if isinstance(prompt, str):
-                prompt += "\nReturn ONLY a valid JSON object. No other text or explanation."
-            elif isinstance(prompt, list):
-                # For multimodal prompts, append to the text content
-                for item in prompt:
-                    if item["type"] == "text":
-                        item["text"] += "\nReturn ONLY a valid JSON object. No other text or explanation."
-                        break
-
-        if isinstance(prompt, list):
-            # If we receive a multimodal prompt list, place it directly in the user message
-            messages.append({"role": "user", "content": prompt})
-        else:
-            messages.append({"role": "user", "content": prompt})
-
-        final_response_text = None
-        reflection_count = 0
-        start_time = time.time()
+    def chat(self, prompt, temperature=0.2, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False):
+        reasoning_steps = reasoning_steps or self.reasoning_steps
+        # Search for existing knowledge if any knowledge is provided
+        if self.knowledge:
+            search_results = self.knowledge.search(prompt, agent_id=self.agent_id)
+            if search_results:
+                # Check if search_results is a list of dictionaries or strings
+                if isinstance(search_results, dict) and 'results' in search_results:
+                    # Extract memory content from the results
+                    knowledge_content = "\n".join([result['memory'] for result in search_results['results']])
+                else:
+                    # If search_results is a list of strings, join them directly
+                    knowledge_content = "\n".join(search_results)
+                
+                # Append found knowledge to the prompt
+                prompt = f"{prompt}\n\nKnowledge: {knowledge_content}"
 
-        while True:
+        if self._using_custom_llm:
             try:
-                if self.verbose:
-                    # Handle both string and list prompts for instruction display
-                    display_text = prompt
-                    if isinstance(prompt, list):
-                        # Extract text content from multimodal prompt
-                        display_text = next((item["text"] for item in prompt if item["type"] == "text"), "")
-                    
-                    if display_text and str(display_text).strip():
-                        display_instruction(f"Agent {self.name} is processing prompt: {display_text}", console=self.console)
+                # Pass everything to LLM class
+                response_text = self.llm_instance.get_response(
+                    prompt=prompt,
+                    system_prompt=f"{self.backstory}\n\nYour Role: {self.role}\n\nYour Goal: {self.goal}" if self.use_system_prompt else None,
+                    chat_history=self.chat_history,
+                    temperature=temperature,
+                    tools=tools,
+                    output_json=output_json,
+                    output_pydantic=output_pydantic,
+                    verbose=self.verbose,
+                    markdown=self.markdown,
+                    self_reflect=self.self_reflect,
+                    max_reflect=self.max_reflect,
+                    min_reflect=self.min_reflect,
+                    console=self.console,
+                    agent_name=self.name,
+                    agent_role=self.role,
+                    agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in self.tools],
+                    execute_tool_fn=self.execute_tool,  # Pass tool execution function
+                    reasoning_steps=reasoning_steps
+                )
 
-                response = self._chat_completion(messages, temperature=temperature, tools=tools if tools else None)
-                if not response:
-                    return None
+                self.chat_history.append({"role": "user", "content": prompt})
+                self.chat_history.append({"role": "assistant", "content": response_text})
 
-                tool_calls = getattr(response.choices[0].message, 'tool_calls', None)
-                response_text = response.choices[0].message.content.strip()
+                return response_text
+            except Exception as e:
+                display_error(f"Error in LLM chat: {e}")
+                return None
+        else:
+            if self.use_system_prompt:
+                system_prompt = f"""{self.backstory}\n
+Your Role: {self.role}\n
+Your Goal: {self.goal}
+                """
+                if output_json:
+                    system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(output_json.model_json_schema())}"
+                elif output_pydantic:
+                    system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(output_pydantic.model_json_schema())}"
+            else:
+                system_prompt = None
 
-                if tool_calls:
-                    messages.append({
-                        "role": "assistant",
-                        "content": response_text,
-                        "tool_calls": tool_calls
-                    })
-                    
-                    for tool_call in tool_calls:
-                        function_name = tool_call.function.name
-                        arguments = json.loads(tool_call.function.arguments)
+            messages = []
+            if system_prompt:
+                messages.append({"role": "system", "content": system_prompt})
+            messages.extend(self.chat_history)
 
-                        if self.verbose:
-                            display_tool_call(f"Agent {self.name} is calling function '{function_name}' with arguments: {arguments}", console=self.console)
+            # Modify prompt if output_json or output_pydantic is specified
+            original_prompt = prompt
+            if output_json or output_pydantic:
+                if isinstance(prompt, str):
+                    prompt += "\nReturn ONLY a valid JSON object. No other text or explanation."
+                elif isinstance(prompt, list):
+                    # For multimodal prompts, append to the text content
+                    for item in prompt:
+                        if item["type"] == "text":
+                            item["text"] += "\nReturn ONLY a valid JSON object. No other text or explanation."
+                            break
 
-                        tool_result = self.execute_tool(function_name, arguments)
+            if isinstance(prompt, list):
+                # If we receive a multimodal prompt list, place it directly in the user message
+                messages.append({"role": "user", "content": prompt})
+            else:
+                messages.append({"role": "user", "content": prompt})
 
-                        if tool_result:
-                            if self.verbose:
-                                display_tool_call(f"Function '{function_name}' returned: {tool_result}", console=self.console)
-                            messages.append({
-                                "role": "tool",
-                                "tool_call_id": tool_call.id,
-                                "content": json.dumps(tool_result)
-                            })
-                        else:
-                            messages.append({
-                                "role": "tool",
-                                "tool_call_id": tool_call.id,
-                                "content": "Function returned an empty output"
-                            })
+            final_response_text = None
+            reflection_count = 0
+            start_time = time.time()
+
+            while True:
+                try:
+                    if self.verbose:
+                        # Handle both string and list prompts for instruction display
+                        display_text = prompt
+                        if isinstance(prompt, list):
+                            # Extract text content from multimodal prompt
+                            display_text = next((item["text"] for item in prompt if item["type"] == "text"), "")
                         
-                    response = self._chat_completion(messages, temperature=temperature)
+                        if display_text and str(display_text).strip():
+                            # Pass agent information to display_instruction
+                            agent_tools = [t.__name__ if hasattr(t, '__name__') else str(t) for t in self.tools]
+                            display_instruction(
+                                f"Agent {self.name} is processing prompt: {display_text}", 
+                                console=self.console,
+                                agent_name=self.name,
+                                agent_role=self.role,
+                                agent_tools=agent_tools
+                            )
+
+                    response = self._chat_completion(messages, temperature=temperature, tools=tools if tools else None, reasoning_steps=reasoning_steps)
                     if not response:
                         return None
+
+                    tool_calls = getattr(response.choices[0].message, 'tool_calls', None)
                     response_text = response.choices[0].message.content.strip()
 
-                # Handle output_json or output_pydantic if specified
-                if output_json or output_pydantic:
-                    # Add to chat history and return raw response
-                    self.chat_history.append({"role": "user", "content": original_prompt})
-                    self.chat_history.append({"role": "assistant", "content": response_text})
-                    if self.verbose:
-                        display_interaction(original_prompt, response_text, markdown=self.markdown, 
-                                         generation_time=time.time() - start_time, console=self.console)
-                    return response_text
+                    if tool_calls:
+                        messages.append({
+                            "role": "assistant",
+                            "content": response_text,
+                            "tool_calls": tool_calls
+                        })
+                        
+                        for tool_call in tool_calls:
+                            function_name = tool_call.function.name
+                            arguments = json.loads(tool_call.function.arguments)
 
-                if not self.self_reflect:
-                    self.chat_history.append({"role": "user", "content": original_prompt})
-                    self.chat_history.append({"role": "assistant", "content": response_text})
-                    if self.verbose:
-                        logging.info(f"Agent {self.name} final response: {response_text}")
-                    display_interaction(original_prompt, response_text, markdown=self.markdown, generation_time=time.time() - start_time, console=self.console)
-                    return response_text
+                            if self.verbose:
+                                display_tool_call(f"Agent {self.name} is calling function '{function_name}' with arguments: {arguments}", console=self.console)
+
+                            tool_result = self.execute_tool(function_name, arguments)
+
+                            if tool_result:
+                                if self.verbose:
+                                    display_tool_call(f"Function '{function_name}' returned: {tool_result}", console=self.console)
+                                messages.append({
+                                    "role": "tool",
+                                    "tool_call_id": tool_call.id,
+                                    "content": json.dumps(tool_result)
+                                })
+                            else:
+                                messages.append({
+                                    "role": "tool",
+                                    "tool_call_id": tool_call.id,
+                                    "content": "Function returned an empty output"
+                                })
+                            
+                        response = self._chat_completion(messages, temperature=temperature)
+                        if not response:
+                            return None
+                        response_text = response.choices[0].message.content.strip()
+
+                    # Handle output_json or output_pydantic if specified
+                    if output_json or output_pydantic:
+                        # Add to chat history and return raw response
+                        self.chat_history.append({"role": "user", "content": original_prompt})
+                        self.chat_history.append({"role": "assistant", "content": response_text})
+                        if self.verbose:
+                            display_interaction(original_prompt, response_text, markdown=self.markdown, 
+                                             generation_time=time.time() - start_time, console=self.console)
+                        return response_text
+
+                    if not self.self_reflect:
+                        self.chat_history.append({"role": "user", "content": original_prompt})
+                        self.chat_history.append({"role": "assistant", "content": response_text})
+                        if self.verbose:
+                            logging.debug(f"Agent {self.name} final response: {response_text}")
+                        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
 
-                reflection_prompt = f"""
+                    reflection_prompt = f"""
 Reflect on your previous response: '{response_text}'.
 Identify any flaws, improvements, or actions.
 Provide a "satisfactory" status ('yes' or 'no').
 Output MUST be JSON with 'reflection' and 'satisfactory'.
-                """
-                logging.debug(f"{self.name} reflection attempt {reflection_count+1}, sending prompt: {reflection_prompt}")
-                messages.append({"role": "user", "content": reflection_prompt})
+                    """
+                    logging.debug(f"{self.name} reflection attempt {reflection_count+1}, sending prompt: {reflection_prompt}")
+                    messages.append({"role": "user", "content": reflection_prompt})
 
-                try:
-                    reflection_response = client.beta.chat.completions.parse(
-                        model=self.reflect_llm if self.reflect_llm else self.llm,
-                        messages=messages,
-                        temperature=temperature,
-                        response_format=ReflectionOutput
-                    )
+                    try:
+                        reflection_response = client.beta.chat.completions.parse(
+                            model=self.reflect_llm if self.reflect_llm else self.llm,
+                            messages=messages,
+                            temperature=temperature,
+                            response_format=ReflectionOutput
+                        )
 
-                    reflection_output = reflection_response.choices[0].message.parsed
+                        reflection_output = reflection_response.choices[0].message.parsed
 
-                    if self.verbose:
-                        display_self_reflection(f"Agent {self.name} self reflection (using {self.reflect_llm if self.reflect_llm else self.llm}): reflection='{reflection_output.reflection}' satisfactory='{reflection_output.satisfactory}'", console=self.console)
+                        if self.verbose:
+                            display_self_reflection(f"Agent {self.name} self reflection (using {self.reflect_llm if self.reflect_llm else self.llm}): reflection='{reflection_output.reflection}' satisfactory='{reflection_output.satisfactory}'", console=self.console)
 
-                    messages.append({"role": "assistant", "content": f"Self Reflection: {reflection_output.reflection} Satisfactory?: {reflection_output.satisfactory}"})
+                        messages.append({"role": "assistant", "content": f"Self Reflection: {reflection_output.reflection} Satisfactory?: {reflection_output.satisfactory}"})
 
-                    # Only consider satisfactory after minimum reflections
-                    if reflection_output.satisfactory == "yes" and reflection_count >= self.min_reflect - 1:
-                        if self.verbose:
-                            display_self_reflection("Agent marked the response as satisfactory after meeting minimum reflections", console=self.console)
-                        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
+                        # Only consider satisfactory after minimum reflections
+                        if reflection_output.satisfactory == "yes" and reflection_count >= self.min_reflect - 1:
+                            if self.verbose:
+                                display_self_reflection("Agent marked the response as satisfactory after meeting minimum reflections", console=self.console)
+                            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
 
-                    # Check if we've hit max reflections
-                    if reflection_count >= self.max_reflect - 1:
-                        if self.verbose:
-                            display_self_reflection("Maximum reflection count reached, returning current response", console=self.console)
-                        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
+                        # Check if we've hit max reflections
+                        if reflection_count >= self.max_reflect - 1:
+                            if self.verbose:
+                                display_self_reflection("Maximum reflection count reached, returning current response", console=self.console)
+                            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
 
-                    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"})
-                    response = self._chat_completion(messages, temperature=temperature, tools=None, stream=True)
-                    response_text = response.choices[0].message.content.strip()
-                    reflection_count += 1
-                    continue  # Continue the loop for more reflections
+                        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"})
+                        response = self._chat_completion(messages, temperature=temperature, tools=None, stream=True)
+                        response_text = response.choices[0].message.content.strip()
+                        reflection_count += 1
+                        continue  # Continue the loop for more reflections
 
+                    except Exception as e:
+                        display_error(f"Error in parsing self-reflection json {e}. Retrying", console=self.console)
+                        logging.error("Reflection parsing failed.", exc_info=True)
+                        messages.append({"role": "assistant", "content": f"Self Reflection failed."})
+                        reflection_count += 1
+                        continue  # Continue even after error to try again
+                    
                 except Exception as e:
-                    display_error(f"Error in parsing self-reflection json {e}. Retrying", console=self.console)
-                    logging.error("Reflection parsing failed.", exc_info=True)
-                    messages.append({"role": "assistant", "content": f"Self Reflection failed."})
-                    reflection_count += 1
-                    continue  # Continue even after error to try again
-                
-            except Exception as e:
-                display_error(f"Error in chat: {e}", console=self.console)
-                return None 
+                    display_error(f"Error in chat: {e}", console=self.console)
+                    return None 
 
     def clean_json_output(self, output: str) -> str:
         """Clean and extract JSON from response text."""
@@ -599,85 +886,158 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
             cleaned = cleaned[len("```"):].strip()
         if cleaned.endswith("```"):
             cleaned = cleaned[:-3].strip()
-        return cleaned 
+        return cleaned  
 
-    async def achat(self, prompt, temperature=0.2, tools=None, output_json=None, output_pydantic=None):
-        """Async version of chat method"""
+    async def achat(self, prompt: str, temperature=0.2, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False):
+        """Async version of chat method. TODO: Requires Syncing with chat method.""" 
+        reasoning_steps = reasoning_steps or self.reasoning_steps
         try:
-            # Build system prompt
-            system_prompt = self.system_prompt
-            if output_json:
-                system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {output_json.schema_json()}"
-            elif output_pydantic:
-                system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {output_pydantic.schema_json()}"
-
-            # Build messages
-            if isinstance(prompt, str):
-                messages = [
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": prompt + ("\nReturn ONLY a valid JSON object. No other text or explanation." if (output_json or output_pydantic) else "")}
-                ]
+            # Search for existing knowledge if any knowledge is provided
+            if self.knowledge:
+                search_results = self.knowledge.search(prompt, agent_id=self.agent_id)
+                if search_results:
+                    if isinstance(search_results, dict) and 'results' in search_results:
+                        knowledge_content = "\n".join([result['memory'] for result in search_results['results']])
+                    else:
+                        knowledge_content = "\n".join(search_results)
+                    prompt = f"{prompt}\n\nKnowledge: {knowledge_content}"
+
+            if self._using_custom_llm:
+                try:
+                    response_text = await self.llm_instance.get_response_async(
+                        prompt=prompt,
+                        system_prompt=f"{self.backstory}\n\nYour Role: {self.role}\n\nYour Goal: {self.goal}" if self.use_system_prompt else None,
+                        chat_history=self.chat_history,
+                        temperature=temperature,
+                        tools=tools,
+                        output_json=output_json,
+                        output_pydantic=output_pydantic,
+                        verbose=self.verbose,
+                        markdown=self.markdown,
+                        self_reflect=self.self_reflect,
+                        max_reflect=self.max_reflect,
+                        min_reflect=self.min_reflect,
+                        console=self.console,
+                        agent_name=self.name,
+                        agent_role=self.role,
+                        agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in self.tools],
+                        execute_tool_fn=self.execute_tool_async,
+                        reasoning_steps=reasoning_steps
+                    )
+
+                    self.chat_history.append({"role": "user", "content": prompt})
+                    self.chat_history.append({"role": "assistant", "content": response_text})
+
+                    return response_text
+                except Exception as e:
+                    display_error(f"Error in LLM chat: {e}")
+                    return None
+
+            # For OpenAI client
+            if self.use_system_prompt:
+                system_prompt = f"""{self.backstory}\n
+Your Role: {self.role}\n
+Your Goal: {self.goal}
+                """
+                if output_json:
+                    system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(output_json.model_json_schema())}"
+                elif output_pydantic:
+                    system_prompt += f"\nReturn ONLY a JSON object that matches this Pydantic model: {json.dumps(output_pydantic.model_json_schema())}"
             else:
-                # For multimodal prompts
-                messages = [
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": prompt}
-                ]
-                if output_json or output_pydantic:
-                    # Add JSON instruction to text content
-                    for item in messages[-1]["content"]:
+                system_prompt = None
+
+            messages = []
+            if system_prompt:
+                messages.append({"role": "system", "content": system_prompt})
+            messages.extend(self.chat_history)
+
+            # Modify prompt if output_json or output_pydantic is specified
+            original_prompt = prompt
+            if output_json or output_pydantic:
+                if isinstance(prompt, str):
+                    prompt += "\nReturn ONLY a valid JSON object. No other text or explanation."
+                elif isinstance(prompt, list):
+                    for item in prompt:
                         if item["type"] == "text":
                             item["text"] += "\nReturn ONLY a valid JSON object. No other text or explanation."
                             break
 
-            # Format tools if provided
-            formatted_tools = []
-            if tools:
-                for tool in tools:
-                    if isinstance(tool, str):
-                        tool_def = self._generate_tool_definition(tool)
-                        if tool_def:
-                            formatted_tools.append(tool_def)
-                    elif isinstance(tool, dict):
-                        formatted_tools.append(tool)
-                    elif hasattr(tool, "to_openai_tool"):
-                        formatted_tools.append(tool.to_openai_tool())
-                    elif callable(tool):
-                        formatted_tools.append(self._generate_tool_definition(tool.__name__))
-
-            # Create async OpenAI client
-            async_client = AsyncOpenAI()
-
-            # Make the API call based on the type of request
-            if tools:
-                response = await async_client.chat.completions.create(
-                    model=self.llm,
-                    messages=messages,
-                    temperature=temperature,
-                    tools=formatted_tools
-                )
-                return await self._achat_completion(response, tools)
-            elif output_json or output_pydantic:
-                response = await async_client.chat.completions.create(
-                    model=self.llm,
-                    messages=messages,
-                    temperature=temperature,
-                    response_format={"type": "json_object"}
-                )
-                # Return the raw response
-                return response.choices[0].message.content
+            if isinstance(prompt, list):
+                messages.append({"role": "user", "content": prompt})
             else:
-                response = await async_client.chat.completions.create(
-                    model=self.llm,
-                    messages=messages,
-                    temperature=temperature
-                )
-                return response.choices[0].message.content
+                messages.append({"role": "user", "content": prompt})
+
+            reflection_count = 0
+            start_time = time.time()
+
+            while True:
+                try:
+                    if self.verbose:
+                        display_text = prompt
+                        if isinstance(prompt, list):
+                            display_text = next((item["text"] for item in prompt if item["type"] == "text"), "")
+                        
+                        if display_text and str(display_text).strip():
+                            agent_tools = [t.__name__ if hasattr(t, '__name__') else str(t) for t in self.tools]
+                            await adisplay_instruction(
+                                f"Agent {self.name} is processing prompt: {display_text}",
+                                console=self.console,
+                                agent_name=self.name,
+                                agent_role=self.role,
+                                agent_tools=agent_tools
+                            )
+
+                    # Format tools if provided
+                        formatted_tools = []
+                        if tools:
+                            for tool in tools:
+                                if isinstance(tool, str):
+                                    tool_def = self._generate_tool_definition(tool)
+                                    if tool_def:
+                                        formatted_tools.append(tool_def)
+                                elif isinstance(tool, dict):
+                                    formatted_tools.append(tool)
+                                elif hasattr(tool, "to_openai_tool"):
+                                    formatted_tools.append(tool.to_openai_tool())
+                                elif callable(tool):
+                                    formatted_tools.append(self._generate_tool_definition(tool.__name__))
+
+                        # Create async OpenAI client
+                        async_client = AsyncOpenAI()
+
+                        # Make the API call based on the type of request
+                        if tools:
+                            response = await async_client.chat.completions.create(
+                                model=self.llm,
+                                messages=messages,
+                                temperature=temperature,
+                                tools=formatted_tools
+                            )
+                            return await self._achat_completion(response, tools)
+                        elif output_json or output_pydantic:
+                            response = await async_client.chat.completions.create(
+                                model=self.llm,
+                                messages=messages,
+                                temperature=temperature,
+                                response_format={"type": "json_object"}
+                            )
+                            # Return the raw response
+                            return response.choices[0].message.content
+                        else:
+                            response = await async_client.chat.completions.create(
+                                model=self.llm,
+                                messages=messages,
+                                temperature=temperature
+                            )
+                            return response.choices[0].message.content
+                except Exception as e:
+                    display_error(f"Error in chat completion: {e}")
+                    return None
         except Exception as e:
-            display_error(f"Error in chat completion: {e}")
+            display_error(f"Error in achat: {e}")
             return None
 
-    async def _achat_completion(self, response, tools):
+    async def _achat_completion(self, response, tools, reasoning_steps=False):
         """Async version of _chat_completion method"""
         try:
             message = response.choices[0].message
@@ -723,9 +1083,42 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
                         final_response = await async_client.chat.completions.create(
                             model=self.llm,
                             messages=messages,
-                            temperature=0.2
+                            temperature=0.2,
+                            stream=True
                         )
-                        return final_response.choices[0].message.content
+                        full_response_text = ""
+                        reasoning_content = ""
+                        chunks = []
+                        start_time = time.time()
+                        
+                        with Live(
+                            display_generating("", start_time),
+                            console=self.console,
+                            refresh_per_second=4,
+                            transient=True,
+                            vertical_overflow="ellipsis",
+                            auto_refresh=True
+                        ) as live:
+                            async for chunk in final_response:
+                                chunks.append(chunk)
+                                if chunk.choices[0].delta.content:
+                                    full_response_text += chunk.choices[0].delta.content
+                                    live.update(display_generating(full_response_text, start_time))
+                                
+                                if reasoning_steps and hasattr(chunk.choices[0].delta, "reasoning_content"):
+                                    rc = chunk.choices[0].delta.reasoning_content
+                                    if rc:
+                                        reasoning_content += rc
+                                        live.update(display_generating(f"{full_response_text}\n[Reasoning: {reasoning_content}]", start_time))
+                        
+                        self.console.print()
+                        
+                        final_response = process_stream_chunks(chunks)
+                        # Return only reasoning content if reasoning_steps is True
+                        if reasoning_steps and hasattr(final_response.choices[0].message, 'reasoning_content'):
+                            return final_response.choices[0].message.reasoning_content
+                        return final_response.choices[0].message.content if final_response else full_response_text
+
                     except Exception as e:
                         display_error(f"Error in final chat completion: {e}")
                         return formatted_results
@@ -733,8 +1126,59 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
             return None
         except Exception as e:
             display_error(f"Error in _achat_completion: {e}")
-            return None 
+            return None
+
+    async def astart(self, prompt: str, **kwargs):
+        """Async version of start method"""
+        return await self.achat(prompt, **kwargs)
 
     def run(self):
         """Alias for start() method"""
-        return self.start() 
+        return self.start() 
+
+    def start(self, prompt: str, **kwargs):
+        """Start the agent with a prompt. This is a convenience method that wraps chat()."""
+        return self.chat(prompt, **kwargs) 
+
+    async def execute_tool_async(self, function_name: str, arguments: Dict[str, Any]) -> Any:
+        """Async version of execute_tool"""
+        try:
+            logging.info(f"Executing async tool: {function_name} with arguments: {arguments}")
+            # Try to find the function in the agent's tools list first
+            func = None
+            for tool in self.tools:
+                if (callable(tool) and getattr(tool, '__name__', '') == function_name):
+                    func = tool
+                    break
+            
+            if func is None:
+                logging.error(f"Function {function_name} not found in tools")
+                return {"error": f"Function {function_name} not found in tools"}
+
+            try:
+                if inspect.iscoroutinefunction(func):
+                    logging.debug(f"Executing async function: {function_name}")
+                    result = await func(**arguments)
+                else:
+                    logging.debug(f"Executing sync function in executor: {function_name}")
+                    loop = asyncio.get_event_loop()
+                    result = await loop.run_in_executor(None, lambda: func(**arguments))
+                
+                # Ensure result is JSON serializable
+                logging.debug(f"Raw result from tool: {result}")
+                if result is None:
+                    return {"result": None}
+                try:
+                    json.dumps(result)  # Test serialization
+                    return result
+                except TypeError:
+                    logging.warning(f"Result not JSON serializable, converting to string: {result}")
+                    return {"result": str(result)}
+
+            except Exception as e:
+                logging.error(f"Error executing {function_name}: {str(e)}", exc_info=True)
+                return {"error": f"Error executing {function_name}: {str(e)}"}
+
+        except Exception as e:
+            logging.error(f"Error in execute_tool_async: {str(e)}", exc_info=True)
+            return {"error": f"Error in execute_tool_async: {str(e)}"}