mirror of https://github.com/aliasrobotics/cai.git
Qwen executes tool calls
This commit is contained in:
luijait
parent
383bae681d
commit
3cda833bd3
|
|
@ -8,6 +8,8 @@ import litellm
|
|||
import tiktoken
|
||||
import inspect
|
||||
import hashlib
|
||||
import re
|
||||
import asyncio
|
||||
|
||||
from collections.abc import AsyncIterator, Iterable
|
||||
from dataclasses import dataclass, field
|
||||
|
|
@ -540,7 +542,27 @@ class OpenAIChatCompletionsModel(Model):
|
|||
streaming_text_buffer = ""
|
||||
# For tool call streaming, accumulate tool_calls to add to message_history at the end
|
||||
streamed_tool_calls = []
|
||||
|
||||
# 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 = ""
|
||||
is_ollama = False
|
||||
|
||||
model_str = str(self.model).lower()
|
||||
is_ollama = self.is_ollama or "ollama" in model_str or ":" in model_str or "qwen" in model_str
|
||||
|
||||
# Add a small delay to make sure any previous tool outputs are fully rendered
|
||||
# This helps prevent overlapping of panels in the terminal
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
# Add visual separation before agent output
|
||||
if streaming_context and should_show_rich_stream:
|
||||
# If we're using rich context, we'll add separation through that
|
||||
pass
|
||||
else:
|
||||
# Print clear visual separator
|
||||
print("\n")
|
||||
|
||||
async for chunk in stream:
|
||||
if not state.started:
|
||||
state.started = True
|
||||
|
|
@ -593,6 +615,10 @@ class OpenAIChatCompletionsModel(Model):
|
|||
content = delta['content']
|
||||
|
||||
if content:
|
||||
# For Ollama, we need to accumulate the full content to check for function calls
|
||||
if is_ollama:
|
||||
ollama_full_content += content
|
||||
|
||||
# Add to the streaming text buffer
|
||||
streaming_text_buffer += content
|
||||
|
||||
|
|
@ -776,6 +802,320 @@ class OpenAIChatCompletionsModel(Model):
|
|||
streamed_tool_calls.append(tool_call_msg)
|
||||
add_to_message_history(tool_call_msg)
|
||||
|
||||
# Special handling for Ollama - check if accumulated text contains a valid function call
|
||||
if is_ollama and ollama_full_content and len(state.function_calls) == 0:
|
||||
# Look for JSON object that might be a function call
|
||||
try:
|
||||
# Try to extract a JSON object from the content
|
||||
json_start = ollama_full_content.find('{')
|
||||
json_end = ollama_full_content.rfind('}') + 1
|
||||
|
||||
logger.debug(f"Ollama content length: {len(ollama_full_content)}, JSON start: {json_start}, JSON end: {json_end}")
|
||||
|
||||
if json_start >= 0 and json_end > json_start:
|
||||
json_str = ollama_full_content[json_start:json_end]
|
||||
logger.debug(f"Extracted potential JSON: {json_str[:100]}...")
|
||||
|
||||
# Special check for generic_linux_command format
|
||||
if "generic_linux_command" in json_str and "arguments" in json_str:
|
||||
logger.debug("Detected generic_linux_command pattern - using special handling")
|
||||
|
||||
# Try regex pattern matching to handle potentially malformed JSON
|
||||
cmd_pattern = re.search(r'"command"\s*:\s*"([^"]+)"', json_str)
|
||||
args_pattern = re.search(r'"args"\s*:\s*"([^"]*)"', json_str)
|
||||
ctf_pattern = re.search(r'"ctf"\s*:\s*"([^"]*)"', json_str)
|
||||
|
||||
if cmd_pattern:
|
||||
# Create a tool call specifically for generic_linux_command
|
||||
command = cmd_pattern.group(1)
|
||||
args = args_pattern.group(1) if args_pattern else ""
|
||||
ctf = ctf_pattern.group(1) if ctf_pattern else "<CTF_ENV>"
|
||||
|
||||
tool_call_id = f"call_{hashlib.md5(('generic_linux_command' + str(time.time())).encode()).hexdigest()[:8]}"
|
||||
|
||||
# Create a proper arguments object
|
||||
command_args = {
|
||||
"command": command,
|
||||
"args": args,
|
||||
"ctf": ctf,
|
||||
"async_mode": False,
|
||||
"session_id": ""
|
||||
}
|
||||
|
||||
arguments_str = json.dumps(command_args)
|
||||
logger.debug(f"Created generic_linux_command with args: {arguments_str}")
|
||||
|
||||
# Add it to our function_calls state
|
||||
state.function_calls[0] = ResponseFunctionToolCall(
|
||||
id=FAKE_RESPONSES_ID,
|
||||
arguments=arguments_str,
|
||||
name="generic_linux_command",
|
||||
type="function_call",
|
||||
call_id=tool_call_id,
|
||||
)
|
||||
|
||||
# Display the tool call in CLI
|
||||
from cai.util import cli_print_agent_messages
|
||||
try:
|
||||
# Create a message-like object to display the function call
|
||||
tool_msg = type('ToolCallWrapper', (), {
|
||||
'content': None,
|
||||
'tool_calls': [
|
||||
type('ToolCallDetail', (), {
|
||||
'function': type('FunctionDetail', (), {
|
||||
'name': "generic_linux_command",
|
||||
'arguments': arguments_str
|
||||
}),
|
||||
'id': tool_call_id,
|
||||
'type': 'function'
|
||||
})
|
||||
]
|
||||
})
|
||||
|
||||
# Print the tool call using the CLI utility
|
||||
cli_print_agent_messages(
|
||||
agent_name=getattr(self, 'agent_name', 'Agent'),
|
||||
message=tool_msg,
|
||||
counter=getattr(self, 'interaction_counter', 0),
|
||||
model=str(self.model),
|
||||
debug=False,
|
||||
interaction_input_tokens=estimated_input_tokens,
|
||||
interaction_output_tokens=estimated_output_tokens,
|
||||
interaction_reasoning_tokens=0,
|
||||
total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens,
|
||||
total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens,
|
||||
total_reasoning_tokens=0,
|
||||
interaction_cost=None,
|
||||
total_cost=None
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error displaying tool call in CLI: {e}")
|
||||
|
||||
# Add to message history
|
||||
tool_call_msg = {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": tool_call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "generic_linux_command",
|
||||
"arguments": arguments_str
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
streamed_tool_calls.append(tool_call_msg)
|
||||
add_to_message_history(tool_call_msg)
|
||||
logger.debug(f"Added generic_linux_command with args: {arguments_str}")
|
||||
|
||||
|
||||
# Standard JSON parsing for other function calls
|
||||
parsed = json.loads(json_str)
|
||||
|
||||
# Check if it looks like a function call
|
||||
if ('name' in parsed and 'arguments' in parsed and
|
||||
isinstance(parsed['arguments'], dict)):
|
||||
|
||||
logger.debug(f"Found valid function call in Ollama output: {json_str}")
|
||||
|
||||
# Create a tool call from the JSON object
|
||||
tool_call_id = f"call_{hashlib.md5((parsed['name'] + str(time.time())).encode()).hexdigest()[:8]}"
|
||||
|
||||
# Add it to our function_calls state
|
||||
state.function_calls[0] = ResponseFunctionToolCall(
|
||||
id=FAKE_RESPONSES_ID,
|
||||
arguments=json.dumps(parsed['arguments']),
|
||||
name=parsed['name'],
|
||||
type="function_call",
|
||||
call_id=tool_call_id,
|
||||
)
|
||||
|
||||
# Ensure arguments is a valid JSON string
|
||||
arguments_str = ""
|
||||
if isinstance(parsed['arguments'], dict):
|
||||
arguments_str = json.dumps(parsed['arguments'])
|
||||
elif isinstance(parsed['arguments'], str):
|
||||
# If it's already a string, check if it's valid JSON
|
||||
try:
|
||||
# Try parsing to validate
|
||||
json.loads(parsed['arguments'])
|
||||
arguments_str = parsed['arguments']
|
||||
except:
|
||||
# If not valid JSON, encode it as a JSON string
|
||||
arguments_str = json.dumps(parsed['arguments'])
|
||||
else:
|
||||
# For any other type, convert to string and then JSON
|
||||
arguments_str = json.dumps(str(parsed['arguments']))
|
||||
|
||||
logger.debug(f"Final arguments string: {arguments_str}")
|
||||
|
||||
# Update with the properly formatted arguments
|
||||
state.function_calls[0].arguments = arguments_str
|
||||
|
||||
# Log the tool call for development purposes
|
||||
logger.debug(f"Adding tool call: {parsed['name']} with arguments: {arguments_str}")
|
||||
|
||||
# Display the tool call in CLI
|
||||
from cai.util import cli_print_agent_messages
|
||||
try:
|
||||
# Create a message-like object to display the function call
|
||||
tool_msg = type('ToolCallWrapper', (), {
|
||||
'content': None,
|
||||
'tool_calls': [
|
||||
type('ToolCallDetail', (), {
|
||||
'function': type('FunctionDetail', (), {
|
||||
'name': parsed['name'],
|
||||
'arguments': arguments_str
|
||||
}),
|
||||
'id': tool_call_id,
|
||||
'type': 'function'
|
||||
})
|
||||
]
|
||||
})
|
||||
|
||||
# Print the tool call using the CLI utility
|
||||
cli_print_agent_messages(
|
||||
agent_name=getattr(self, 'agent_name', 'Agent'),
|
||||
message=tool_msg,
|
||||
counter=getattr(self, 'interaction_counter', 0),
|
||||
model=str(self.model),
|
||||
debug=False,
|
||||
interaction_input_tokens=estimated_input_tokens,
|
||||
interaction_output_tokens=estimated_output_tokens,
|
||||
interaction_reasoning_tokens=0, # Not available for Ollama
|
||||
total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens,
|
||||
total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens,
|
||||
total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0),
|
||||
interaction_cost=None,
|
||||
total_cost=None,
|
||||
tool_output=None # Will be shown once the tool is executed
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error displaying tool call in CLI: {e}")
|
||||
|
||||
# Add to message history
|
||||
tool_call_msg = {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": tool_call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": parsed['name'],
|
||||
"arguments": arguments_str
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
streamed_tool_calls.append(tool_call_msg)
|
||||
add_to_message_history(tool_call_msg)
|
||||
|
||||
logger.debug(f"Added function call: {parsed['name']} with args: {json.dumps(parsed['arguments'])}")
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to parse potential Ollama function call: {e}")
|
||||
|
||||
# Even if JSON parsing fails, try to extract generic_linux_command with regex
|
||||
if "generic_linux_command" in ollama_full_content:
|
||||
logger.debug("JSON parsing failed but detected generic_linux_command - trying regex fallback")
|
||||
try:
|
||||
# Use regex to extract command components even from malformed output
|
||||
cmd_pattern = re.search(r'"command"\s*:\s*"([^"]+)"', ollama_full_content)
|
||||
args_pattern = re.search(r'"args"\s*:\s*"([^"]*)"', ollama_full_content)
|
||||
ctf_pattern = re.search(r'"ctf"\s*:\s*"([^"]*)"', ollama_full_content)
|
||||
|
||||
if cmd_pattern:
|
||||
# Create a tool call specifically for generic_linux_command
|
||||
command = cmd_pattern.group(1)
|
||||
args = args_pattern.group(1) if args_pattern else ""
|
||||
ctf = ctf_pattern.group(1) if ctf_pattern else "<CTF_ENV>"
|
||||
|
||||
tool_call_id = f"call_{hashlib.md5(('generic_linux_command' + str(time.time())).encode()).hexdigest()[:8]}"
|
||||
|
||||
# Create a proper arguments object
|
||||
command_args = {
|
||||
"command": command,
|
||||
"args": args,
|
||||
"ctf": ctf,
|
||||
"async_mode": False,
|
||||
"session_id": ""
|
||||
}
|
||||
|
||||
arguments_str = json.dumps(command_args)
|
||||
logger.debug(f"Fallback created generic_linux_command with args: {arguments_str}")
|
||||
|
||||
# Add it to our function_calls state
|
||||
state.function_calls[0] = ResponseFunctionToolCall(
|
||||
id=FAKE_RESPONSES_ID,
|
||||
arguments=arguments_str,
|
||||
name="generic_linux_command",
|
||||
type="function_call",
|
||||
call_id=tool_call_id,
|
||||
)
|
||||
|
||||
# Display the tool call in CLI
|
||||
from cai.util import cli_print_agent_messages
|
||||
try:
|
||||
# Create a message-like object for display
|
||||
tool_msg = type('ToolCallWrapper', (), {
|
||||
'content': None,
|
||||
'tool_calls': [
|
||||
type('ToolCallDetail', (), {
|
||||
'function': type('FunctionDetail', (), {
|
||||
'name': "generic_linux_command",
|
||||
'arguments': arguments_str
|
||||
}),
|
||||
'id': tool_call_id,
|
||||
'type': 'function'
|
||||
})
|
||||
]
|
||||
})
|
||||
|
||||
cli_print_agent_messages(
|
||||
agent_name=getattr(self, 'agent_name', 'Agent'),
|
||||
message=tool_msg,
|
||||
counter=getattr(self, 'interaction_counter', 0),
|
||||
model=str(self.model),
|
||||
debug=False,
|
||||
interaction_input_tokens=estimated_input_tokens,
|
||||
interaction_output_tokens=estimated_output_tokens,
|
||||
interaction_reasoning_tokens=0,
|
||||
total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens,
|
||||
total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens,
|
||||
total_reasoning_tokens=0,
|
||||
interaction_cost=None,
|
||||
total_cost=None
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error displaying fallback tool call in CLI: {e}")
|
||||
|
||||
# Add to message history
|
||||
tool_call_msg = {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": tool_call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "generic_linux_command",
|
||||
"arguments": arguments_str
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
streamed_tool_calls.append(tool_call_msg)
|
||||
add_to_message_history(tool_call_msg)
|
||||
|
||||
logger.debug(f"Added fallback generic_linux_command")
|
||||
except Exception as regex_err:
|
||||
logger.error(f"Regex fallback also failed: {regex_err}")
|
||||
|
||||
function_call_starting_index = 0
|
||||
if state.text_content_index_and_output:
|
||||
function_call_starting_index += 1
|
||||
|
|
@ -952,13 +1292,26 @@ class OpenAIChatCompletionsModel(Model):
|
|||
direct_stats = final_stats.copy()
|
||||
direct_stats["interaction_cost"] = float(interaction_cost)
|
||||
direct_stats["total_cost"] = float(total_cost)
|
||||
|
||||
# Add a small delay to avoid overlapping outputs
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
# Use the direct copy with guaranteed float costs
|
||||
finish_agent_streaming(streaming_context, direct_stats)
|
||||
|
||||
# Add visual separation after agent output completes
|
||||
print("\n")
|
||||
# If we're not using rich streaming and not suppressing output, use old method
|
||||
elif not self.suppress_final_output and final_response.output and any(isinstance(item, ResponseOutputMessage) for item in final_response.output):
|
||||
# Find the assistant message to print
|
||||
for item in final_response.output:
|
||||
if isinstance(item, ResponseOutputMessage) and item.role == 'assistant':
|
||||
# Add a small delay to avoid overlapping outputs
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
# Print clear visual separator before message
|
||||
print("\n")
|
||||
|
||||
cli_print_agent_messages(
|
||||
agent_name=getattr(self, 'agent_name', 'Agent'),
|
||||
message=item,
|
||||
|
|
@ -974,6 +1327,9 @@ class OpenAIChatCompletionsModel(Model):
|
|||
interaction_cost=interaction_cost,
|
||||
total_cost=total_cost,
|
||||
)
|
||||
|
||||
# Add visual separation after message
|
||||
print("\n")
|
||||
break
|
||||
|
||||
# --- Add assistant tool call(s) to message_history at the end of streaming ---
|
||||
|
|
@ -1194,6 +1550,7 @@ class OpenAIChatCompletionsModel(Model):
|
|||
# Use the specialized Qwen approach first
|
||||
return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls)
|
||||
except Exception as qwen_e:
|
||||
print(qwen_e)
|
||||
# If that fails, try our direct OpenAI approach
|
||||
qwen_params = kwargs.copy()
|
||||
qwen_params["api_base"] = get_ollama_api_base()
|
||||
|
|
@ -1351,117 +1708,80 @@ class OpenAIChatCompletionsModel(Model):
|
|||
# Standard OpenAI handling for non-streaming
|
||||
ret = litellm.completion(**kwargs)
|
||||
return ret
|
||||
|
||||
|
||||
async def _fetch_response_litellm_ollama(
|
||||
self,
|
||||
kwargs: dict,
|
||||
model_settings: ModelSettings,
|
||||
tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven,
|
||||
stream: bool,
|
||||
parallel_tool_calls: bool
|
||||
parallel_tool_calls: bool,
|
||||
provider="ollama"
|
||||
) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]:
|
||||
# Extract only supported parameters for Ollama
|
||||
ollama_supported_params = {
|
||||
"model": kwargs.get("model", ""),
|
||||
"messages": kwargs.get("messages", []),
|
||||
"stream": kwargs.get("stream", False)
|
||||
}
|
||||
|
||||
# Safely add optional parameters only if they exist in kwargs
|
||||
if "temperature" in kwargs:
|
||||
ollama_supported_params["temperature"] = kwargs["temperature"] if kwargs["temperature"] is not NOT_GIVEN else None
|
||||
if "top_p" in kwargs:
|
||||
ollama_supported_params["top_p"] = kwargs["top_p"] if kwargs["top_p"] is not NOT_GIVEN else None
|
||||
if "max_tokens" in kwargs:
|
||||
ollama_supported_params["max_tokens"] = kwargs["max_tokens"] if kwargs["max_tokens"] is not NOT_GIVEN else None
|
||||
|
||||
# Add stream parameter - default to False if not present
|
||||
ollama_supported_params["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"]
|
||||
|
||||
# IMPORTANT: For tool calls with Ollama (especially with Qwen),
|
||||
# we need to pass the tools and tool_choice as part of the request
|
||||
# This is needed for both streaming and non-streaming modes
|
||||
# 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")
|
||||
|
||||
# Include tool_choice if present and not NOT_GIVEN
|
||||
if "tool_choice" in kwargs and kwargs.get("tool_choice") is not NOT_GIVEN:
|
||||
ollama_supported_params["tool_choice"] = kwargs.get("tool_choice")
|
||||
|
||||
# Modify the messages to remove system message for Ollama
|
||||
if (ollama_supported_params["messages"] and
|
||||
len(ollama_supported_params["messages"]) > 0 and
|
||||
ollama_supported_params["messages"][0].get("role") == "system"):
|
||||
# Extract the system message
|
||||
system_content = ollama_supported_params["messages"][0].get("content", "")
|
||||
# Remove it from the messages
|
||||
ollama_supported_params["messages"] = ollama_supported_params["messages"][1:]
|
||||
# If there are user messages, prepend system to first user
|
||||
if (ollama_supported_params["messages"] and
|
||||
len(ollama_supported_params["messages"]) > 0 and
|
||||
ollama_supported_params["messages"][0].get("role") == "user"):
|
||||
# Prepend the system instruction to the first user message, with a separator
|
||||
user_content = ollama_supported_params["messages"][0].get("content", "")
|
||||
if isinstance(user_content, str):
|
||||
ollama_supported_params["messages"][0]["content"] = f"System: {system_content}\n\nUser: {user_content}"
|
||||
|
||||
# Remove None values
|
||||
ollama_kwargs = {k: v for k, v in ollama_supported_params.items() if v is not None}
|
||||
|
||||
# Detect if this is a Qwen model
|
||||
# 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:
|
||||
# For streaming with Ollama, we need to create a Response object first
|
||||
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),
|
||||
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,
|
||||
)
|
||||
|
||||
# Special handling for Qwen when streaming to ensure tool calls are properly handled
|
||||
if is_qwen and os.getenv('CAI_STREAM', 'false').lower() == 'true':
|
||||
# Make sure we have the right custom provider in streaming mode
|
||||
stream_obj = await litellm.acompletion(
|
||||
**ollama_kwargs,
|
||||
api_base=get_ollama_api_base(),
|
||||
custom_llm_provider="openai" # Use openai provider for best compatibility with Qwen
|
||||
)
|
||||
else:
|
||||
# Standard Ollama streaming
|
||||
stream_obj = await litellm.acompletion(
|
||||
**ollama_kwargs,
|
||||
api_base=get_ollama_api_base().rstrip('/v1'),
|
||||
custom_llm_provider="ollama"
|
||||
)
|
||||
# Get streaming response
|
||||
stream_obj = await litellm.acompletion(
|
||||
**ollama_kwargs,
|
||||
api_base=api_base,
|
||||
custom_llm_provider=provider,
|
||||
)
|
||||
return response, stream_obj
|
||||
else:
|
||||
# Non-streaming mode
|
||||
# For Qwen models, use the openai provider when tools are present
|
||||
if is_qwen and "tools" in ollama_kwargs:
|
||||
ret = litellm.completion(
|
||||
**ollama_kwargs,
|
||||
api_base=get_ollama_api_base(),
|
||||
custom_llm_provider="openai" # Use openai provider for better tool handling
|
||||
)
|
||||
else:
|
||||
# Standard Ollama completion
|
||||
ret = litellm.completion(
|
||||
**ollama_kwargs,
|
||||
api_base=get_ollama_api_base().rstrip('/v1'),
|
||||
custom_llm_provider="ollama"
|
||||
)
|
||||
return ret
|
||||
|
||||
|
||||
# Get completion response
|
||||
return litellm.completion(
|
||||
**ollama_kwargs,
|
||||
api_base=api_base,
|
||||
custom_llm_provider=provider,
|
||||
)
|
||||
|
||||
def _get_client(self) -> AsyncOpenAI:
|
||||
if self._client is None:
|
||||
|
|
@ -1494,6 +1814,33 @@ class OpenAIChatCompletionsModel(Model):
|
|||
'arguments': function_call.get('arguments', '')
|
||||
}
|
||||
}]
|
||||
|
||||
# Handle special Ollama generic_linux_command format
|
||||
if isinstance(delta, dict) and 'content' in delta:
|
||||
content = delta['content']
|
||||
# Try to detect if the content is a JSON string with function call format
|
||||
try:
|
||||
if isinstance(content, str) and '{' in content and '}' in content:
|
||||
# Try to extract JSON from the content (it might be embedded in text)
|
||||
json_start = content.find('{')
|
||||
json_end = content.rfind('}') + 1
|
||||
if json_start >= 0 and json_end > json_start:
|
||||
json_str = content[json_start:json_end]
|
||||
parsed = json.loads(json_str)
|
||||
if 'name' in parsed and 'arguments' in parsed:
|
||||
# This looks like a function call in JSON format
|
||||
return [{
|
||||
'index': 0,
|
||||
'id': f"call_{time.time_ns()}", # Generate a unique ID
|
||||
'type': 'function',
|
||||
'function': {
|
||||
'name': parsed['name'],
|
||||
'arguments': json.dumps(parsed['arguments']) if isinstance(parsed['arguments'], dict) else parsed['arguments']
|
||||
}
|
||||
}]
|
||||
except Exception:
|
||||
# If JSON parsing fails, just continue with normal processing
|
||||
pass
|
||||
|
||||
# Anthropic-style tool_use format
|
||||
if hasattr(delta, 'tool_use') and delta.tool_use:
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import asyncio
|
|||
import copy
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, cast
|
||||
import os
|
||||
|
||||
from openai.types.responses import ResponseCompletedEvent
|
||||
|
||||
|
|
@ -778,6 +779,23 @@ class Runner:
|
|||
run_config: RunConfig,
|
||||
tool_use_tracker: AgentToolUseTracker,
|
||||
) -> SingleStepResult:
|
||||
# Log the raw model response output, focusing on tool calls
|
||||
logger.debug(f"[_get_single_step_result_from_response] Raw new_response.id: {new_response.referenceable_id}")
|
||||
if new_response.output:
|
||||
for i, item in enumerate(new_response.output):
|
||||
item_type = type(item).__name__
|
||||
logger.debug(f"[_get_single_step_result_from_response] Raw output item [{i}] type: {item_type}")
|
||||
if hasattr(item, "name") and hasattr(item, "arguments"): # For ResponseFunctionToolCall
|
||||
logger.debug(
|
||||
f"[_get_single_step_result_from_response] Raw ResponseFunctionToolCall item [{i}]: "
|
||||
f"Name='{getattr(item, 'name', 'N/A')}', "
|
||||
f"Args='{getattr(item, 'arguments', 'N/A')}', "
|
||||
f"CallID='{getattr(item, 'call_id', 'N/A')}'"
|
||||
)
|
||||
elif hasattr(item, "text"): # For ResponseOutputText
|
||||
logger.debug(f"[_get_single_step_result_from_response] Raw ResponseOutputText item [{i}]: Text='{getattr(item, 'text', '')[:100]}...'")
|
||||
|
||||
|
||||
processed_response = RunImpl.process_model_response(
|
||||
agent=agent,
|
||||
all_tools=all_tools,
|
||||
|
|
@ -786,6 +804,69 @@ class Runner:
|
|||
handoffs=handoffs,
|
||||
)
|
||||
|
||||
# Log the processed response, focusing on tools_used
|
||||
logger.debug(f"[_get_single_step_result_from_response] Processed response type: {type(processed_response).__name__}")
|
||||
if hasattr(processed_response, 'is_final_output'):
|
||||
logger.debug(f"[_get_single_step_result_from_response] Processed response: is_final_output={processed_response.is_final_output}")
|
||||
else:
|
||||
logger.debug(f"[_get_single_step_result_from_response] Processed response does not have is_final_output attribute")
|
||||
|
||||
if hasattr(processed_response, 'final_output_from_llm') and processed_response.final_output_from_llm is not None:
|
||||
logger.debug(f"[_get_single_step_result_from_response] Processed response: final_output_from_llm='{str(processed_response.final_output_from_llm)[:100]}...'")
|
||||
|
||||
# Log tools used with robust type checking
|
||||
if hasattr(processed_response, 'tools_used') and processed_response.tools_used:
|
||||
# Log summarizing number of tools used first
|
||||
logger.debug(f"[_get_single_step_result_from_response] Found {len(processed_response.tools_used)} tools used")
|
||||
|
||||
# Add spacing between blocks in terminal output
|
||||
if os.environ.get('CAI_STREAM', 'false').lower() == 'true':
|
||||
print("") # Visual separator only when streaming is enabled
|
||||
|
||||
for i, tool_call in enumerate(processed_response.tools_used):
|
||||
try:
|
||||
# Safely extract tool name with multiple fallbacks
|
||||
tool_name = "Unknown"
|
||||
try:
|
||||
if hasattr(tool_call, 'tool'):
|
||||
if isinstance(tool_call.tool, str):
|
||||
tool_name = tool_call.tool
|
||||
elif hasattr(tool_call.tool, 'name'):
|
||||
tool_name = tool_call.tool.name
|
||||
else:
|
||||
tool_name = str(tool_call.tool)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Safely extract call_id
|
||||
call_id = "Unknown"
|
||||
try:
|
||||
if hasattr(tool_call, 'call_id'):
|
||||
call_id = str(tool_call.call_id)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Safely extract parsed_args
|
||||
parsed_args = "Unknown"
|
||||
try:
|
||||
if hasattr(tool_call, 'parsed_args'):
|
||||
parsed_args = str(tool_call.parsed_args)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.debug(
|
||||
f"[_get_single_step_result_from_response] Processed tool_call [{i}]: "
|
||||
f"Name='{tool_name}', "
|
||||
f"CallID='{call_id}', "
|
||||
f"ParsedArgs='{parsed_args}'"
|
||||
)
|
||||
except Exception as e:
|
||||
# Last resort fallback for any unexpected structure
|
||||
logger.debug(f"[_get_single_step_result_from_response] Could not process tool_call [{i}]: {str(e)}")
|
||||
else:
|
||||
logger.debug("[_get_single_step_result_from_response] Processed response: No tools_used.")
|
||||
|
||||
|
||||
tool_use_tracker.add_tool_use(agent, processed_response.tools_used)
|
||||
|
||||
return await RunImpl.execute_tools_and_side_effects(
|
||||
|
|
|
|||
|
|
@ -222,18 +222,41 @@ def terminate_session(session_id):
|
|||
|
||||
def _run_ctf(ctf, command, stdout=False, timeout=100, stream=False, call_id=None):
|
||||
try:
|
||||
# Ensure the command is executed in a shell that supports command
|
||||
# chaining
|
||||
output = ctf.get_shell(command, timeout=timeout)
|
||||
# exploit_logger.log_ok()
|
||||
|
||||
if stdout:
|
||||
print("\033[32m" + output + "\033[0m")
|
||||
return output # output if output else result.stder
|
||||
# Check the type of ctf object to handle various formats properly
|
||||
if ctf is None:
|
||||
return f"Error: CTF environment is None"
|
||||
|
||||
# Handle string CTF values (like "CTF_ENV")
|
||||
if isinstance(ctf, str):
|
||||
# If in local mode, fallback to running local command
|
||||
return _run_local(command, stdout, timeout, stream, call_id)
|
||||
|
||||
# Handle dict format that might come from Ollama/Qwen models
|
||||
if isinstance(ctf, dict):
|
||||
# Check if this dict has a get_shell key or function
|
||||
if callable(getattr(ctf, 'get_shell', None)):
|
||||
# This is a proper CTF object
|
||||
output = ctf.get_shell(command, timeout=timeout)
|
||||
if stdout:
|
||||
print("\033[32m" + output + "\033[0m")
|
||||
return output
|
||||
else:
|
||||
# This is a dict but not a proper CTF object, fallback to local
|
||||
return _run_local(command, stdout, timeout, stream, call_id)
|
||||
|
||||
# Original code path for proper CTF objects
|
||||
if hasattr(ctf, 'get_shell') and callable(ctf.get_shell):
|
||||
output = ctf.get_shell(command, timeout=timeout)
|
||||
if stdout:
|
||||
print("\033[32m" + output + "\033[0m")
|
||||
return output
|
||||
|
||||
# Fallback for any other case
|
||||
return _run_local(command, stdout, timeout, stream, call_id)
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
print(color(f"Error executing CTF command: {e}", fg="red"))
|
||||
# exploit_logger.log_error(str(e))
|
||||
return f"Error executing CTF command: {str(e)}"
|
||||
# Fallback to local execution
|
||||
return _run_local(command, stdout, timeout, stream, call_id)
|
||||
|
||||
|
||||
def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None):
|
||||
|
|
@ -328,7 +351,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
|
|||
header.append(")", style="yellow")
|
||||
tool_time = 0
|
||||
start_time = time.time()
|
||||
total_time = time.time() - START_TIME
|
||||
total_time = time.time() - START_TIME if START_TIME else 0
|
||||
timing_info = []
|
||||
if total_time:
|
||||
timing_info.append(f"Total: {format_time(total_time)}")
|
||||
|
|
@ -339,15 +362,24 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
|
|||
|
||||
content = Text()
|
||||
|
||||
# Get console width and calculate appropriate panel width
|
||||
console_width = console.width if hasattr(console, 'width') else 100
|
||||
panel_width = min(int(console_width * 0.95), console_width - 4) # Use 95% of width or leave 4 chars margin
|
||||
|
||||
# Create the panel with a specific width to avoid overflow
|
||||
panel = Panel(
|
||||
Text.assemble(header, "\n\n", content),
|
||||
title="[bold green]Tool Execution[/bold green]",
|
||||
subtitle="[bold green]Live Output[/bold green]",
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
box=ROUNDED
|
||||
box=ROUNDED,
|
||||
width=panel_width
|
||||
)
|
||||
|
||||
# Print clear separator before panel to avoid overlap with previous content
|
||||
console.print("\n\n")
|
||||
|
||||
# Start Live display
|
||||
with Live(panel, console=console, refresh_per_second=4) as live:
|
||||
# Stream stdout in real-time
|
||||
|
|
@ -364,7 +396,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
|
|||
|
||||
# Update tool_time and header with new timing info
|
||||
tool_time = time.time() - start_time
|
||||
total_time = time.time() - START_TIME
|
||||
total_time = time.time() - START_TIME if START_TIME else 0
|
||||
# Remove any previous timing info from header (rebuild header)
|
||||
timing_info = []
|
||||
if total_time:
|
||||
|
|
@ -386,7 +418,8 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
|
|||
subtitle="[bold green]Live Output[/bold green]",
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
box=ROUNDED
|
||||
box=ROUNDED,
|
||||
width=panel_width
|
||||
)
|
||||
live.update(panel)
|
||||
# Check if process is done
|
||||
|
|
@ -405,7 +438,8 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
|
|||
subtitle="[bold green]Live Output[/bold green]",
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
box=ROUNDED
|
||||
box=ROUNDED,
|
||||
width=panel_width
|
||||
)
|
||||
live.update(panel)
|
||||
|
||||
|
|
@ -415,12 +449,16 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
|
|||
title="[bold green]Tool Execution[/bold green]",
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
box=ROUNDED
|
||||
box=ROUNDED,
|
||||
width=panel_width
|
||||
)
|
||||
live.update(panel)
|
||||
|
||||
# Wait a moment for the panel to be displayed properly
|
||||
time.sleep(0.5)
|
||||
|
||||
# Print clear separator after panel to avoid overlap with next content
|
||||
console.print("\n\n")
|
||||
else:
|
||||
# Fallback to simpler streaming with cli_print_tool_output
|
||||
# Parse command into command and args (same as rich mode)
|
||||
|
|
@ -436,6 +474,9 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
|
|||
tool_args["args"] = args
|
||||
# Note: Omitted empty values and async_mode=False as it's default
|
||||
|
||||
# Print separator to avoid overlap with previous content
|
||||
print("\n")
|
||||
|
||||
# Initial notification - just once
|
||||
cli_print_tool_output(tool_name, tool_args, "Command started...", call_id=call_id)
|
||||
|
||||
|
|
@ -473,6 +514,9 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
|
|||
final_output += f"\nCommand exited with code {return_code}"
|
||||
|
||||
cli_print_tool_output(tool_name, tool_args, final_output, call_id=call_id)
|
||||
|
||||
# Print separator to avoid overlap with next content
|
||||
print("\n")
|
||||
|
||||
# Return the full output
|
||||
return ''.join(output_buffer)
|
||||
|
|
@ -547,7 +591,22 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
if stream and not call_id:
|
||||
call_id = str(uuid.uuid4())[:8]
|
||||
|
||||
# Otherwise, run command normally
|
||||
if ctf:
|
||||
return _run_ctf(ctf, command, stdout, timeout, stream, call_id)
|
||||
return _run_local(command, stdout, timeout, stream, call_id, tool_name)
|
||||
# Determine whether to use CTF or local execution
|
||||
use_ctf = False
|
||||
if ctf is not None:
|
||||
# Check if ctf is a proper object that can handle shell commands
|
||||
if (hasattr(ctf, 'get_shell') and callable(ctf.get_shell)) or isinstance(ctf, dict) or isinstance(ctf, str):
|
||||
use_ctf = True
|
||||
|
||||
# Run the command using the appropriate method
|
||||
if use_ctf:
|
||||
try:
|
||||
# Try with CTF first
|
||||
return _run_ctf(ctf, command, stdout, timeout, stream, call_id)
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
# Fallback to local if CTF fails
|
||||
print(color(f"CTF execution failed, falling back to local: {str(e)}", fg="yellow"))
|
||||
return _run_local(command, stdout, timeout, stream, call_id, tool_name)
|
||||
else:
|
||||
# Use local execution
|
||||
return _run_local(command, stdout, timeout, stream, call_id, tool_name)
|
||||
|
|
|
|||
Loading…
Reference in New Issue