diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 0134fcbf..c5138d14 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -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 "" + + 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 "" + + 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: diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py index 84bb664b..a0937d28 100644 --- a/src/cai/sdk/agents/run.py +++ b/src/cai/sdk/agents/run.py @@ -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( diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index e1e58778..95b0f223 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -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)