diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 533b2173..4d054a7a 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -164,6 +164,54 @@ class _StreamingState: # Add a new function for consistent token counting using tiktoken +def _check_reasoning_compatibility(messages): + """ + Check if message history is compatible with Claude reasoning/thinking. + + According to Claude 4 docs, when reasoning is enabled, the final assistant + message must start with a thinking block. If there are assistant messages + with regular text content, reasoning should be disabled. + + Args: + messages: List of message dictionaries + + Returns: + bool: True if compatible with reasoning, False otherwise + """ + if not messages: + return True # Empty messages are compatible + + # Find the last assistant message + last_assistant_msg = None + for msg in reversed(messages): + if msg.get("role") == "assistant": + last_assistant_msg = msg + break + + if not last_assistant_msg: + return True # No assistant messages, compatible + + # Check if the last assistant message has regular text content + content = last_assistant_msg.get("content") + if content: + # If it's a string with text content, not compatible + if isinstance(content, str) and content.strip(): + return False + # If it's a list, check for text content blocks + elif isinstance(content, list): + for block in content: + if isinstance(block, dict): + if block.get("type") == "text" and block.get("text", "").strip(): + return False + + # Check if message has tool_calls (these are compatible) + if last_assistant_msg.get("tool_calls"): + return True + + # If no content or only thinking blocks, it's compatible + return True + + def count_tokens_with_tiktoken(text_or_messages): """ Count tokens consistently using tiktoken library. @@ -305,6 +353,60 @@ class OpenAIChatCompletionsModel(Model): "role": "system", }, ) + + # Add support for prompt caching for claude (not automatically applied) + # Gemini supports it too + # https://www.anthropic.com/news/token-saving-updates + # Maximize cache efficiency by using up to 4 cache_control blocks + if ((str(self.model).startswith("claude") or + "gemini" in str(self.model)) and + len(converted_messages) > 0): + + # Strategy: Cache the most valuable messages for maximum savings + # 1. System message (always first priority) + # 2. Long user messages (high token count) + # 3. Assistant messages with tool calls (complex context) + # 4. Recent context (last message) + + cache_candidates = [] + + # Always cache system message if present + for i, msg in enumerate(converted_messages): + if msg.get("role") == "system": + cache_candidates.append((i, len(str(msg.get("content", ""))), "system")) + break + + # Find long user messages and assistant messages with tool calls + for i, msg in enumerate(converted_messages): + content_len = len(str(msg.get("content", ""))) + role = msg.get("role") + + if role == "user" and content_len > 500: # Long user messages + cache_candidates.append((i, content_len, "user")) + elif role == "assistant" and msg.get("tool_calls"): # Tool calls + cache_candidates.append((i, content_len + 200, "assistant_tools")) # Bonus for tool calls + + # Always consider the last message for recent context + if len(converted_messages) > 1: + last_idx = len(converted_messages) - 1 + last_msg = converted_messages[last_idx] + last_content_len = len(str(last_msg.get("content", ""))) + cache_candidates.append((last_idx, last_content_len, "recent")) + + # Sort by value (content length) and select top 4 unique indices + cache_candidates.sort(key=lambda x: x[1], reverse=True) + selected_indices = [] + for idx, _, msg_type in cache_candidates: + if idx not in selected_indices: + selected_indices.append(idx) + if len(selected_indices) >= 4: # Max 4 cache blocks + break + + # Apply cache_control to selected messages + for idx in selected_indices: + msg_copy = converted_messages[idx].copy() + msg_copy["cache_control"] = {"type": "ephemeral"} + converted_messages[idx] = msg_copy # --- Add to message_history: user, system, and assistant tool call messages --- # Add system prompt to message_history if system_instructions: @@ -689,6 +791,60 @@ class OpenAIChatCompletionsModel(Model): "role": "system", }, ) + + # Add support for prompt caching for claude (not automatically applied) + # Gemini supports it too + # https://www.anthropic.com/news/token-saving-updates + # Maximize cache efficiency by using up to 4 cache_control blocks + if ((str(self.model).startswith("claude") or + "gemini" in str(self.model)) and + len(converted_messages) > 0): + + # Strategy: Cache the most valuable messages for maximum savings + # 1. System message (always first priority) + # 2. Long user messages (high token count) + # 3. Assistant messages with tool calls (complex context) + # 4. Recent context (last message) + + cache_candidates = [] + + # Always cache system message if present + for i, msg in enumerate(converted_messages): + if msg.get("role") == "system": + cache_candidates.append((i, len(str(msg.get("content", ""))), "system")) + break + + # Find long user messages and assistant messages with tool calls + for i, msg in enumerate(converted_messages): + content_len = len(str(msg.get("content", ""))) + role = msg.get("role") + + if role == "user" and content_len > 500: # Long user messages + cache_candidates.append((i, content_len, "user")) + elif role == "assistant" and msg.get("tool_calls"): # Tool calls + cache_candidates.append((i, content_len + 200, "assistant_tools")) # Bonus for tool calls + + # Always consider the last message for recent context + if len(converted_messages) > 1: + last_idx = len(converted_messages) - 1 + last_msg = converted_messages[last_idx] + last_content_len = len(str(last_msg.get("content", ""))) + cache_candidates.append((last_idx, last_content_len, "recent")) + + # Sort by value (content length) and select top 4 unique indices + cache_candidates.sort(key=lambda x: x[1], reverse=True) + selected_indices = [] + for idx, _, msg_type in cache_candidates: + if idx not in selected_indices: + selected_indices.append(idx) + if len(selected_indices) >= 4: # Max 4 cache blocks + break + + # Apply cache_control to selected messages + for idx in selected_indices: + msg_copy = converted_messages[idx].copy() + msg_copy["cache_control"] = {"type": "ephemeral"} + converted_messages[idx] = msg_copy # --- Add to message_history: user, system prompts --- if system_instructions: sys_msg = { @@ -765,6 +921,16 @@ class OpenAIChatCompletionsModel(Model): # For tool call streaming, accumulate tool_calls to add to message_history at the end streamed_tool_calls = [] + # Initialize Claude thinking display if applicable + thinking_context = None + from cai.util import start_claude_thinking_if_applicable + if should_show_rich_stream: # Only show thinking in rich streaming mode + thinking_context = start_claude_thinking_if_applicable( + str(self.model), + self.agent_name, + self.interaction_counter + ) + # Ollama specific: accumulate full content to check for function calls at the end # Some Ollama models output the function call as JSON in the text content ollama_full_content = "" @@ -826,6 +992,34 @@ class OpenAIChatCompletionsModel(Model): if not delta: continue + # Handle Claude reasoning content first (before regular content) + reasoning_content = None + if hasattr(delta, 'reasoning_content') and delta.reasoning_content is not None: + reasoning_content = delta.reasoning_content + elif isinstance(delta, dict) and 'reasoning_content' in delta and delta['reasoning_content'] is not None: + reasoning_content = delta['reasoning_content'] + + # Also check for thinking_blocks structure + thinking_blocks = None + if hasattr(delta, 'thinking_blocks') and delta.thinking_blocks is not None: + thinking_blocks = delta.thinking_blocks + elif isinstance(delta, dict) and 'thinking_blocks' in delta and delta['thinking_blocks'] is not None: + thinking_blocks = delta['thinking_blocks'] + + # Extract reasoning content from thinking blocks if available + if thinking_blocks and not reasoning_content: + for block in thinking_blocks: + if isinstance(block, dict) and block.get('type') == 'thinking': + reasoning_content = block.get('thinking', '') + break + + if reasoning_content and thinking_context: + # Update the thinking display + from cai.util import update_claude_thinking_content + update_claude_thinking_content(thinking_context, reasoning_content) + + + # Handle text content = None if hasattr(delta, 'content') and delta.content is not None: @@ -1181,6 +1375,14 @@ class OpenAIChatCompletionsModel(Model): finish_agent_streaming(streaming_context, None) except Exception: pass + + # Ensure thinking context is cleaned up in case of errors + if thinking_context: + try: + from cai.util import finish_claude_thinking_display + finish_claude_thinking_display(thinking_context) + except Exception: + pass raise e # Special handling for Ollama - check if accumulated text contains a valid function call @@ -1537,6 +1739,11 @@ class OpenAIChatCompletionsModel(Model): # Removed extra newline after streaming completes to avoid blank lines pass + # Finish Claude thinking display if it was active + if thinking_context: + from cai.util import finish_claude_thinking_display + finish_claude_thinking_display(thinking_context) + if tracing.include_data(): span_generation.span_data.output = [final_response.model_dump()] @@ -1597,6 +1804,14 @@ class OpenAIChatCompletionsModel(Model): except Exception: pass + # Ensure thinking context is cleaned up in case of errors + if 'thinking_context' in locals() and thinking_context: + try: + from cai.util import finish_claude_thinking_display + finish_claude_thinking_display(thinking_context) + except Exception: + pass + # Stop active timer and start idle timer when streaming errors out stop_active_timer() start_idle_timer() @@ -1657,6 +1872,60 @@ class OpenAIChatCompletionsModel(Model): "role": "system", }, ) + + # Add support for prompt caching for claude (not automatically applied) + # Gemini supports it too + # https://www.anthropic.com/news/token-saving-updates + # Maximize cache efficiency by using up to 4 cache_control blocks + if ((str(self.model).startswith("claude") or + "gemini" in str(self.model)) and + len(converted_messages) > 0): + + # Strategy: Cache the most valuable messages for maximum savings + # 1. System message (always first priority) + # 2. Long user messages (high token count) + # 3. Assistant messages with tool calls (complex context) + # 4. Recent context (last message) + + cache_candidates = [] + + # Always cache system message if present + for i, msg in enumerate(converted_messages): + if msg.get("role") == "system": + cache_candidates.append((i, len(str(msg.get("content", ""))), "system")) + break + + # Find long user messages and assistant messages with tool calls + for i, msg in enumerate(converted_messages): + content_len = len(str(msg.get("content", ""))) + role = msg.get("role") + + if role == "user" and content_len > 500: # Long user messages + cache_candidates.append((i, content_len, "user")) + elif role == "assistant" and msg.get("tool_calls"): # Tool calls + cache_candidates.append((i, content_len + 200, "assistant_tools")) # Bonus for tool calls + + # Always consider the last message for recent context + if len(converted_messages) > 1: + last_idx = len(converted_messages) - 1 + last_msg = converted_messages[last_idx] + last_content_len = len(str(last_msg.get("content", ""))) + cache_candidates.append((last_idx, last_content_len, "recent")) + + # Sort by value (content length) and select top 4 unique indices + cache_candidates.sort(key=lambda x: x[1], reverse=True) + selected_indices = [] + for idx, _, msg_type in cache_candidates: + if idx not in selected_indices: + selected_indices.append(idx) + if len(selected_indices) >= 4: # Max 4 cache blocks + break + + # Apply cache_control to selected messages + for idx in selected_indices: + msg_copy = converted_messages[idx].copy() + msg_copy["cache_control"] = {"type": "ephemeral"} + converted_messages[idx] = msg_copy if tracing.include_data(): span.span_data.input = converted_messages @@ -1751,6 +2020,27 @@ class OpenAIChatCompletionsModel(Model): # Remove tool_choice if no tools are specified if not converted_tools: kwargs.pop("tool_choice", None) + + # Add extended reasoning support for Claude models + # Supports Claude 3.7, Claude 4, and any model with "thinking" in the name + has_reasoning_capability = ("3.7" in model_str or "4" in model_str or "thinking" in model_str) + + if has_reasoning_capability: + # Clean the model name by removing "thinking" before sending to API + clean_model = kwargs["model"] + if isinstance(clean_model, str) and "thinking" in clean_model.lower(): + # Remove "thinking" and clean up any extra spaces/separators + clean_model = re.sub(r'[_-]?thinking[_-]?', '', clean_model, flags=re.IGNORECASE) + clean_model = re.sub(r'[-_]{2,}', '-', clean_model) # Clean up multiple separators + clean_model = clean_model.strip('-_') # Clean up leading/trailing separators + kwargs["model"] = clean_model + + # Check if message history is compatible with reasoning + messages = kwargs.get("messages", []) + is_compatible = _check_reasoning_compatibility(messages) + + if is_compatible: + kwargs["reasoning_effort"] = "low" # Use reasoning_effort instead of thinking elif provider == "gemini": kwargs.pop("parallel_tool_calls", None) # Add any specific gemini settings if needed @@ -1764,6 +2054,27 @@ class OpenAIChatCompletionsModel(Model): # Remove tool_choice if no tools are specified if not converted_tools: kwargs.pop("tool_choice", None) + + # Add extended reasoning support for Claude models + # Supports Claude 3.7, Claude 4, and any model with "thinking" in the name + has_reasoning_capability = ("3.7" in model_str or "4" in model_str or "thinking" in model_str) + + if has_reasoning_capability: + # Clean the model name by removing "thinking" before sending to API + clean_model = kwargs["model"] + if isinstance(clean_model, str) and "thinking" in clean_model.lower(): + # Remove "thinking" and clean up any extra spaces/separators + clean_model = re.sub(r'[_-]?thinking[_-]?', '', clean_model, flags=re.IGNORECASE) + clean_model = re.sub(r'[-_]{2,}', '-', clean_model) # Clean up multiple separators + clean_model = clean_model.strip('-_') # Clean up leading/trailing separators + kwargs["model"] = clean_model + + # Check if message history is compatible with reasoning + messages = kwargs.get("messages", []) + is_compatible = _check_reasoning_compatibility(messages) + + if is_compatible: + kwargs["reasoning_effort"] = "low" # Use reasoning_effort instead of thinking elif "gemini" in model_str: kwargs.pop("parallel_tool_calls", None) elif "qwen" in model_str or ":" in model_str: @@ -1802,6 +2113,41 @@ class OpenAIChatCompletionsModel(Model): return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) except litellm.exceptions.BadRequestError as e: + error_msg = str(e) + + # Handle Claude reasoning/thinking compatibility errors + if ("Expected `thinking` or `redacted_thinking`, but found `text`" in error_msg or + "When `thinking` is enabled, a final `assistant` message must start with a thinking block" in error_msg): + print(f"Claude 4 thinking error: {e}") + print("Disabling reasoning_effort due to incompatible message history") + + # Retry without reasoning_effort + retry_kwargs = kwargs.copy() + retry_kwargs.pop("reasoning_effort", None) + + try: + if stream: + response = Response( + id=FAKE_RESPONSES_ID, + created_at=time.time(), + model=self.model, + object="response", + output=[], + tool_choice="auto" if tool_choice is None or tool_choice == NOT_GIVEN else cast(Literal["auto", "required", "none"], tool_choice), + top_p=model_settings.top_p, + temperature=model_settings.temperature, + tools=[], + parallel_tool_calls=parallel_tool_calls or False, + ) + stream_obj = await litellm.acompletion(**retry_kwargs) + return response, stream_obj + else: + ret = litellm.completion(**retry_kwargs) + return ret + except Exception as retry_e: + # If retry also fails, raise the original error + raise e + #print(color("BadRequestError encountered: " + str(e), fg="yellow")) if "LLM Provider NOT provided" in str(e): model_str = str(self.model).lower() @@ -1867,6 +2213,25 @@ class OpenAIChatCompletionsModel(Model): provider_kwargs["custom_llm_provider"] = "anthropic" provider_kwargs.pop("store", None) # Claude doesn't support store parameter provider_kwargs.pop("parallel_tool_calls", None) # Claude doesn't support parallel tool calls + + # Add extended reasoning support for Claude models + # Supports Claude 3.7, Claude 4, and any model with "thinking" in the name + if ("3.7" in model_str or "4" in model_str or "thinking" in model_str): + # Clean the model name by removing "thinking" before sending to API + clean_model = provider_kwargs["model"] + if isinstance(clean_model, str) and "thinking" in clean_model.lower(): + # Remove "thinking" and clean up any extra spaces/separators + clean_model = re.sub(r'[_-]?thinking[_-]?', '', clean_model, flags=re.IGNORECASE) + clean_model = re.sub(r'[-_]{2,}', '-', clean_model) # Clean up multiple separators + clean_model = clean_model.strip('-_') # Clean up leading/trailing separators + provider_kwargs["model"] = clean_model + + # Check if message history is compatible with reasoning + messages = provider_kwargs.get("messages", []) + is_compatible = _check_reasoning_compatibility(messages) + + if is_compatible: + provider_kwargs["reasoning_effort"] = "low" # Use reasoning_effort instead of thinking elif provider == "gemini": provider_kwargs["custom_llm_provider"] = "gemini" provider_kwargs.pop("store", None) # Gemini doesn't support store parameter diff --git a/src/cai/util.py b/src/cai/util.py index 2e77c45a..45206e7c 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -44,6 +44,11 @@ _timing_lock = threading.Lock() # Set up a global tracker for live streaming panels _LIVE_STREAMING_PANELS = {} +# ======================== CLAUDE THINKING STREAMING FUNCTIONS ======================== + +# Global tracker for Claude thinking streaming panels +_CLAUDE_THINKING_PANELS = {} + def start_active_timer(): """ Start measuring active time (when LLM is processing or tool is executing). @@ -3020,3 +3025,277 @@ def setup_ctf(): bg="blue")) return ctf, messages + +def create_claude_thinking_context(agent_name, counter, model): + """ + Create a streaming context for Claude thinking/reasoning display. + This creates a dedicated panel that shows Claude's internal reasoning process. + + Args: + agent_name: The name of the agent + counter: The interaction counter + model: The model name + + Returns: + A dictionary with the streaming context for thinking display + """ + import uuid + from rich.live import Live + from rich.panel import Panel + from rich.text import Text + from rich.box import ROUNDED + from rich.console import Group + import shutil + + # Generate unique thinking context ID + thinking_id = f"thinking_{agent_name}_{counter}_{str(uuid.uuid4())[:8]}" + + # Check if we already have an active thinking panel + if thinking_id in _CLAUDE_THINKING_PANELS: + return _CLAUDE_THINKING_PANELS[thinking_id] + + try: + timestamp = datetime.now().strftime("%H:%M:%S") + + # Terminal size for better display + terminal_width, _ = shutil.get_terminal_size((100, 24)) + panel_width = min(terminal_width - 4, 120) + + # Create the thinking panel header + header = Text() + header.append("🧠 ", style="bold yellow") + header.append(f"Claude Reasoning [{counter}]", style="bold yellow") + header.append(f" | {agent_name}", style="bold cyan") + header.append(f" | {timestamp}", style="dim") + + # Initial thinking content + thinking_content = Text("Thinking...", style="italic dim") + + # Create the panel for thinking + panel = Panel( + Group(header, Text("\n"), thinking_content), + title="[bold yellow]🧠 Thinking Process[/bold yellow]", + border_style="yellow", + box=ROUNDED, + padding=(1, 2), + width=panel_width, + expand=True + ) + + # Create Live display object + live = Live(panel, refresh_per_second=8, console=console, auto_refresh=True) + + context = { + "thinking_id": thinking_id, + "live": live, + "panel": panel, + "header": header, + "thinking_content": thinking_content, + "timestamp": timestamp, + "model": model, + "agent_name": agent_name, + "panel_width": panel_width, + "is_started": False, + "accumulated_thinking": "", + } + + # Store in global tracker + _CLAUDE_THINKING_PANELS[thinking_id] = context + + return context + + except Exception as e: + print(f"Error creating Claude thinking context: {e}") + return None + +def update_claude_thinking_content(context, thinking_delta): + """ + Update the Claude thinking content with new reasoning text. + + Args: + context: The thinking context created by create_claude_thinking_context + thinking_delta: The new thinking text to add + """ + if not context: + return False + + try: + # Accumulate the thinking text + context["accumulated_thinking"] += thinking_delta + + # Create syntax highlighted thinking content + from rich.syntax import Syntax + from rich.text import Text + from rich.console import Group + + # Try to format as markdown-like reasoning + thinking_text = context["accumulated_thinking"] + + # Create formatted thinking display + if len(thinking_text) > 500: + # For long thinking, use syntax highlighting + thinking_display = Syntax( + thinking_text, + "markdown", + theme="monokai", + background_color="#2E2E2E", + word_wrap=True, + line_numbers=False + ) + else: + # For short thinking, use regular text with styling + thinking_display = Text(thinking_text, style="white") + + # Update the panel content + updated_panel = Panel( + Group( + context["header"], + Text("\n"), + thinking_display + ), + title="[bold yellow]🧠 Thinking Process[/bold yellow]", + border_style="yellow", + box=ROUNDED, + padding=(1, 2), + width=context.get("panel_width", 100), + expand=True + ) + + # Start the display if not already started + if not context.get("is_started", False): + try: + context["live"].start() + context["is_started"] = True + except Exception as e: + print(f"Error starting Claude thinking display: {e}") + return False + + # Update the live display + context["live"].update(updated_panel) + context["panel"] = updated_panel + context["live"].refresh() + + return True + + except Exception as e: + print(f"Error updating Claude thinking content: {e}") + return False + +def finish_claude_thinking_display(context): + """ + Finish the Claude thinking display session. + + Args: + context: The thinking context to finish + """ + if not context: + return False + + # Clean up from global tracker + thinking_id = context.get("thinking_id") + if thinking_id and thinking_id in _CLAUDE_THINKING_PANELS: + del _CLAUDE_THINKING_PANELS[thinking_id] + + try: + # Import required classes + from rich.text import Text + from rich.syntax import Syntax + from rich.console import Group + + # Add final formatting to show completion + final_header = Text() + final_header.append("🧠 ", style="bold green") + final_header.append(f"Claude Reasoning Complete", style="bold green") + final_header.append(f" | {context['agent_name']}", style="bold cyan") + final_header.append(f" | {context['timestamp']}", style="dim") + + thinking_text = context["accumulated_thinking"] + + if thinking_text.strip(): + # Create final formatted display + final_thinking_display = Syntax( + thinking_text, + "markdown", + theme="monokai", + background_color="#2E2E2E", + word_wrap=True, + line_numbers=False + ) + else: + final_thinking_display = Text("No reasoning captured", style="dim italic") + + # Create final panel + final_panel = Panel( + Group( + final_header, + Text("\n"), + final_thinking_display + ), + title="[bold green]🧠 Thinking Complete[/bold green]", + border_style="green", + box=ROUNDED, + padding=(1, 2), + width=context.get("panel_width", 100), + expand=True + ) + + # Update one last time + if context.get("is_started", False): + context["live"].update(final_panel) + + # Give a moment for the final panel to be seen + import time + time.sleep(0.3) + + # Stop the live display + context["live"].stop() + + return True + + except Exception as e: + print(f"Error finishing Claude thinking display: {e}") + return False + +def detect_claude_thinking_in_stream(model_name): + """ + Detect if a model should show thinking/reasoning display. + Only applies to Claude models with reasoning capability. + + Args: + model_name: The model name to check + + Returns: + bool: True if thinking display should be shown + """ + if not model_name: + return False + + model_str = str(model_name).lower() + + # Check for Claude models with reasoning capability + # According to LiteLLM docs, Claude 3.7, Claude 4 and models with "thinking" support reasoning + has_reasoning = ( + "claude" in model_str and ( + "3.7" in model_str or + "4" in model_str or + "thinking" in model_str + ) + ) + + return has_reasoning + +def start_claude_thinking_if_applicable(model_name, agent_name, counter): + """ + Start Claude thinking display if the model supports it. + + Args: + model_name: The model name + agent_name: The agent name + counter: The interaction counter + + Returns: + The thinking context if created, None otherwise + """ + if detect_claude_thinking_in_stream(model_name): + return create_claude_thinking_context(agent_name, counter, model_name) + return None