RomanNum3ral
/
cai
mirror of https://github.com/aliasrobotics/cai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Streaming finished

This commit is contained in:
luijait 2025-05-11 12:54:18 +02:00
parent a5addaf2ab
commit cdefd5fe1f
3 changed files with 311 additions and 113 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
pyproject.toml
View File

@ -8,7 +8,6 @@ license = {text = "Dual-licensed MIT and Proprietary"}
authors = [{ name = "Alias Robotics", email = "research@aliasrobotics.com" }]
dependencies = [
# logging and visualization
"flask>=2.0, <3",
"folium>=0.15.0, <1",
"matplotlib>=3.0, <4",
"numpy>=1.21, <3",
@ -26,7 +25,7 @@ dependencies = [
"prompt_toolkit>=3.0.39",
"dotenv>=0.9.9",
"litellm>=1.63.7",
"mako>=1.3.9",
"mako>=1.3.8",
"mcp; python_version >= '3.10'",
"mkdocs>=1.6.0",
"mkdocs-material>=9.6.0",

7
src/cai/sdk/agents/models/openai_chatcompletions.py
View File

@ -420,6 +420,10 @@ class OpenAIChatCompletionsModel(Model):
# Only display the agent message if we haven't already shown the tool output
if should_display_message:
# Ensure we're in non-streaming mode for proper markdown parsing
previous_stream_setting = os.environ.get('CAI_STREAM', 'false')
os.environ['CAI_STREAM'] = 'false' # Force non-streaming mode for markdown parsing
# Print the agent message for CLI display
cli_print_agent_messages(
agent_name=getattr(self, 'agent_name', 'Agent'),
@ -444,6 +448,9 @@ class OpenAIChatCompletionsModel(Model):
tool_output=None, # Don't pass tool output here, we're using direct display
suppress_empty=True # Suppress empty panels
)
# Restore previous streaming setting
os.environ['CAI_STREAM'] = previous_stream_setting
# --- Add assistant tool call to message_history if present ---
# If the response contains tool_calls, add them to message_history as assistant messages

414
src/cai/util.py
View File

@ -970,28 +970,116 @@ def parse_message_content(message):
"""
Parse a message object to extract its textual content.
Only processes messages that don't have tool calls.
Detects markdown code blocks and applies syntax highlighting in non-streaming mode.
Also formats other markdown elements like headers, lists, and text formatting.
Args:
message: Can be a string or a Message object with content attribute
Returns:
str: The extracted content as a string
str or rich.console.Group: The extracted content as a string or as a rich Group with Syntax highlighting
"""
# Check if this is a duplicate print from OpenAIChatCompletionsModel
# If message is already a string, return it
from rich.console import Group
from rich.syntax import Syntax
from rich.text import Text
from rich.markdown import Markdown
import re
# Extract the raw content
raw_content = ""
# If message is already a string, use it
if isinstance(message, str):
return message
raw_content = message
# If message is a Message object with content attribute
if hasattr(message, 'content') and message.content is not None:
return message.content
elif hasattr(message, 'content') and message.content is not None:
raw_content = message.content
# If message is a dict with content key
if isinstance(message, dict) and 'content' in message:
return message['content']
elif isinstance(message, dict) and 'content' in message:
raw_content = message['content']
# If we can't extract content, convert to string
return str(message)
else:
raw_content = str(message)
# Check if streaming is enabled
streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true'
# Only apply markdown formatting in non-streaming mode
if not streaming_enabled and raw_content:
# Check if content contains markdown code blocks with improved regex
code_block_pattern = r'```(\w*)\s*([\s\S]*?)\s*```'
matches = re.findall(code_block_pattern, raw_content, re.DOTALL)
if matches:
# Prepare to process markdown with code blocks highlighted
elements = []
last_end = 0
# Find all code blocks with improved regex pattern
for match in re.finditer(r'```(\w*)\s*([\s\S]*?)\s*```', raw_content, re.DOTALL):
# Get text before the code block
start = match.start()
if start > last_end:
text_before = raw_content[last_end:start]
# Process markdown in the text before the code block
if text_before.strip():
md = Markdown(text_before)
elements.append(md)
# Process the code block
lang = match.group(1) or "text"
code = match.group(2)
# Use the language mapping helper to get proper syntax highlighting
syntax_lang = get_language_from_code_block(lang)
# Create syntax highlighted code
syntax = Syntax(
code,
syntax_lang,
theme="monokai",
line_numbers=True,
word_wrap=True,
background_color="#272822"
)
elements.append(syntax)
last_end = match.end()
# Add any remaining text after the last code block
if last_end < len(raw_content):
text_after = raw_content[last_end:]
# Process markdown in the text after the code block
if text_after.strip():
md = Markdown(text_after)
elements.append(md)
return Group(*elements)
else:
# If no code blocks, but still contains markdown, use Rich's markdown renderer
# Check for markdown elements (headers, lists, formatting)
has_markdown = any([
# Headers
re.search(r'^#{1,6}\s+\w+', raw_content, re.MULTILINE),
# Lists
re.search(r'^\s*[-*+]\s+\w+', raw_content, re.MULTILINE),
re.search(r'^\s*\d+\.\s+\w+', raw_content, re.MULTILINE),
# Bold/Italic
'**' in raw_content,
'*' in raw_content and not '**' in raw_content,
'__' in raw_content,
'_' in raw_content and not '__' in raw_content,
# Links
re.search(r'\[.+?\]\(.+?\)', raw_content)
])
if has_markdown:
return Group(Markdown(raw_content))
# For streaming mode or no markdown, return the raw content
return raw_content
def parse_message_tool_call(message, tool_output=None):
"""
@ -1152,23 +1240,33 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli
if suppress_empty and not parsed_message and not tool_panels:
return
# Check if parsed_message is empty or "null"
is_empty_message = (parsed_message == "null" or parsed_message == "" or
(isinstance(parsed_message, str) and not parsed_message.strip()))
# Also skip if the only message is "null" or empty
if parsed_message == "null" or parsed_message == "":
if is_empty_message:
if suppress_empty and not tool_panels:
return
# Check if we have Group content from markdown parsing
is_rich_content = False
from rich.console import Group
if isinstance(parsed_message, Group):
is_rich_content = True
# Special handling for Reasoner Agent
if agent_name == "Reasoner Agent":
text.append(f"[{counter}] ", style="bold red")
text.append(f"Agent: {agent_name} ", style="bold yellow")
if parsed_message:
if parsed_message and not is_rich_content:
text.append(f">> {parsed_message} ", style="green")
text.append(f"[{timestamp}", style="dim")
if model:
text.append(f" ({os.getenv('CAI_SUPPORT_MODEL')})",
style="bold blue")
text.append("]", style="dim")
elif not parsed_message:
elif is_empty_message:
# When parsed_message is empty, only include timestamp and model info
text.append(f"Agent: {agent_name} ", style="bold green")
text.append(f"[{timestamp}", style="dim")
@ -1178,7 +1276,7 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli
else:
text.append(f"[{counter}] ", style="bold cyan")
text.append(f"Agent: {agent_name} ", style="bold green")
if parsed_message:
if parsed_message and not is_rich_content:
text.append(f">> {parsed_message} ", style="yellow")
text.append(f"[{timestamp}", style="dim")
if model:
@ -1206,19 +1304,52 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli
total_cost
)
# Only append token information if there is a parsed message
if parsed_message:
if parsed_message and not is_rich_content:
text.append(tokens_text)
panel = Panel(
text,
border_style="red" if agent_name == "Reasoner Agent" else "blue",
box=ROUNDED,
padding=(0, 1),
title=("[bold]Reasoning Analysis[/bold]"
if agent_name == "Reasoner Agent"
else "[bold]Agent Interaction[/bold]"),
title_align="left"
)
# Create the panel content based on whether we have rich content or not
from rich.panel import Panel
from rich.console import Group
if is_rich_content:
# For rich content, create a Group with the header, content, and tokens
panel_content = []
panel_content.append(text)
# Add spacing between header and content for better readability
panel_content.append(Text("\n"))
# Add the Group with highlighted content
panel_content.append(parsed_message)
# Add token information at the bottom with proper spacing
if tokens_text:
panel_content.append(Text("\n"))
panel_content.append(tokens_text)
panel = Panel(
Group(*panel_content),
border_style="red" if agent_name == "Reasoner Agent" else "blue",
box=ROUNDED,
padding=(1, 1), # Increased padding for better appearance
title=("[bold]Reasoning Analysis[/bold]"
if agent_name == "Reasoner Agent"
else "[bold]Agent Interaction[/bold]"),
title_align="left"
)
else:
# For regular text content, use the original panel format
panel = Panel(
text,
border_style="red" if agent_name == "Reasoner Agent" else "blue",
box=ROUNDED,
padding=(0, 1),
title=("[bold]Reasoning Analysis[/bold]"
if agent_name == "Reasoner Agent"
else "[bold]Agent Interaction[/bold]"),
title_align="left"
)
#console.print("\n")
console.print(panel)
@ -1490,6 +1621,7 @@ def finish_agent_streaming(context, final_stats=None):
return False
def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execution_info=None, token_info=None, streaming=False):
"""
Print a tool call output to the command line.
@ -1512,6 +1644,10 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
if not output and not call_id and not streaming:
return
# Skip early for execute_code tool in non-streaming mode
if tool_name == "execute_code" and not streaming:
return
# Set up global tracker for streaming sessions
if not hasattr(cli_print_tool_output, '_streaming_sessions'):
cli_print_tool_output._streaming_sessions = {}
@ -1723,20 +1859,20 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
console.print(panel)
except ImportError:
# Fall back to simple formatting if Rich is not available
_print_simple_tool_output(tool_name, args, output, execution_info, token_info)
pass
# Helper function to create tool panel content
def _create_tool_panel_content(tool_name, args, output, execution_info=None, token_info=None):
"""Create the header and content for a tool output panel."""
from rich.text import Text
from rich.syntax import Syntax # Import Syntax for highlighting
from rich.syntax import Syntax # Import Syntax for highlighting
from rich.panel import Panel
from rich.console import Group
from rich.box import ROUNDED
# Format arguments for display
args_str = _format_tool_args(args)
# Format arguments for display, passing tool_name for specific formatting
args_str = _format_tool_args(args, tool_name=tool_name)
# Get timing information
timing_info, tool_time = _get_timing_info(execution_info)
@ -1778,73 +1914,72 @@ def _create_tool_panel_content(tool_name, args, output, execution_info=None, tok
# Determine if we need specialized content formatting
group_content = [header]
# Special handling for execute_code tool
if tool_name == "execute_code":
try:
# Parse args to get language and code
language = "python" # Default
code = ""
filename = "exploit"
if tool_name == "execute_code" and isinstance(args, dict):
command_type = args.get("command")
actual_code = args.get("code")
language = args.get("language", "python") # Default to python
# For execute command, show code and output panels
if command_type == "execute" and actual_code:
# Ensure language is a string for Syntax
language_str = get_language_from_code_block(str(language))
if isinstance(args, dict):
language = args.get("language", "python")
code = args.get("code", "")
filename = args.get("filename", "exploit")
elif isinstance(args, str):
# Try to extract language and code from args string
import re
lang_match = re.search(r"language\s*=\s*['\"]?([a-zA-Z0-9+]+)['\"]?", args)
if lang_match:
language = lang_match.group(1)
# Create syntax highlighted code panel if we have code
if code:
syntax = Syntax(code, language, theme="monokai", line_numbers=True,
background_color="#272822", indent_guides=True)
code_panel = Panel(
syntax,
title=f"Code ({language})",
border_style="cyan",
title_align="left",
box=ROUNDED,
padding=(0, 1)
)
code_syntax = Syntax(actual_code, language_str, theme="monokai",
line_numbers=True, background_color="#272822",
indent_guides=True, word_wrap=True)
code_panel = Panel(
code_syntax,
title=f"Code ({language_str})",
border_style="cyan",
title_align="left",
box=ROUNDED,
padding=(0,1)
)
group_content.extend([Text("\n"), code_panel])
# Panel for the output of the executed code
if output:
# Try to highlight output as text, or specific language if known (e.g. json)
output_lang = "text"
try:
json.loads(output) # Check if output is JSON
output_lang = "json"
except json.JSONDecodeError:
pass # Not JSON, keep as text
# Create output panel with proper highlighting
output_syntax = Syntax(output, "text", theme="monokai",
background_color="#272822")
output_syntax = Syntax(output, output_lang, theme="monokai",
background_color="#272822", word_wrap=True)
output_panel = Panel(
output_syntax,
title="Output",
border_style="green",
title_align="left",
box=ROUNDED,
padding=(0, 1)
padding=(0,1)
)
# Add code and output panels to content
group_content.append("\n")
group_content.append(code_panel)
group_content.append("\n")
group_content.append(output_panel)
# Add token info if available
if token_content:
group_content.append("\n")
group_content.append(token_content)
return header, Group(*group_content)
except Exception:
# Fallback if syntax highlighting fails
pass
# Special handling for generic_linux_command
elif tool_name == "generic_linux_command" or "command" in tool_name:
group_content.extend([Text("\n"), output_panel])
# For other commands (like cat) or if no code, just show output if any
elif output:
output_syntax = Syntax(output, "text", theme="monokai",
background_color="#272822", word_wrap=True)
output_panel = Panel(
output_syntax,
title="Output",
border_style="green",
title_align="left",
box=ROUNDED,
padding=(0,1)
)
group_content.extend([Text("\n"), output_panel])
# Special handling for generic_linux_command or any command containing 'command'
elif "command" in tool_name.lower() or "shell" in tool_name.lower():
try:
# Highlight the output as bash
output_syntax = Syntax(output, "bash", theme="monokai",
background_color="#272822")
background_color="#272822", word_wrap=True)
# Create a panel for the formatted output
output_panel = Panel(
@ -1857,34 +1992,20 @@ def _create_tool_panel_content(tool_name, args, output, execution_info=None, tok
)
# Assemble content with highlighted output
group_content.append("\n")
group_content.append(output_panel)
group_content.extend([Text("\n"), output_panel])
# Add token info if available
if token_content:
group_content.append("\n")
group_content.append(token_content)
return header, Group(*group_content)
except Exception:
# Fallback if syntax highlighting fails
pass
# Default content assembly for other tools
content = Text()
content.append(header)
content.append("\n\n")
content.append(output)
# Fallback if syntax highlighting fails, just add raw output
group_content.extend([Text("\n"), Text(output)])
# Add token info if available
if token_content:
content.append("\n\n")
content.append(token_content)
group_content.extend([Text("\n"), token_content])
return header, content
return header, Group(*group_content)
# Helper function to format tool arguments
def _format_tool_args(args):
def _format_tool_args(args, tool_name=None):
"""Format tool arguments as a clean string."""
# If args is already a string, it might be pre-formatted or a simple arg string
if isinstance(args, str):
@ -1894,7 +2015,7 @@ def _format_tool_args(args):
parsed_dict = json.loads(args)
# Recursively call with the parsed dict for consistent formatting
return _format_tool_args(parsed_dict)
return _format_tool_args(parsed_dict, tool_name=tool_name)
except json.JSONDecodeError:
# Not valid JSON, or not a dict; return as is
return args
@ -1913,11 +2034,17 @@ def _format_tool_args(args):
# Skip special flags
if key in ["async_mode", "streaming"] and not value:
continue
value_str = str(value)
# Format the value
if isinstance(value, str):
arg_parts.append(f"{key}={value}")
# Truncate long string values that are not specifically handled above
if len(value_str) > 70 and key not in ["code", "args"]:
value_str = value_str[:67] + "..."
arg_parts.append(f"{key}={value_str}")
else:
arg_parts.append(f"{key}={value}")
arg_parts.append(f"{key}={value_str}")
return ", ".join(arg_parts)
else:
return str(args)
@ -1947,7 +2074,6 @@ def _get_timing_info(execution_info=None):
timing_info.append(f"Tool: {format_time(tool_time)}")
return timing_info, tool_time
# Helper function to create token info display
def _create_token_info_display(token_info=None):
"""Create token information display text."""
@ -2279,4 +2405,70 @@ def print_message_history(messages, title="Message History"):
# Display the panel
console.print(panel)
return len(messages) # Return message count for convenience
return len(messages) # Return message count for convenience
def get_language_from_code_block(lang_identifier):
"""
Maps a language identifier from a markdown code block to a proper syntax
highlighting language name. Handles common aliases and defaults.
Args:
lang_identifier (str): Language identifier from markdown code block
Returns:
str: Proper language name for syntax highlighting
"""
# Convert to lowercase and strip whitespace
lang = lang_identifier.lower().strip() if lang_identifier else ""
# Map common language aliases to their proper names
lang_map = {
# Empty strings or unknown
"": "text",
# Python variants
"py": "python",
"python3": "python",
# JavaScript variants
"js": "javascript",
"jsx": "jsx",
"ts": "typescript",
"tsx": "tsx",
"typescript": "typescript",
# Shell variants
"sh": "bash",
"shell": "bash",
"console": "bash",
"terminal": "bash",
# Web languages
"html": "html",
"css": "css",
"json": "json",
"xml": "xml",
"yml": "yaml",
"yaml": "yaml",
# C family
"c": "c",
"cpp": "cpp",
"c++": "cpp",
"csharp": "csharp",
"cs": "csharp",
"java": "java",
# Other common languages
"go": "go",
"golang": "go",
"ruby": "ruby",
"rb": "ruby",
"rust": "rust",
"php": "php",
"sql": "sql",
"diff": "diff",
"markdown": "markdown",
"md": "markdown",
# Default fallback
"text": "text",
"plaintext": "text",
"txt": "text",
}
# Return mapped language or default to the original if not in map
return lang_map.get(lang, lang or "text")