praisonaiagents 0.0.141__py3-none-any.whl → 0.0.143__py3-none-any.whl

praisonaiagents/llm/openai_client.py

@@ -1049,6 +1049,150 @@ class OpenAIClient:
                 break
         
         return final_response
+        
+    def chat_completion_with_tools_stream(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str = "gpt-4o",
+        temperature: float = 0.7,
+        tools: Optional[List[Any]] = None,
+        execute_tool_fn: Optional[Callable] = None,
+        reasoning_steps: bool = False,
+        verbose: bool = True,
+        max_iterations: int = 10,
+        **kwargs
+    ):
+        """
+        Create a streaming chat completion with tool support.
+        
+        This method yields chunks of the response as they are generated,
+        enabling real-time streaming to the user.
+        
+        Args:
+            messages: List of message dictionaries
+            model: Model to use
+            temperature: Temperature for generation
+            tools: List of tools (can be callables, dicts, or strings)
+            execute_tool_fn: Function to execute tools
+            reasoning_steps: Whether to show reasoning
+            verbose: Whether to show verbose output
+            max_iterations: Maximum tool calling iterations
+            **kwargs: Additional API parameters
+            
+        Yields:
+            String chunks of the response as they are generated
+        """
+        # Format tools for OpenAI API
+        formatted_tools = self.format_tools(tools)
+        
+        # Continue tool execution loop until no more tool calls are needed
+        iteration_count = 0
+        
+        while iteration_count < max_iterations:
+            try:
+                # Create streaming response
+                response_stream = self._sync_client.chat.completions.create(
+                    model=model,
+                    messages=messages,
+                    temperature=temperature,
+                    tools=formatted_tools if formatted_tools else None,
+                    stream=True,
+                    **kwargs
+                )
+                
+                full_response_text = ""
+                reasoning_content = ""
+                chunks = []
+                
+                # Stream the response chunk by chunk
+                for chunk in response_stream:
+                    chunks.append(chunk)
+                    if chunk.choices and chunk.choices[0].delta.content:
+                        content = chunk.choices[0].delta.content
+                        full_response_text += content
+                        yield content
+                    
+                    # Handle reasoning content if enabled
+                    if reasoning_steps and chunk.choices and hasattr(chunk.choices[0].delta, "reasoning_content"):
+                        rc = chunk.choices[0].delta.reasoning_content
+                        if rc:
+                            reasoning_content += rc
+                            yield f"[Reasoning: {rc}]"
+                
+                # Process the complete response to check for tool calls
+                final_response = process_stream_chunks(chunks)
+                
+                if not final_response:
+                    return
+                
+                # Check for tool calls
+                tool_calls = getattr(final_response.choices[0].message, 'tool_calls', None)
+                
+                if tool_calls and execute_tool_fn:
+                    # Convert ToolCall dataclass objects to dict for JSON serialization
+                    serializable_tool_calls = []
+                    for tc in tool_calls:
+                        if isinstance(tc, ToolCall):
+                            # Convert dataclass to dict
+                            serializable_tool_calls.append({
+                                "id": tc.id,
+                                "type": tc.type,
+                                "function": tc.function
+                            })
+                        else:
+                            # Already an OpenAI object, keep as is
+                            serializable_tool_calls.append(tc)
+                    
+                    messages.append({
+                        "role": "assistant", 
+                        "content": final_response.choices[0].message.content,
+                        "tool_calls": serializable_tool_calls
+                    })
+                    
+                    for tool_call in tool_calls:
+                        # Handle both ToolCall dataclass and OpenAI object
+                        try:
+                            if isinstance(tool_call, ToolCall):
+                                function_name = tool_call.function["name"]
+                                arguments = json.loads(tool_call.function["arguments"])
+                            else:
+                                function_name = tool_call.function.name
+                                arguments = json.loads(tool_call.function.arguments)
+                        except json.JSONDecodeError as e:
+                            if verbose:
+                                yield f"\n[Error parsing arguments for {function_name if 'function_name' in locals() else 'unknown function'}: {str(e)}]"
+                            continue
+                        
+                        if verbose:
+                            yield f"\n[Calling function: {function_name}]"
+                        
+                        # Execute the tool with error handling
+                        try:
+                            tool_result = execute_tool_fn(function_name, arguments)
+                            results_str = json.dumps(tool_result) if tool_result else "Function returned an empty output"
+                        except Exception as e:
+                            results_str = f"Error executing function: {str(e)}"
+                            if verbose:
+                                yield f"\n[Function error: {str(e)}]"
+                        
+                        if verbose:
+                            yield f"\n[Function result: {results_str}]"
+                        
+                        messages.append({
+                            "role": "tool",
+                            "tool_call_id": tool_call.id if hasattr(tool_call, 'id') else tool_call['id'],
+                            "content": results_str
+                        })
+                    
+                    # Continue the loop to allow more tool calls
+                    iteration_count += 1
+                else:
+                    # No tool calls, we're done
+                    break
+                    
+            except Exception as e:
+                yield f"Error: {str(e)}"
+                break
     
     def parse_structured_output(
         self,

praisonaiagents/mcp/mcp.py

@@ -25,6 +25,8 @@ class MCPToolRunner(threading.Thread):
         self.initialized = threading.Event()
         self.tools = []
         self.timeout = timeout
+        self._tool_timings = {}
+        self._timings_lock = threading.Lock()
         self.start()
         
     def run(self):
@@ -75,25 +77,63 @@ class MCPToolRunner(threading.Thread):
     
     def call_tool(self, tool_name, arguments):
         """Call an MCP tool and wait for the result."""
+        # Import telemetry here to avoid circular imports
+        try:
+            from ..telemetry.telemetry import get_telemetry
+            telemetry = get_telemetry()
+        except (ImportError, AttributeError):
+            telemetry = None
+        
+        # Check initialization first (without timing)
         if not self.initialized.is_set():
             self.initialized.wait(timeout=self.timeout)
             if not self.initialized.is_set():
+                # Track initialization timeout failure
+                if telemetry:
+                    telemetry.track_tool_usage(tool_name, success=False, execution_time=0)
                 return f"Error: MCP initialization timed out after {self.timeout} seconds"
         
-        # Put request in queue
-        self.queue.put((tool_name, arguments))
-        
-        # Wait for result
-        success, result = self.result_queue.get()
-        if not success:
-            return f"Error: {result}"
-        
-        # Process result
-        if hasattr(result, 'content') and result.content:
-            if hasattr(result.content[0], 'text'):
-                return result.content[0].text
-            return str(result.content[0])
-        return str(result)
+        # Start timing after initialization check
+        start_time = time.time()
+        is_success = False
+        try:
+            # Put request in queue
+            self.queue.put((tool_name, arguments))
+            
+            # Wait for result
+            success, result = self.result_queue.get()
+            if not success:
+                return f"Error: {result}"
+            
+            # Process result
+            if hasattr(result, 'content') and result.content:
+                if hasattr(result.content[0], 'text'):
+                    processed_result = result.content[0].text
+                else:
+                    processed_result = str(result.content[0])
+            else:
+                processed_result = str(result)
+            
+            is_success = True
+            return processed_result
+            
+        except Exception as e:
+            return f"Error: {str(e)}"
+        finally:
+            # Track timing regardless of success/failure
+            end_time = time.time()
+            execution_time = end_time - start_time
+            
+            # Log timing information for debugging
+            logging.debug(f"Tool '{tool_name}' execution time: {execution_time:.3f} seconds")
+            
+            # Store timing in thread-safe manner
+            with self._timings_lock:
+                self._tool_timings[tool_name] = execution_time
+            
+            # Track tool usage with timing information
+            if telemetry:
+                telemetry.track_tool_usage(tool_name, success=is_success, execution_time=execution_time)
     
     def shutdown(self):
         """Signal the thread to shut down."""