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

Merge branch 'nostream_ctrlc' into '0.4.0' 

Nostream ctrlc

See merge request aliasrobotics/alias_research/cai!180
This commit is contained in:
Luis Javier Navarrete Lozano 2025-05-19 16:59:43 +00:00
parent 770f6b0ac0 3d3be8c7df
commit ff0bea80ca
6 changed files with 160 additions and 65 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/cai/cli.py
View File

@ -702,7 +702,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf')):
start_idle_timer()
except KeyboardInterrupt:
# No need to clean up streaming context as model handles it
print("\n\033[91mKeyboard interrupt detected\033[0m")
pass
except Exception as e:
import traceback

8
src/cai/sdk/agents/function_schema.py
View File

@ -263,7 +263,13 @@ def function_schema(
# field_name -> (type_annotation, default_value_or_Field(...))
fields: dict[str, Any] = {}
for name, param in filtered_params:
filtered_params_no_ctf = [
(name, param)
for name, param in filtered_params
if name.lower() != 'ctf'
]
for name, param in filtered_params_no_ctf:
ann = type_hints.get(name, param.annotation)
default = param.default

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

@ -343,7 +343,7 @@ class OpenAIChatCompletionsModel(Model):
from cai.util import fix_message_list
converted_messages = fix_message_list(converted_messages)
except Exception as e:
logger.warning(f"Failed to fix message list: {e}")
pass
# Get token count estimate before API call for consistent counting
estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages)
@ -1067,7 +1067,7 @@ class OpenAIChatCompletionsModel(Model):
arguments=arguments_str,
name=parsed['name'],
type="function_call",
call_id=tool_call_id,
call_id=tool_call_id[:40],
)
# Display the tool call in CLI
@ -1089,7 +1089,7 @@ class OpenAIChatCompletionsModel(Model):
'name': parsed['name'],
'arguments': arguments_str
}),
'id': tool_call_id,
'id': tool_call_id[:40],
'type': 'function'
})
]
@ -1171,7 +1171,7 @@ class OpenAIChatCompletionsModel(Model):
yield ResponseOutputItemAddedEvent(
item=ResponseFunctionToolCall(
id=FAKE_RESPONSES_ID,
call_id=function_call.call_id,
call_id=function_call.call_id[:40],
arguments=function_call.arguments,
name=function_call.name,
type="function_call",
@ -1190,7 +1190,7 @@ class OpenAIChatCompletionsModel(Model):
yield ResponseOutputItemDoneEvent(
item=ResponseFunctionToolCall(
id=FAKE_RESPONSES_ID,
call_id=function_call.call_id,
call_id=function_call.call_id[:40],
arguments=function_call.arguments,
name=function_call.name,
type="function_call",
@ -1375,6 +1375,7 @@ class OpenAIChatCompletionsModel(Model):
self.total_cost
)
# Stop active timer and start idle timer when streaming is complete
stop_active_timer()
start_idle_timer()
@ -1462,7 +1463,7 @@ class OpenAIChatCompletionsModel(Model):
if new_length != prev_length:
logger.debug(f"Message list was fixed: {prev_length} -> {new_length} messages")
except Exception as e:
logger.warning(f"Failed to fix message list: {e}")
pass
parallel_tool_calls = (
True if model_settings.parallel_tool_calls and tools and len(tools) > 0 else NOT_GIVEN
@ -1588,7 +1589,7 @@ 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:
# print(color("BadRequestError encountered: " + str(e), fg="yellow"))
#print(color("BadRequestError encountered: " + str(e), fg="yellow"))
if "LLM Provider NOT provided" in str(e):
model_str = str(self.model).lower()
provider = None
@ -1701,7 +1702,7 @@ class OpenAIChatCompletionsModel(Model):
kwargs["messages"] = fixed_messages
except Exception as fix_error:
print(f"Failed to fix message sequence: {fix_error}")
pass
return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls)
@ -1775,29 +1776,92 @@ class OpenAIChatCompletionsModel(Model):
stream: bool,
parallel_tool_calls: bool
) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]:
"""Handle standard LiteLLM API calls for OpenAI and compatible models."""
if stream:
# Standard LiteLLM handling for streaming
ret = litellm.completion(**kwargs)
stream_obj = await litellm.acompletion(**kwargs)
"""
Handle standard LiteLLM API calls for OpenAI and compatible models.
If a ContextWindowExceededError occurs due to a tool_call id being
too long, truncate all tool_call ids in the messages to 40 characters
and retry once silently.
"""
try:
if stream:
# Standard LiteLLM handling for streaming
ret = litellm.completion(**kwargs)
stream_obj = await litellm.acompletion(**kwargs)
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,
)
return response, stream_obj
else:
# Standard OpenAI handling for non-streaming
ret = litellm.completion(**kwargs)
return ret
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,
)
return response, stream_obj
else:
# Standard OpenAI handling for non-streaming
ret = litellm.completion(**kwargs)
return ret
except Exception as e:
error_msg = str(e)
# Handle both OpenAI and Anthropic error messages for tool_call_id
if (
"string too long" in error_msg
or "Invalid 'messages" in error_msg
and "tool_call_id" in error_msg
and "maximum length" in error_msg
):
# Truncate all tool_call ids in all messages to 40 characters
messages = kwargs.get("messages", [])
for msg in messages:
# Truncate tool_call_id in the message itself if present
if (
"tool_call_id" in msg
and isinstance(msg["tool_call_id"], str)
and len(msg["tool_call_id"]) > 40
):
msg["tool_call_id"] = msg["tool_call_id"][:40]
# Truncate tool_call ids in tool_calls if present
if "tool_calls" in msg and isinstance(msg["tool_calls"], list):
for tool_call in msg["tool_calls"]:
if (
isinstance(tool_call, dict)
and "id" in tool_call
and isinstance(tool_call["id"], str)
and len(tool_call["id"]) > 40
):
tool_call["id"] = tool_call["id"][:40]
kwargs["messages"] = messages
# Retry once, silently
if stream:
ret = litellm.completion(**kwargs)
stream_obj = await litellm.acompletion(**kwargs)
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,
)
return response, stream_obj
else:
ret = litellm.completion(**kwargs)
return ret
else:
raise
async def _fetch_response_litellm_ollama(
self,
@ -1808,40 +1872,60 @@ class OpenAIChatCompletionsModel(Model):
parallel_tool_calls: bool,
provider="ollama"
) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]:
"""
Fetches a response from an Ollama or Qwen model using LiteLLM, ensuring
that the 'format' parameter is not set to a JSON string, which can cause
issues with the Ollama API.
Args:
kwargs (dict): Parameters for the completion request.
model_settings (ModelSettings): Model configuration.
tool_choice (ChatCompletionToolChoiceOptionParam | NotGiven): Tool choice.
stream (bool): Whether to stream the response.
parallel_tool_calls (bool): Whether to allow parallel tool calls.
provider (str): Provider name, defaults to "ollama".
Returns:
ChatCompletion or tuple[Response, AsyncStream[ChatCompletionChunk]]:
The completion response or a tuple for streaming.
"""
# Extract only supported parameters for Ollama
ollama_supported_params = {
"model": kwargs.get("model", ""),
"messages": kwargs.get("messages", []),
"stream": kwargs.get("stream", False)
}
# Add optional parameters if they exist and are not NOT_GIVEN
for param in ["temperature", "top_p", "max_tokens"]:
if param in kwargs and kwargs[param] is not NOT_GIVEN:
ollama_supported_params[param] = kwargs[param]
# Add extra headers if available
if "extra_headers" in kwargs:
ollama_supported_params["extra_headers"] = kwargs["extra_headers"]
# Add tools and tool_choice for compatibility with Qwen
if "tools" in kwargs and kwargs.get("tools") and kwargs.get("tools") is not NOT_GIVEN:
ollama_supported_params["tools"] = kwargs.get("tools")
if "tool_choice" in kwargs and kwargs.get("tool_choice") is not NOT_GIVEN:
ollama_supported_params["tool_choice"] = kwargs.get("tool_choice")
# Remove None values
ollama_kwargs = {k: v for k, v in ollama_supported_params.items() if v is not None}
# Add tools for compatibility with Qwen
if (
"tools" in kwargs
and kwargs.get("tools")
and kwargs.get("tools") is not NOT_GIVEN
):
ollama_supported_params["tools"] = kwargs.get("tools")
# Remove None values and filter out 'response_format'
ollama_kwargs = {
k: v for k, v in ollama_supported_params.items()
if v is not None and k != "response_format"
}
# Check if this is a Qwen model
model_str = str(self.model).lower()
is_qwen = "qwen" in model_str
api_base = get_ollama_api_base()
if "ollama" in provider:
api_base = api_base.rstrip('/v1')
# Create response object for streaming
if stream:
response = Response(
id=FAKE_RESPONSES_ID,
@ -1849,8 +1933,9 @@ class OpenAIChatCompletionsModel(Model):
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),
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=[],
@ -1864,8 +1949,6 @@ class OpenAIChatCompletionsModel(Model):
)
return response, stream_obj
else:
# Get completion response
return litellm.completion(
**ollama_kwargs,
@ -2035,7 +2118,7 @@ class _Converter:
items.append(
ResponseFunctionToolCall(
id=FAKE_RESPONSES_ID,
call_id=tool_call.id,
call_id=tool_call.id[:40],
arguments=tool_call.function.arguments,
name=tool_call.function.name,
type="function_call",
@ -2254,7 +2337,7 @@ class _Converter:
tool_calls_param.append(
ChatCompletionMessageToolCallParam(
id=tc.get("id", ""),
id=tc.get("id", "")[:40],
type=tc.get("type", "function"),
function={
"name": function_details.get("name", "unknown_function"),
@ -2366,7 +2449,7 @@ class _Converter:
asst = ensure_assistant_message()
tool_calls = list(asst.get("tool_calls", []))
new_tool_call = ChatCompletionMessageToolCallParam(
id=file_search["id"],
id=file_search["id"][:40],
type="function",
function={
"name": "file_search_call",
@ -2409,7 +2492,7 @@ class _Converter:
arguments = json.dumps(arguments)
new_tool_call = ChatCompletionMessageToolCallParam(
id=func_call["call_id"],
id=func_call["call_id"][:40],
type="function",
function={
"name": func_call["name"],

5
src/cai/sdk/agents/tool.py
View File

@ -272,6 +272,11 @@ def function_tool(
async def _on_invoke_tool(ctx: RunContextWrapper[Any], input: str) -> Any:
try:
return await _on_invoke_tool_impl(ctx, input)
except KeyboardInterrupt:
logger.info(
f"Tool {schema.name} execution was interrupted by the user (Ctrl+C)."
)
return "execution was interrupted by the user (Ctrl+C)"
except Exception as e:
if failure_error_function is None:
raise

18
src/cai/tools/common.py
View File

@ -127,7 +127,7 @@ class ShellSession: # pylint: disable=too-many-instance-attributes
except Exception as e:
self.output_buffer.append(f"Error starting container session: {str(e)}")
self.is_running = False
return
return str(e)
# --- Start in CTF ---
if self.ctf:
@ -141,8 +141,8 @@ class ShellSession: # pylint: disable=too-many-instance-attributes
self.output_buffer.append(output)
except Exception as e: # pylint: disable=broad-except
self.output_buffer.append(f"Error executing CTF command: {str(e)}")
self.is_running = False
return
self.is_running = False
return str(e)
# --- Start Locally (Host) ---
try:
@ -167,7 +167,7 @@ class ShellSession: # pylint: disable=too-many-instance-attributes
except Exception as e: # pylint: disable=broad-except
self.output_buffer.append(f"Error starting local session: {str(e)}")
self.is_running = False
return str(e)
def _read_output(self):
"""Read output from the process"""
try:
@ -195,6 +195,7 @@ class ShellSession: # pylint: disable=too-many-instance-attributes
except Exception as e:
self.output_buffer.append(f"Error in read_output loop: {str(e)}")
self.is_running = False
return str(e)
def is_process_running(self):
@ -645,6 +646,8 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
if stdout:
print("\033[32m" + error_msg + "\033[0m")
return error_msg
return error_msg
except Exception as e: # pylint: disable=broad-except
@ -705,6 +708,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
Returns:
str: Command output, status message, or session ID.
"""
if ctf and not hasattr(ctf, "get_shell"):
ctf = None
# Use the active timer during tool execution
stop_idle_timer()
start_active_timer()
@ -1256,13 +1261,10 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
custom_args=args
)
# Switch back to idle mode after local command completes
stop_active_timer()
start_idle_timer()
return result
except Exception as e:
# Ensure we switch back to idle mode if any unexpected error occurs
stop_active_timer()
start_idle_timer()
raise e
raise

5
src/cai/util.py
View File

@ -844,7 +844,6 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912
i += 2
else:
i += 1
return sanitized_messages
def cli_print_tool_call(tool_name="", args="", output="", prefix=" "):
@ -2103,8 +2102,8 @@ def print_message_history(messages, title="Message History"):
table = Table(show_header=True, header_style="bold magenta", expand=True)
table.add_column("#", style="dim", width=3)
table.add_column("Role", style="cyan", width=10)
table.add_column("Content", width=40)
table.add_column("Metadata", width=30)
table.add_column("Content", width=1000)
table.add_column("Metadata", width=1000)
# Process each message
for i, msg in enumerate(messages):