From 7aabfe856f36ecdef395ea31e47af8d35409aa0f Mon Sep 17 00:00:00 2001 From: lidia9 Date: Mon, 19 May 2025 10:42:11 +0200 Subject: [PATCH 01/11] FIX CAI_PRICE_LIMIT --- src/cai/repl/ui/toolbar.py | 2 ++ src/cai/sdk/agents/exceptions.py | 6 ++++++ src/cai/sdk/agents/run.py | 9 +++++++++ src/cai/util.py | 20 +++++++++++++------- 4 files changed, 30 insertions(+), 7 deletions(-) diff --git a/src/cai/repl/ui/toolbar.py b/src/cai/repl/ui/toolbar.py index eb44d9a0..6f3edde6 100644 --- a/src/cai/repl/ui/toolbar.py +++ b/src/cai/repl/ui/toolbar.py @@ -123,6 +123,8 @@ def update_toolbar_in_background(): os.getenv('CAI_MODEL', 'default')} | " f"Max Turns: { os.getenv('CAI_MAX_TURNS', 'inf')} | " + f"Price Limit: { + os.getenv('CAI_PRICE_LIMIT', 'inf')} | " f"{current_time_with_tz}" ) toolbar_cache['last_update'] = datetime.datetime.now() diff --git a/src/cai/sdk/agents/exceptions.py b/src/cai/sdk/agents/exceptions.py index 820390db..170bdd76 100644 --- a/src/cai/sdk/agents/exceptions.py +++ b/src/cai/sdk/agents/exceptions.py @@ -61,3 +61,9 @@ class OutputGuardrailTripwireTriggered(AgentsException): super().__init__( f"Guardrail {guardrail_result.guardrail.__class__.__name__} triggered tripwire" ) + + +class PriceLimitExceeded(AgentsException): + """Raised when the maximum price limit is exceeded.""" + def __init__(self, current_cost: float, price_limit: float): + super().__init__(f"Maximum price limit (${price_limit:.4f}) exceeded. Current cost: ${current_cost:.4f}") diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py index 6089d331..15d6eb70 100644 --- a/src/cai/sdk/agents/run.py +++ b/src/cai/sdk/agents/run.py @@ -58,6 +58,15 @@ if max_turns_env is not None: else: DEFAULT_MAX_TURNS = float("inf") +price_limit_env = os.getenv("CAI_PRICE_LIMIT") +if price_limit_env is not None: + try: + DEFAULT_PRICE_LIMIT = float(price_limit_env) + except ValueError: + DEFAULT_PRICE_LIMIT = float("inf") +else: + DEFAULT_PRICE_LIMIT = float("inf") + @dataclass class RunConfig: diff --git a/src/cai/util.py b/src/cai/util.py index a33cb05d..cf94f2ab 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -228,16 +228,22 @@ class CostTracker: # Track the last calculation to debug inconsistencies last_interaction_cost: float = 0.0 last_total_cost: float = 0.0 - - def reset_interaction_stats(self): - """Reset stats for a new interaction""" - self.interaction_input_tokens = 0 - self.interaction_output_tokens = 0 - self.interaction_reasoning_tokens = 0 - self.interaction_cost = 0.0 + + def check_price_limit(self, new_cost: float) -> None: + """Check if adding the new cost would exceed the price limit.""" + from cai.sdk.agents.exceptions import PriceLimitExceeded + from cai.sdk.agents.run import DEFAULT_PRICE_LIMIT + + if DEFAULT_PRICE_LIMIT != float("inf"): + total_cost = self.session_total_cost + new_cost + if total_cost > DEFAULT_PRICE_LIMIT: + raise PriceLimitExceeded(total_cost, DEFAULT_PRICE_LIMIT) def update_session_cost(self, new_cost: float) -> None: """Add cost to session total and log the update""" + # Check price limit before updating + self.check_price_limit(new_cost) + old_total = self.session_total_cost self.session_total_cost += new_cost From 81e14210fe4c25bdbff936d3ebc3e31a64488922 Mon Sep 17 00:00:00 2001 From: luijait Date: Mon, 19 May 2025 10:48:24 +0200 Subject: [PATCH 02/11] Fix commands --- .../agents/models/openai_chatcompletions.py | 354 ++++++++---------- src/cai/tools/common.py | 2 +- 2 files changed, 158 insertions(+), 198 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 40899e5a..ea8c41df 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -289,68 +289,44 @@ class OpenAIChatCompletionsModel(Model): stop_idle_timer() start_active_timer() + # Process current input (user messages, tool results from previous turn) + # and add them to the global message_history. + # _Converter.items_to_messages converts input to ChatCompletionMessageParam list + current_turn_chat_completion_params = _Converter.items_to_messages(input) + for msg_param in current_turn_chat_completion_params: + # Ensure these are plain dicts for add_to_message_history if they are typed objects + # add_to_message_history expects dicts. ChatCompletionMessageParam are TypedDicts. + add_to_message_history(cast(dict, msg_param)) + + # Ensure system instructions are in history (add_to_message_history handles duplicates) + if system_instructions: + sys_msg_for_history = {"role": "system", "content": system_instructions} + add_to_message_history(sys_msg_for_history) + + # Original logic for preparing converted_messages for token counting and logging (local scope) + # This specific `converted_messages` variable is for logging, the API call will use message_history. + # However, for consistent logging, it should reflect what's sent. + # Let's defer defining this until after _fetch_response returns the actual sent messages. + with generation_span( model=str(self.model), model_config=dataclasses.asdict(model_settings) | {"base_url": str(self._client.base_url)}, disabled=tracing.is_disabled(), ) as span_generation: - # Prepare the messages for consistent token counting - converted_messages = _Converter.items_to_messages(input) - if system_instructions: - converted_messages.insert( - 0, - { - "content": system_instructions, - "role": "system", - }, - ) - # --- Add to message_history: user, system, and assistant tool call messages --- - # Add system prompt to message_history - if system_instructions: - sys_msg = { - "role": "system", - "content": system_instructions - } - add_to_message_history(sys_msg) - - # Add user prompt(s) to message_history - if isinstance(input, str): - user_msg = { - "role": "user", - "content": input - } - add_to_message_history(user_msg) - # Log the user message - self.logger.log_user_message(input) - elif isinstance(input, list): - for item in input: - # Try to extract user messages - if isinstance(item, dict): - if item.get("role") == "user": - user_msg = { - "role": "user", - "content": item.get("content", "") - } - add_to_message_history(user_msg) - # Log the user message - if item.get("content"): - self.logger.log_user_message(item.get("content")) + # Get token count estimate using the current state of message_history for better accuracy + # as this reflects what will be sent to _fetch_response. + # Note: fix_message_list might alter it further, this is an estimate. + estimating_messages = list(message_history) # Use a copy + # Append current_turn_chat_completion_params for estimation if not already fully reflected + # This part is tricky as add_to_message_history might deduplicate. + # For simplicity, let's assume message_history is now the source for estimation. + estimated_input_tokens, _ = count_tokens_with_tiktoken(list(message_history)) - # IMPORTANT: Ensure the message list has valid tool call/result pairs - # This needs to happen before the API call to prevent errors - try: - 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}") - - # Get token count estimate before API call for consistent counting - estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) - - response = await self._fetch_response( - system_instructions, - input, + # _fetch_response will now use the global message_history + api_response, messages_sent_to_api = await self._fetch_response( + # system_instructions, # Removed + # input, # Removed model_settings, tools, output_schema, @@ -360,84 +336,79 @@ class OpenAIChatCompletionsModel(Model): stream=False, ) + # Use messages_sent_to_api for logging and further processing if needed + final_converted_messages_for_log = messages_sent_to_api + if _debug.DONT_LOG_MODEL_DATA: logger.debug("Received model response") else: logger.debug( - f"LLM resp:\n{json.dumps(response.choices[0].message.model_dump(), indent=2)}\n" + f"LLM resp:\\n{json.dumps(api_response.choices[0].message.model_dump(), indent=2)}\\n" ) # Ensure we have reasonable token counts - if response.usage: - input_tokens = response.usage.prompt_tokens - output_tokens = response.usage.completion_tokens - total_tokens = response.usage.total_tokens + if api_response.usage: + input_tokens = api_response.usage.prompt_tokens + output_tokens = api_response.usage.completion_tokens + # total_tokens = api_response.usage.total_tokens # total_tokens might be unused # Use estimated tokens if API returns zeroes or implausible values - if input_tokens == 0 or input_tokens < (len(str(input)) // 10): # Sanity check + # Compare against the length of the messages actually sent + if input_tokens == 0 or input_tokens < (len(json.dumps(messages_sent_to_api)) // 20): # Heuristic input_tokens = estimated_input_tokens - total_tokens = input_tokens + output_tokens - - # # Debug information - # print(f"\nDEBUG CONSISTENT TOKEN COUNTS - API tokens: input={input_tokens}, output={output_tokens}, total={total_tokens}") - # print(f"Estimated tokens were: input={estimated_input_tokens}") + # total_tokens = input_tokens + output_tokens else: # If no usage info, use our estimates input_tokens = estimated_input_tokens - output_tokens = 0 - total_tokens = input_tokens - # print(f"\nDEBUG CONSISTENT TOKEN COUNTS - No API tokens, using estimates: input={input_tokens}, output={output_tokens}") + output_tokens = 0 # Output tokens can't be estimated accurately before response + # total_tokens = input_tokens # Update token totals for CLI display self.total_input_tokens += input_tokens - self.total_output_tokens += output_tokens - if (response.usage and - hasattr(response.usage, 'completion_tokens_details') and - response.usage.completion_tokens_details and - hasattr(response.usage.completion_tokens_details, 'reasoning_tokens')): - self.total_reasoning_tokens += response.usage.completion_tokens_details.reasoning_tokens + self.total_output_tokens += output_tokens # This should be from actual response usage + if hasattr(api_response.usage, 'completion_tokens') and api_response.usage.completion_tokens is not None: + self.total_output_tokens = self.total_output_tokens - output_tokens + api_response.usage.completion_tokens # adjust if output_tokens was 0 + output_tokens = api_response.usage.completion_tokens + + + if (api_response.usage and + hasattr(api_response.usage, 'completion_tokens_details') and + api_response.usage.completion_tokens_details and + hasattr(api_response.usage.completion_tokens_details, 'reasoning_tokens')): + self.total_reasoning_tokens += api_response.usage.completion_tokens_details.reasoning_tokens # Check if this message contains tool calls - tool_output = None + # tool_output = None # tool_output seems unused here should_display_message = True - if (hasattr(response.choices[0].message, 'tool_calls') and - response.choices[0].message.tool_calls): + if (hasattr(api_response.choices[0].message, 'tool_calls') and + api_response.choices[0].message.tool_calls): - # For each tool call in the message, get corresponding output if available - for tool_call in response.choices[0].message.tool_calls: + for tool_call in api_response.choices[0].message.tool_calls: call_id = tool_call.id - - # If we're using direct tool output display with cli_print_tool_output, - # and we've already displayed this tool call output, we can skip displaying - # the assistant message to avoid duplication if (hasattr(_Converter, 'tool_outputs') and call_id in _Converter.tool_outputs and hasattr(_Converter, 'recent_tool_calls') and call_id in _Converter.recent_tool_calls): - # We've already displayed this tool and its output directly should_display_message = False break - # 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 + os.environ['CAI_STREAM'] = 'false' - # Print the agent message for CLI display cli_print_agent_messages( agent_name=getattr(self, 'agent_name', 'Agent'), - message=response.choices[0].message, + message=api_response.choices[0].message, counter=getattr(self, 'interaction_counter', 0), model=str(self.model), debug=False, interaction_input_tokens=input_tokens, interaction_output_tokens=output_tokens, interaction_reasoning_tokens=( - response.usage.completion_tokens_details.reasoning_tokens - if response.usage and hasattr(response.usage, 'completion_tokens_details') - and response.usage.completion_tokens_details - and hasattr(response.usage.completion_tokens_details, 'reasoning_tokens') + api_response.usage.completion_tokens_details.reasoning_tokens + if api_response.usage and hasattr(api_response.usage, 'completion_tokens_details') + and api_response.usage.completion_tokens_details + and hasattr(api_response.usage.completion_tokens_details, 'reasoning_tokens') else 0 ), total_input_tokens=getattr(self, 'total_input_tokens', 0), @@ -445,118 +416,96 @@ class OpenAIChatCompletionsModel(Model): total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), interaction_cost=None, total_cost=None, - tool_output=None, # Don't pass tool output here, we're using direct display - suppress_empty=True # Suppress empty panels + tool_output=None, + suppress_empty=True ) - # 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 - assistant_msg = response.choices[0].message - if hasattr(assistant_msg, "tool_calls") and assistant_msg.tool_calls: - for tool_call in assistant_msg.tool_calls: - # Compose a message for the tool call - tool_call_msg = { + assistant_msg_from_api = api_response.choices[0].message + if hasattr(assistant_msg_from_api, "tool_calls") and assistant_msg_from_api.tool_calls: + for tool_call_param in assistant_msg_from_api.tool_calls: + tool_call_dict = { "role": "assistant", - "content": None, + "content": None, # Or assistant_msg_from_api.content if it can coexist "tool_calls": [ { - "id": tool_call.id, - "type": tool_call.type, + "id": tool_call_param.id, + "type": tool_call_param.type, "function": { - "name": tool_call.function.name, - "arguments": tool_call.function.arguments + "name": tool_call_param.function.name, + "arguments": tool_call_param.function.arguments } } ] } + add_to_message_history(tool_call_dict) - add_to_message_history(tool_call_msg) - - # Save the tool call details for later matching with output - # This is important for non-streaming mode to track tool calls properly if not hasattr(_Converter, 'recent_tool_calls'): _Converter.recent_tool_calls = {} - - # Store the tool call by ID for later reference - import time - _Converter.recent_tool_calls[tool_call.id] = { - 'name': tool_call.function.name, - 'arguments': tool_call.function.arguments, - 'start_time': time.time(), - 'execution_info': { - 'start_time': time.time() - } + _Converter.recent_tool_calls[tool_call_param.id] = { + 'name': tool_call_param.function.name, + 'arguments': tool_call_param.function.arguments, + 'start_time': time.time(), # This time might be slightly off; tool call already received + 'execution_info': {'start_time': time.time()} } - # Log the assistant tool call message - tool_calls_list = [] - for tool_call in assistant_msg.tool_calls: - tool_calls_list.append({ - "id": tool_call.id, - "type": tool_call.type, - "function": { - "name": tool_call.function.name, - "arguments": tool_call.function.arguments - } - }) - self.logger.log_assistant_message(None, tool_calls_list) - # If the assistant message is just text, add it as well - elif hasattr(assistant_msg, "content") and assistant_msg.content: - asst_msg = { + tool_calls_for_log = [{ + "id": tc.id, "type": tc.type, + "function": {"name": tc.function.name, "arguments": tc.function.arguments} + } for tc in assistant_msg_from_api.tool_calls] + self.logger.log_assistant_message(None, tool_calls_for_log) + + elif hasattr(assistant_msg_from_api, "content") and assistant_msg_from_api.content: + asst_msg_for_history = { "role": "assistant", - "content": assistant_msg.content + "content": assistant_msg_from_api.content } - add_to_message_history(asst_msg) - # Log the assistant message - self.logger.log_assistant_message(assistant_msg.content) + add_to_message_history(asst_msg_for_history) + self.logger.log_assistant_message(assistant_msg_from_api.content) - # Log the complete response for the session self.logger.rec_training_data( { "model": str(self.model), - "messages": converted_messages, + "messages": final_converted_messages_for_log, # Use the messages actually sent "stream": False, "tools": [t.params_json_schema for t in tools] if tools else [], "tool_choice": model_settings.tool_choice }, - response, + api_response, # This is ChatCompletion, not the Response object self.total_cost ) - usage = ( + usage_obj = ( Usage( requests=1, input_tokens=input_tokens, - output_tokens=output_tokens, + output_tokens=output_tokens, # Ensure this is the final output_tokens total_tokens=input_tokens + output_tokens, ) - if response.usage or input_tokens > 0 + if api_response.usage or input_tokens > 0 else Usage() ) if tracing.include_data(): - span_generation.span_data.output = [response.choices[0].message.model_dump()] + span_generation.span_data.output = [assistant_msg_from_api.model_dump()] span_generation.span_data.usage = { - "input_tokens": usage.input_tokens, - "output_tokens": usage.output_tokens, + "input_tokens": usage_obj.input_tokens, + "output_tokens": usage_obj.output_tokens, } - items = _Converter.message_to_output_items(response.choices[0].message) + items = _Converter.message_to_output_items(assistant_msg_from_api) - # For non-streaming responses, make sure we also log token usage with compatible field names - # This ensures both streaming and non-streaming use consistent naming - if not hasattr(response, 'usage'): - response.usage = {} - if hasattr(response.usage, 'prompt_tokens') and not hasattr(response.usage, 'input_tokens'): - response.usage.input_tokens = response.usage.prompt_tokens - if hasattr(response.usage, 'completion_tokens') and not hasattr(response.usage, 'output_tokens'): - response.usage.output_tokens = response.usage.completion_tokens + # Ensure usage compatibility for ModelResponse + final_usage = {} + if api_response.usage: + final_usage['input_tokens'] = api_response.usage.prompt_tokens + final_usage['output_tokens'] = api_response.usage.completion_tokens + final_usage['total_tokens'] = api_response.usage.total_tokens + # else: ModelResponse will use its default Usage() if no API usage return ModelResponse( output=items, - usage=usage, + usage=Usage(**final_usage) if final_usage else usage_obj, # Pass constructed Usage referenceable_id=None, ) @@ -1362,7 +1311,7 @@ class OpenAIChatCompletionsModel(Model): self.logger.rec_training_data( { "model": str(self.model), - "messages": converted_messages, + "messages": final_converted_messages_for_log, # Use the messages actually sent "stream": True, "tools": [t.params_json_schema for t in tools] if tools else [], "tool_choice": model_settings.tool_choice @@ -1392,8 +1341,8 @@ class OpenAIChatCompletionsModel(Model): @overload async def _fetch_response( self, - system_instructions: str | None, - input: str | list[TResponseInputItem], + # system_instructions: str | None, # REMOVED + # input: str | list[TResponseInputItem], # REMOVED model_settings: ModelSettings, tools: list[Tool], output_schema: AgentOutputSchema | None, @@ -1401,13 +1350,13 @@ class OpenAIChatCompletionsModel(Model): span: Span[GenerationSpanData], tracing: ModelTracing, stream: Literal[True], - ) -> tuple[Response, AsyncStream[ChatCompletionChunk]]: ... + ) -> tuple[Response, AsyncStream[ChatCompletionChunk], list[dict]]: ... # Added messages_sent_to_api @overload async def _fetch_response( self, - system_instructions: str | None, - input: str | list[TResponseInputItem], + # system_instructions: str | None, # REMOVED + # input: str | list[TResponseInputItem], # REMOVED model_settings: ModelSettings, tools: list[Tool], output_schema: AgentOutputSchema | None, @@ -1415,12 +1364,12 @@ class OpenAIChatCompletionsModel(Model): span: Span[GenerationSpanData], tracing: ModelTracing, stream: Literal[False], - ) -> ChatCompletion: ... + ) -> tuple[ChatCompletion, list[dict]]: ... # Added messages_sent_to_api async def _fetch_response( self, - system_instructions: str | None, - input: str | list[TResponseInputItem], + # system_instructions: str | None, # REMOVED + # input: str | list[TResponseInputItem], # REMOVED model_settings: ModelSettings, tools: list[Tool], output_schema: AgentOutputSchema | None, @@ -1428,38 +1377,37 @@ class OpenAIChatCompletionsModel(Model): span: Span[GenerationSpanData], tracing: ModelTracing, stream: bool = False, - ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]: + ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]] | tuple[ChatCompletion, list[dict]] | tuple[Response, AsyncStream[ChatCompletionChunk], list[dict]]: # start by re-fetching self.is_ollama self.is_ollama = os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() == 'true' - converted_messages = _Converter.items_to_messages(input) + # Messages for API are now derived from the global message_history + messages_for_api = list(message_history) # Create a mutable copy - if system_instructions: - converted_messages.insert( - 0, - { - "content": system_instructions, - "role": "system", - }, - ) if tracing.include_data(): - span.span_data.input = converted_messages + span.span_data.input = messages_for_api # Log the full list - # IMPORTANT: Always sanitize the message list to prevent tool call errors - # This is critical to fix common errors with tool/assistant sequences + # IMPORTANT: Always sanitize the message list (which is the full history) try: from cai.util import fix_message_list - prev_length = len(converted_messages) - converted_messages = fix_message_list(converted_messages) - new_length = len(converted_messages) + prev_length = len(messages_for_api) + # Critical: fix_message_list operates on the full history now + messages_for_api_fixed = fix_message_list(messages_for_api) + new_length = len(messages_for_api_fixed) - # Log if the message list was changed significantly if new_length != prev_length: - logger.debug(f"Message list was fixed: {prev_length} -> {new_length} messages") + logger.debug(f"Message list was fixed: {prev_length} -> {new_length} messages. Old: {messages_for_api}, New: {messages_for_api_fixed}") + messages_for_api = messages_for_api_fixed # Use the fixed list except Exception as e: - logger.warning(f"Failed to fix message list: {e}") - + # Log more detailed error if fix_message_list fails + logger.error(f"CRITICAL: fix_message_list failed on message history: {e}", exc_info=True) + logger.error(f"Message history that caused failure: {json.dumps(messages_for_api, indent=2)}") + # It's crucial to understand why fix_message_list might fail here. + # Re-raise or handle as an AgentsException if appropriate. + raise AgentsException(f"Message history sanitization failed processing full history: {e}") from e + + # parallel_tool_calls, tool_choice, response_format, converted_tools... parallel_tool_calls = ( True if model_settings.parallel_tool_calls and tools and len(tools) > 0 else NOT_GIVEN ) @@ -1474,7 +1422,7 @@ class OpenAIChatCompletionsModel(Model): logger.debug("Calling LLM") else: logger.debug( - f"{json.dumps(converted_messages, indent=2)}\n" + f"{json.dumps(messages_for_api, indent=2)}\n" f"Tools:\n{json.dumps(converted_tools, indent=2)}\n" f"Stream: {stream}\n" f"Tool choice: {tool_choice}\n" @@ -1495,7 +1443,7 @@ class OpenAIChatCompletionsModel(Model): # Prepare kwargs for the API call kwargs = { "model": agent_model if agent_model else self.model, - "messages": converted_messages, + "messages": messages_for_api, # Use the fully prepared and fixed message list "tools": converted_tools or NOT_GIVEN, "temperature": self._non_null_or_not_given(model_settings.temperature), "top_p": self._non_null_or_not_given(model_settings.top_p), @@ -1579,9 +1527,21 @@ class OpenAIChatCompletionsModel(Model): try: if self.is_ollama: - return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + # Adjust Ollama fetch to return messages_for_api + if stream: + response_obj, stream_obj = await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + return response_obj, stream_obj, messages_for_api + else: + completion = await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + return completion, messages_for_api else: - return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + # Adjust OpenAI fetch to return messages_for_api + if stream: + response_obj, stream_obj = await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + return response_obj, stream_obj, messages_for_api + else: + completion = await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + return completion, messages_for_api except litellm.exceptions.BadRequestError as e: # print(color("BadRequestError encountered: " + str(e), fg="yellow")) @@ -1770,11 +1730,11 @@ class OpenAIChatCompletionsModel(Model): tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven, stream: bool, parallel_tool_calls: bool - ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]: + ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]: # Return type will be wrapped by _fetch_response """Handle standard LiteLLM API calls for OpenAI and compatible models.""" if stream: # Standard LiteLLM handling for streaming - ret = litellm.completion(**kwargs) + # ret = litellm.completion(**kwargs) # This was likely a typo, should be acompletion for stream stream_obj = await litellm.acompletion(**kwargs) response = Response( @@ -1803,7 +1763,7 @@ class OpenAIChatCompletionsModel(Model): stream: bool, parallel_tool_calls: bool, provider="ollama" - ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]: + ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]: # Return type will be wrapped by _fetch_response # Extract only supported parameters for Ollama ollama_supported_params = { "model": kwargs.get("model", ""), diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 2650a827..1846d56b 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -1249,7 +1249,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg command, stdout, timeout, - stream=stream, + stream=True, call_id=call_id, tool_name=tool_name, workspace_dir=_get_workspace_dir(), From d744d8f7bfea559ed1ec1708cf5c1983242a41f7 Mon Sep 17 00:00:00 2001 From: lidia9 Date: Mon, 19 May 2025 11:08:37 +0200 Subject: [PATCH 03/11] FIX CAI LIMIT PRICE update with /config --- src/cai/util.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/cai/util.py b/src/cai/util.py index cf94f2ab..bda8a98d 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -232,12 +232,17 @@ class CostTracker: def check_price_limit(self, new_cost: float) -> None: """Check if adding the new cost would exceed the price limit.""" from cai.sdk.agents.exceptions import PriceLimitExceeded - from cai.sdk.agents.run import DEFAULT_PRICE_LIMIT - - if DEFAULT_PRICE_LIMIT != float("inf"): + import os + price_limit_env = os.getenv("CAI_PRICE_LIMIT") + try: + price_limit = float(price_limit_env) if price_limit_env is not None else float("inf") + except ValueError: + price_limit = float("inf") + + if price_limit != float("inf"): total_cost = self.session_total_cost + new_cost - if total_cost > DEFAULT_PRICE_LIMIT: - raise PriceLimitExceeded(total_cost, DEFAULT_PRICE_LIMIT) + if total_cost > price_limit: + raise PriceLimitExceeded(total_cost, price_limit) def update_session_cost(self, new_cost: float) -> None: """Add cost to session total and log the update""" From 25bf8f8595d88d61c0351b8e7bd7f3d261ef7417 Mon Sep 17 00:00:00 2001 From: luijait Date: Mon, 19 May 2025 12:29:38 +0200 Subject: [PATCH 04/11] Pricing --- pricing.json | 27 ++++++++++ .../agents/models/openai_chatcompletions.py | 52 +++++++++++++------ src/cai/util.py | 21 ++++++-- 3 files changed, 82 insertions(+), 18 deletions(-) create mode 100644 pricing.json diff --git a/pricing.json b/pricing.json new file mode 100644 index 00000000..f0443799 --- /dev/null +++ b/pricing.json @@ -0,0 +1,27 @@ +{ + "alias0": { + "max_tokens": 128000, + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "input_cost_per_token": 0.000005, + "output_cost_per_token": 0.00005, + "cache_creation_input_token_cost": 0.000005, + "cache_read_input_token_cost": 0.0000005, + "search_context_cost_per_query": { + "search_context_size_low": 1e-2, + "search_context_size_medium": 1e-2, + "search_context_size_high": 1e-2 + }, + "litellm_provider": "openai", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 159, + "supports_assistant_prefill": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "deprecation_date": "2026-02-01", + "supports_tool_choice": true + } +} \ No newline at end of file diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index ea8c41df..f0fe332b 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -625,9 +625,7 @@ class OpenAIChatCompletionsModel(Model): # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) - response, stream = await self._fetch_response( - system_instructions, - input, + response, stream, messages_sent_to_api = await self._fetch_response( model_settings, tools, output_schema, @@ -1237,8 +1235,8 @@ class OpenAIChatCompletionsModel(Model): total_cost = calculate_model_cost(model_name, total_input, total_output) # Explicit conversion to float with fallback to ensure they're never None or 0 - interaction_cost = max(float(interaction_cost if interaction_cost is not None else 0.0), 0.00001) - total_cost = max(float(total_cost if total_cost is not None else 0.0), 0.00001) + interaction_cost = max(float(interaction_cost if interaction_cost is not None else 0.0), 0) + total_cost = max(float(total_cost if total_cost is not None else 0.0), 0) # Store the total cost for future recording self.total_cost = total_cost @@ -1311,7 +1309,7 @@ class OpenAIChatCompletionsModel(Model): self.logger.rec_training_data( { "model": str(self.model), - "messages": final_converted_messages_for_log, # Use the messages actually sent + "messages": messages_sent_to_api, # Use the messages actually sent "stream": True, "tools": [t.params_json_schema for t in tools] if tools else [], "tool_choice": model_settings.tool_choice @@ -1463,7 +1461,7 @@ class OpenAIChatCompletionsModel(Model): model_str = str(kwargs["model"]).lower() if "alias" in model_str: - kwargs["api_base"] = "http://api.aliasrobotics.com:666/" + kwargs["api_base"] = "http://11.0.0.4:4000/" kwargs["custom_llm_provider"] = "openai" kwargs["api_key"] = os.getenv("ALIAS_API_KEY", "REDACTED_ALIAS_KEY") elif "/" in model_str: @@ -1554,7 +1552,13 @@ class OpenAIChatCompletionsModel(Model): if is_qwen: try: # Use the specialized Qwen approach first - return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + ollama_result = await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + if stream: + # ollama_result is (response, stream_obj) + return ollama_result[0], ollama_result[1], messages_for_api + else: + # ollama_result is completion + return ollama_result, messages_for_api except Exception as qwen_e: print(qwen_e) # If that fails, try our direct OpenAI approach @@ -1584,11 +1588,11 @@ class OpenAIChatCompletionsModel(Model): parallel_tool_calls=parallel_tool_calls or False, ) stream_obj = await litellm.acompletion(**qwen_params) - return response, stream_obj + return response, stream_obj, messages_for_api else: # Non-streaming case ret = litellm.completion(**qwen_params) - return ret + return ret, messages_for_api except Exception as direct_e: # All approaches failed, log and raise the original error print(f"All Qwen approaches failed. Original error: {str(e)}, Direct error: {str(direct_e)}") @@ -1659,7 +1663,13 @@ class OpenAIChatCompletionsModel(Model): except Exception as fix_error: print(f"Failed to fix message sequence: {fix_error}") - return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + openai_res = await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + if stream: + # openai_res is (response, stream_obj) + return openai_res[0], openai_res[1], messages_for_api + else: + # openai_res is completion + return openai_res, messages_for_api # this captures an error related to the fact # that the messages list contains an empty @@ -1671,7 +1681,11 @@ class OpenAIChatCompletionsModel(Model): msg if msg.get("content") is not None else {**msg, "content": ""} for msg in kwargs["messages"] ] - return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + openai_res = await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + if stream: + return openai_res[0], openai_res[1], messages_for_api + else: + return openai_res, messages_for_api # Handle Anthropic error for empty text content blocks elif ("text content blocks must be non-empty" in str(e) or @@ -1689,7 +1703,11 @@ class OpenAIChatCompletionsModel(Model): "content": "Empty content block" } for msg in kwargs["messages"] ] - return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + openai_res = await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + if stream: + return openai_res[0], openai_res[1], messages_for_api + else: + return openai_res, messages_for_api else: raise e except litellm.exceptions.RateLimitError as e: @@ -1718,10 +1736,14 @@ class OpenAIChatCompletionsModel(Model): except Exception as e: # pylint: disable=W0718 print(color("Error encountered: " + str(e), fg="yellow")) try: - return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + ollama_result = await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + if stream: + return ollama_result[0], ollama_result[1], messages_for_api + else: + return ollama_result, messages_for_api except Exception as execp: # pylint: disable=W0718 print("Error: " + str(execp)) - return None + raise execp async def _fetch_response_litellm_openai( self, diff --git a/src/cai/util.py b/src/cai/util.py index a33cb05d..c641784e 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -247,7 +247,6 @@ class CostTracker: if os.environ.get("CAI_COST_DISPLAYED", "").lower() == "true": return print(f"\nTotal CAI Session Cost: ${self.session_total_cost:.6f}") - def get_model_pricing(self, model_name: str) -> tuple: """Get and cache pricing information for a model""" # Use the centralized function to standardize model names @@ -256,8 +255,24 @@ class CostTracker: # Check cache first if model_name in self.model_pricing_cache: return self.model_pricing_cache[model_name] + + # Try to load pricing from local pricing.json first + try: + pricing_path = pathlib.Path("pricing.json") + if pricing_path.exists(): + with open(pricing_path, "r", encoding="utf-8") as f: + local_pricing = json.load(f) + pricing_info = local_pricing.get("alias0", {}) + input_cost = pricing_info.get("input_cost_per_token", 0) + output_cost = pricing_info.get("output_cost_per_token", 0) + + # Cache and return local pricing + self.model_pricing_cache[model_name] = (input_cost, output_cost) + return input_cost, output_cost + except Exception as e: + print(f" WARNING: Error loading local pricing.json: {str(e)}") - # Fetch from LiteLLM API + # Fallback to LiteLLM API if local pricing not found LITELLM_URL = ( "https://raw.githubusercontent.com/BerriAI/litellm/main/" "model_prices_and_context_window.json" @@ -280,7 +295,7 @@ class CostTracker: except Exception as e: print(f" WARNING: Error fetching model pricing: {str(e)}") - # Default values if pricing not found + # Default values if no pricing found default_pricing = (0, 0) self.model_pricing_cache[model_name] = default_pricing return default_pricing From f239f732c94d9bdac1f5a17528b9feba50096f0d Mon Sep 17 00:00:00 2001 From: luijait Date: Mon, 19 May 2025 13:26:52 +0200 Subject: [PATCH 05/11] bug --- .../agents/models/openai_chatcompletions.py | 404 +++++++++--------- 1 file changed, 211 insertions(+), 193 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index f0fe332b..8ecb58ff 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -289,44 +289,68 @@ class OpenAIChatCompletionsModel(Model): stop_idle_timer() start_active_timer() - # Process current input (user messages, tool results from previous turn) - # and add them to the global message_history. - # _Converter.items_to_messages converts input to ChatCompletionMessageParam list - current_turn_chat_completion_params = _Converter.items_to_messages(input) - for msg_param in current_turn_chat_completion_params: - # Ensure these are plain dicts for add_to_message_history if they are typed objects - # add_to_message_history expects dicts. ChatCompletionMessageParam are TypedDicts. - add_to_message_history(cast(dict, msg_param)) - - # Ensure system instructions are in history (add_to_message_history handles duplicates) - if system_instructions: - sys_msg_for_history = {"role": "system", "content": system_instructions} - add_to_message_history(sys_msg_for_history) - - # Original logic for preparing converted_messages for token counting and logging (local scope) - # This specific `converted_messages` variable is for logging, the API call will use message_history. - # However, for consistent logging, it should reflect what's sent. - # Let's defer defining this until after _fetch_response returns the actual sent messages. - with generation_span( model=str(self.model), model_config=dataclasses.asdict(model_settings) | {"base_url": str(self._client.base_url)}, disabled=tracing.is_disabled(), ) as span_generation: - # Get token count estimate using the current state of message_history for better accuracy - # as this reflects what will be sent to _fetch_response. - # Note: fix_message_list might alter it further, this is an estimate. - estimating_messages = list(message_history) # Use a copy - # Append current_turn_chat_completion_params for estimation if not already fully reflected - # This part is tricky as add_to_message_history might deduplicate. - # For simplicity, let's assume message_history is now the source for estimation. - estimated_input_tokens, _ = count_tokens_with_tiktoken(list(message_history)) + # Prepare the messages for consistent token counting + converted_messages = _Converter.items_to_messages(input) + if system_instructions: + converted_messages.insert( + 0, + { + "content": system_instructions, + "role": "system", + }, + ) + # --- Add to message_history: user, system, and assistant tool call messages --- + # Add system prompt to message_history + if system_instructions: + sys_msg = { + "role": "system", + "content": system_instructions + } + add_to_message_history(sys_msg) + + # Add user prompt(s) to message_history + if isinstance(input, str): + user_msg = { + "role": "user", + "content": input + } + add_to_message_history(user_msg) + # Log the user message + self.logger.log_user_message(input) + elif isinstance(input, list): + for item in input: + # Try to extract user messages + if isinstance(item, dict): + if item.get("role") == "user": + user_msg = { + "role": "user", + "content": item.get("content", "") + } + add_to_message_history(user_msg) + # Log the user message + if item.get("content"): + self.logger.log_user_message(item.get("content")) - # _fetch_response will now use the global message_history - api_response, messages_sent_to_api = await self._fetch_response( - # system_instructions, # Removed - # input, # Removed + # IMPORTANT: Ensure the message list has valid tool call/result pairs + # This needs to happen before the API call to prevent errors + try: + 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}") + + # Get token count estimate before API call for consistent counting + estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) + + response = await self._fetch_response( + system_instructions, + input, model_settings, tools, output_schema, @@ -336,79 +360,84 @@ class OpenAIChatCompletionsModel(Model): stream=False, ) - # Use messages_sent_to_api for logging and further processing if needed - final_converted_messages_for_log = messages_sent_to_api - if _debug.DONT_LOG_MODEL_DATA: logger.debug("Received model response") else: logger.debug( - f"LLM resp:\\n{json.dumps(api_response.choices[0].message.model_dump(), indent=2)}\\n" + f"LLM resp:\n{json.dumps(response.choices[0].message.model_dump(), indent=2)}\n" ) # Ensure we have reasonable token counts - if api_response.usage: - input_tokens = api_response.usage.prompt_tokens - output_tokens = api_response.usage.completion_tokens - # total_tokens = api_response.usage.total_tokens # total_tokens might be unused + if response.usage: + input_tokens = response.usage.prompt_tokens + output_tokens = response.usage.completion_tokens + total_tokens = response.usage.total_tokens # Use estimated tokens if API returns zeroes or implausible values - # Compare against the length of the messages actually sent - if input_tokens == 0 or input_tokens < (len(json.dumps(messages_sent_to_api)) // 20): # Heuristic + if input_tokens == 0 or input_tokens < (len(str(input)) // 10): # Sanity check input_tokens = estimated_input_tokens - # total_tokens = input_tokens + output_tokens + total_tokens = input_tokens + output_tokens + + # # Debug information + # print(f"\nDEBUG CONSISTENT TOKEN COUNTS - API tokens: input={input_tokens}, output={output_tokens}, total={total_tokens}") + # print(f"Estimated tokens were: input={estimated_input_tokens}") else: # If no usage info, use our estimates input_tokens = estimated_input_tokens - output_tokens = 0 # Output tokens can't be estimated accurately before response - # total_tokens = input_tokens + output_tokens = 0 + total_tokens = input_tokens + # print(f"\nDEBUG CONSISTENT TOKEN COUNTS - No API tokens, using estimates: input={input_tokens}, output={output_tokens}") # Update token totals for CLI display self.total_input_tokens += input_tokens - self.total_output_tokens += output_tokens # This should be from actual response usage - if hasattr(api_response.usage, 'completion_tokens') and api_response.usage.completion_tokens is not None: - self.total_output_tokens = self.total_output_tokens - output_tokens + api_response.usage.completion_tokens # adjust if output_tokens was 0 - output_tokens = api_response.usage.completion_tokens - - - if (api_response.usage and - hasattr(api_response.usage, 'completion_tokens_details') and - api_response.usage.completion_tokens_details and - hasattr(api_response.usage.completion_tokens_details, 'reasoning_tokens')): - self.total_reasoning_tokens += api_response.usage.completion_tokens_details.reasoning_tokens + self.total_output_tokens += output_tokens + if (response.usage and + hasattr(response.usage, 'completion_tokens_details') and + response.usage.completion_tokens_details and + hasattr(response.usage.completion_tokens_details, 'reasoning_tokens')): + self.total_reasoning_tokens += response.usage.completion_tokens_details.reasoning_tokens # Check if this message contains tool calls - # tool_output = None # tool_output seems unused here + tool_output = None should_display_message = True - if (hasattr(api_response.choices[0].message, 'tool_calls') and - api_response.choices[0].message.tool_calls): + if (hasattr(response.choices[0].message, 'tool_calls') and + response.choices[0].message.tool_calls): - for tool_call in api_response.choices[0].message.tool_calls: + # For each tool call in the message, get corresponding output if available + for tool_call in response.choices[0].message.tool_calls: call_id = tool_call.id + + # If we're using direct tool output display with cli_print_tool_output, + # and we've already displayed this tool call output, we can skip displaying + # the assistant message to avoid duplication if (hasattr(_Converter, 'tool_outputs') and call_id in _Converter.tool_outputs and hasattr(_Converter, 'recent_tool_calls') and call_id in _Converter.recent_tool_calls): + # We've already displayed this tool and its output directly should_display_message = False break + # 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' + 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'), - message=api_response.choices[0].message, + message=response.choices[0].message, counter=getattr(self, 'interaction_counter', 0), model=str(self.model), debug=False, interaction_input_tokens=input_tokens, interaction_output_tokens=output_tokens, interaction_reasoning_tokens=( - api_response.usage.completion_tokens_details.reasoning_tokens - if api_response.usage and hasattr(api_response.usage, 'completion_tokens_details') - and api_response.usage.completion_tokens_details - and hasattr(api_response.usage.completion_tokens_details, 'reasoning_tokens') + response.usage.completion_tokens_details.reasoning_tokens + if response.usage and hasattr(response.usage, 'completion_tokens_details') + and response.usage.completion_tokens_details + and hasattr(response.usage.completion_tokens_details, 'reasoning_tokens') else 0 ), total_input_tokens=getattr(self, 'total_input_tokens', 0), @@ -416,96 +445,118 @@ class OpenAIChatCompletionsModel(Model): total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), interaction_cost=None, total_cost=None, - tool_output=None, - suppress_empty=True + 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 - assistant_msg_from_api = api_response.choices[0].message - if hasattr(assistant_msg_from_api, "tool_calls") and assistant_msg_from_api.tool_calls: - for tool_call_param in assistant_msg_from_api.tool_calls: - tool_call_dict = { + # --- Add assistant tool call to message_history if present --- + # If the response contains tool_calls, add them to message_history as assistant messages + assistant_msg = response.choices[0].message + if hasattr(assistant_msg, "tool_calls") and assistant_msg.tool_calls: + for tool_call in assistant_msg.tool_calls: + # Compose a message for the tool call + tool_call_msg = { "role": "assistant", - "content": None, # Or assistant_msg_from_api.content if it can coexist + "content": None, "tool_calls": [ { - "id": tool_call_param.id, - "type": tool_call_param.type, + "id": tool_call.id, + "type": tool_call.type, "function": { - "name": tool_call_param.function.name, - "arguments": tool_call_param.function.arguments + "name": tool_call.function.name, + "arguments": tool_call.function.arguments } } ] } - add_to_message_history(tool_call_dict) + add_to_message_history(tool_call_msg) + + # Save the tool call details for later matching with output + # This is important for non-streaming mode to track tool calls properly if not hasattr(_Converter, 'recent_tool_calls'): _Converter.recent_tool_calls = {} - _Converter.recent_tool_calls[tool_call_param.id] = { - 'name': tool_call_param.function.name, - 'arguments': tool_call_param.function.arguments, - 'start_time': time.time(), # This time might be slightly off; tool call already received - 'execution_info': {'start_time': time.time()} + + # Store the tool call by ID for later reference + import time + _Converter.recent_tool_calls[tool_call.id] = { + 'name': tool_call.function.name, + 'arguments': tool_call.function.arguments, + 'start_time': time.time(), + 'execution_info': { + 'start_time': time.time() + } } - tool_calls_for_log = [{ - "id": tc.id, "type": tc.type, - "function": {"name": tc.function.name, "arguments": tc.function.arguments} - } for tc in assistant_msg_from_api.tool_calls] - self.logger.log_assistant_message(None, tool_calls_for_log) - - elif hasattr(assistant_msg_from_api, "content") and assistant_msg_from_api.content: - asst_msg_for_history = { + # Log the assistant tool call message + tool_calls_list = [] + for tool_call in assistant_msg.tool_calls: + tool_calls_list.append({ + "id": tool_call.id, + "type": tool_call.type, + "function": { + "name": tool_call.function.name, + "arguments": tool_call.function.arguments + } + }) + self.logger.log_assistant_message(None, tool_calls_list) + # If the assistant message is just text, add it as well + elif hasattr(assistant_msg, "content") and assistant_msg.content: + asst_msg = { "role": "assistant", - "content": assistant_msg_from_api.content + "content": assistant_msg.content } - add_to_message_history(asst_msg_for_history) - self.logger.log_assistant_message(assistant_msg_from_api.content) + add_to_message_history(asst_msg) + # Log the assistant message + self.logger.log_assistant_message(assistant_msg.content) + # Log the complete response for the session self.logger.rec_training_data( { "model": str(self.model), - "messages": final_converted_messages_for_log, # Use the messages actually sent + "messages": converted_messages, "stream": False, "tools": [t.params_json_schema for t in tools] if tools else [], "tool_choice": model_settings.tool_choice }, - api_response, # This is ChatCompletion, not the Response object + response, self.total_cost ) - usage_obj = ( + usage = ( Usage( requests=1, input_tokens=input_tokens, - output_tokens=output_tokens, # Ensure this is the final output_tokens + output_tokens=output_tokens, total_tokens=input_tokens + output_tokens, ) - if api_response.usage or input_tokens > 0 + if response.usage or input_tokens > 0 else Usage() ) if tracing.include_data(): - span_generation.span_data.output = [assistant_msg_from_api.model_dump()] + span_generation.span_data.output = [response.choices[0].message.model_dump()] span_generation.span_data.usage = { - "input_tokens": usage_obj.input_tokens, - "output_tokens": usage_obj.output_tokens, + "input_tokens": usage.input_tokens, + "output_tokens": usage.output_tokens, } - items = _Converter.message_to_output_items(assistant_msg_from_api) + items = _Converter.message_to_output_items(response.choices[0].message) - # Ensure usage compatibility for ModelResponse - final_usage = {} - if api_response.usage: - final_usage['input_tokens'] = api_response.usage.prompt_tokens - final_usage['output_tokens'] = api_response.usage.completion_tokens - final_usage['total_tokens'] = api_response.usage.total_tokens - # else: ModelResponse will use its default Usage() if no API usage + # For non-streaming responses, make sure we also log token usage with compatible field names + # This ensures both streaming and non-streaming use consistent naming + if not hasattr(response, 'usage'): + response.usage = {} + if hasattr(response.usage, 'prompt_tokens') and not hasattr(response.usage, 'input_tokens'): + response.usage.input_tokens = response.usage.prompt_tokens + if hasattr(response.usage, 'completion_tokens') and not hasattr(response.usage, 'output_tokens'): + response.usage.output_tokens = response.usage.completion_tokens return ModelResponse( output=items, - usage=Usage(**final_usage) if final_usage else usage_obj, # Pass constructed Usage + usage=usage, referenceable_id=None, ) @@ -625,7 +676,9 @@ class OpenAIChatCompletionsModel(Model): # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) - response, stream, messages_sent_to_api = await self._fetch_response( + response, stream = await self._fetch_response( + system_instructions, + input, model_settings, tools, output_schema, @@ -1235,8 +1288,8 @@ class OpenAIChatCompletionsModel(Model): total_cost = calculate_model_cost(model_name, total_input, total_output) # Explicit conversion to float with fallback to ensure they're never None or 0 - interaction_cost = max(float(interaction_cost if interaction_cost is not None else 0.0), 0) - total_cost = max(float(total_cost if total_cost is not None else 0.0), 0) + interaction_cost = float(interaction_cost if interaction_cost is not None else 0.0) + total_cost = float(total_cost if total_cost is not None else 0.0) # Store the total cost for future recording self.total_cost = total_cost @@ -1309,7 +1362,7 @@ class OpenAIChatCompletionsModel(Model): self.logger.rec_training_data( { "model": str(self.model), - "messages": messages_sent_to_api, # Use the messages actually sent + "messages": converted_messages, "stream": True, "tools": [t.params_json_schema for t in tools] if tools else [], "tool_choice": model_settings.tool_choice @@ -1339,8 +1392,8 @@ class OpenAIChatCompletionsModel(Model): @overload async def _fetch_response( self, - # system_instructions: str | None, # REMOVED - # input: str | list[TResponseInputItem], # REMOVED + system_instructions: str | None, + input: str | list[TResponseInputItem], model_settings: ModelSettings, tools: list[Tool], output_schema: AgentOutputSchema | None, @@ -1348,13 +1401,13 @@ class OpenAIChatCompletionsModel(Model): span: Span[GenerationSpanData], tracing: ModelTracing, stream: Literal[True], - ) -> tuple[Response, AsyncStream[ChatCompletionChunk], list[dict]]: ... # Added messages_sent_to_api + ) -> tuple[Response, AsyncStream[ChatCompletionChunk]]: ... @overload async def _fetch_response( self, - # system_instructions: str | None, # REMOVED - # input: str | list[TResponseInputItem], # REMOVED + system_instructions: str | None, + input: str | list[TResponseInputItem], model_settings: ModelSettings, tools: list[Tool], output_schema: AgentOutputSchema | None, @@ -1362,12 +1415,12 @@ class OpenAIChatCompletionsModel(Model): span: Span[GenerationSpanData], tracing: ModelTracing, stream: Literal[False], - ) -> tuple[ChatCompletion, list[dict]]: ... # Added messages_sent_to_api + ) -> ChatCompletion: ... async def _fetch_response( self, - # system_instructions: str | None, # REMOVED - # input: str | list[TResponseInputItem], # REMOVED + system_instructions: str | None, + input: str | list[TResponseInputItem], model_settings: ModelSettings, tools: list[Tool], output_schema: AgentOutputSchema | None, @@ -1375,37 +1428,38 @@ class OpenAIChatCompletionsModel(Model): span: Span[GenerationSpanData], tracing: ModelTracing, stream: bool = False, - ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]] | tuple[ChatCompletion, list[dict]] | tuple[Response, AsyncStream[ChatCompletionChunk], list[dict]]: + ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]: # start by re-fetching self.is_ollama self.is_ollama = os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() == 'true' - # Messages for API are now derived from the global message_history - messages_for_api = list(message_history) # Create a mutable copy + converted_messages = _Converter.items_to_messages(input) + if system_instructions: + converted_messages.insert( + 0, + { + "content": system_instructions, + "role": "system", + }, + ) if tracing.include_data(): - span.span_data.input = messages_for_api # Log the full list + span.span_data.input = converted_messages - # IMPORTANT: Always sanitize the message list (which is the full history) + # IMPORTANT: Always sanitize the message list to prevent tool call errors + # This is critical to fix common errors with tool/assistant sequences try: from cai.util import fix_message_list - prev_length = len(messages_for_api) - # Critical: fix_message_list operates on the full history now - messages_for_api_fixed = fix_message_list(messages_for_api) - new_length = len(messages_for_api_fixed) + prev_length = len(converted_messages) + converted_messages = fix_message_list(converted_messages) + new_length = len(converted_messages) + # Log if the message list was changed significantly if new_length != prev_length: - logger.debug(f"Message list was fixed: {prev_length} -> {new_length} messages. Old: {messages_for_api}, New: {messages_for_api_fixed}") - messages_for_api = messages_for_api_fixed # Use the fixed list + logger.debug(f"Message list was fixed: {prev_length} -> {new_length} messages") except Exception as e: - # Log more detailed error if fix_message_list fails - logger.error(f"CRITICAL: fix_message_list failed on message history: {e}", exc_info=True) - logger.error(f"Message history that caused failure: {json.dumps(messages_for_api, indent=2)}") - # It's crucial to understand why fix_message_list might fail here. - # Re-raise or handle as an AgentsException if appropriate. - raise AgentsException(f"Message history sanitization failed processing full history: {e}") from e - - # parallel_tool_calls, tool_choice, response_format, converted_tools... + logger.warning(f"Failed to fix message list: {e}") + parallel_tool_calls = ( True if model_settings.parallel_tool_calls and tools and len(tools) > 0 else NOT_GIVEN ) @@ -1420,7 +1474,7 @@ class OpenAIChatCompletionsModel(Model): logger.debug("Calling LLM") else: logger.debug( - f"{json.dumps(messages_for_api, indent=2)}\n" + f"{json.dumps(converted_messages, indent=2)}\n" f"Tools:\n{json.dumps(converted_tools, indent=2)}\n" f"Stream: {stream}\n" f"Tool choice: {tool_choice}\n" @@ -1441,7 +1495,7 @@ class OpenAIChatCompletionsModel(Model): # Prepare kwargs for the API call kwargs = { "model": agent_model if agent_model else self.model, - "messages": messages_for_api, # Use the fully prepared and fixed message list + "messages": converted_messages, "tools": converted_tools or NOT_GIVEN, "temperature": self._non_null_or_not_given(model_settings.temperature), "top_p": self._non_null_or_not_given(model_settings.top_p), @@ -1461,7 +1515,7 @@ class OpenAIChatCompletionsModel(Model): model_str = str(kwargs["model"]).lower() if "alias" in model_str: - kwargs["api_base"] = "http://11.0.0.4:4000/" + kwargs["api_base"] = "http://api.aliasrobotics.com:666/" kwargs["custom_llm_provider"] = "openai" kwargs["api_key"] = os.getenv("ALIAS_API_KEY", "REDACTED_ALIAS_KEY") elif "/" in model_str: @@ -1525,21 +1579,9 @@ class OpenAIChatCompletionsModel(Model): try: if self.is_ollama: - # Adjust Ollama fetch to return messages_for_api - if stream: - response_obj, stream_obj = await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) - return response_obj, stream_obj, messages_for_api - else: - completion = await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) - return completion, messages_for_api + return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) else: - # Adjust OpenAI fetch to return messages_for_api - if stream: - response_obj, stream_obj = await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) - return response_obj, stream_obj, messages_for_api - else: - completion = await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) - return completion, messages_for_api + 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")) @@ -1552,13 +1594,7 @@ class OpenAIChatCompletionsModel(Model): if is_qwen: try: # Use the specialized Qwen approach first - ollama_result = await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) - if stream: - # ollama_result is (response, stream_obj) - return ollama_result[0], ollama_result[1], messages_for_api - else: - # ollama_result is completion - return ollama_result, messages_for_api + 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 @@ -1588,11 +1624,11 @@ class OpenAIChatCompletionsModel(Model): parallel_tool_calls=parallel_tool_calls or False, ) stream_obj = await litellm.acompletion(**qwen_params) - return response, stream_obj, messages_for_api + return response, stream_obj else: # Non-streaming case ret = litellm.completion(**qwen_params) - return ret, messages_for_api + return ret except Exception as direct_e: # All approaches failed, log and raise the original error print(f"All Qwen approaches failed. Original error: {str(e)}, Direct error: {str(direct_e)}") @@ -1663,13 +1699,7 @@ class OpenAIChatCompletionsModel(Model): except Exception as fix_error: print(f"Failed to fix message sequence: {fix_error}") - openai_res = await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) - if stream: - # openai_res is (response, stream_obj) - return openai_res[0], openai_res[1], messages_for_api - else: - # openai_res is completion - return openai_res, messages_for_api + return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) # this captures an error related to the fact # that the messages list contains an empty @@ -1681,11 +1711,7 @@ class OpenAIChatCompletionsModel(Model): msg if msg.get("content") is not None else {**msg, "content": ""} for msg in kwargs["messages"] ] - openai_res = await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) - if stream: - return openai_res[0], openai_res[1], messages_for_api - else: - return openai_res, messages_for_api + return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) # Handle Anthropic error for empty text content blocks elif ("text content blocks must be non-empty" in str(e) or @@ -1703,11 +1729,7 @@ class OpenAIChatCompletionsModel(Model): "content": "Empty content block" } for msg in kwargs["messages"] ] - openai_res = await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) - if stream: - return openai_res[0], openai_res[1], messages_for_api - else: - return openai_res, messages_for_api + return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) else: raise e except litellm.exceptions.RateLimitError as e: @@ -1736,14 +1758,10 @@ class OpenAIChatCompletionsModel(Model): except Exception as e: # pylint: disable=W0718 print(color("Error encountered: " + str(e), fg="yellow")) try: - ollama_result = await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) - if stream: - return ollama_result[0], ollama_result[1], messages_for_api - else: - return ollama_result, messages_for_api + return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) except Exception as execp: # pylint: disable=W0718 print("Error: " + str(execp)) - raise execp + return None async def _fetch_response_litellm_openai( self, @@ -1752,11 +1770,11 @@ class OpenAIChatCompletionsModel(Model): tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven, stream: bool, parallel_tool_calls: bool - ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]: # Return type will be wrapped by _fetch_response + ) -> 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) # This was likely a typo, should be acompletion for stream + ret = litellm.completion(**kwargs) stream_obj = await litellm.acompletion(**kwargs) response = Response( @@ -1785,7 +1803,7 @@ class OpenAIChatCompletionsModel(Model): stream: bool, parallel_tool_calls: bool, provider="ollama" - ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]: # Return type will be wrapped by _fetch_response + ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]: # Extract only supported parameters for Ollama ollama_supported_params = { "model": kwargs.get("model", ""), From be79caf7dc50e0b69a95ce4e348730a8acae77fe Mon Sep 17 00:00:00 2001 From: luijait Date: Mon, 19 May 2025 13:39:03 +0200 Subject: [PATCH 06/11] Pricing in streaming --- src/cai/sdk/agents/models/openai_chatcompletions.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 8ecb58ff..f92436b3 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -14,7 +14,7 @@ import asyncio from collections.abc import AsyncIterator, Iterable from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, cast, overload -from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages, create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming, calculate_model_cost +from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages, create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming, calculate_model_cost, COST_TRACKER from cai.util import start_idle_timer, stop_idle_timer, start_active_timer, stop_active_timer from wasabi import color from cai.sdk.agents.run_to_jsonl import get_session_recorder @@ -1291,6 +1291,10 @@ class OpenAIChatCompletionsModel(Model): interaction_cost = float(interaction_cost if interaction_cost is not None else 0.0) total_cost = float(total_cost if total_cost is not None else 0.0) + # Update the global COST_TRACKER with the cost of this specific interaction + if hasattr(COST_TRACKER, "add_interaction_cost") and interaction_cost > 0.0: + COST_TRACKER.add_interaction_cost(interaction_cost) + # Store the total cost for future recording self.total_cost = total_cost From 18439c504f5291cb7008f59cfb40f1a5f015c8cf Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 13:15:20 +0200 Subject: [PATCH 07/11] add info --- src/cai/cli.py | 36 +++++++++---- src/cai/tools/common.py | 3 ++ src/cai/util.py | 109 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 138 insertions(+), 10 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index a7283c83..c1766b59 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -134,7 +134,10 @@ from cai.util import ( start_idle_timer, stop_idle_timer, start_active_timer, - stop_active_timer + stop_active_timer, + setup_ctf, + check_flag + ) # CAI REPL imports @@ -151,6 +154,11 @@ from cai.internal.components.metrics import process_metrics # Add import for parallel configs at the top of the file from cai.repl.commands.parallel import PARALLEL_CONFIGS, ParallelConfig +from cai import is_pentestperf_available +ctf_global = None +if is_pentestperf_available() and os.getenv('CTF_NAME', None): + ctf, messages_ctf = setup_ctf() + ctf_global = ctf # Load environment variables from .env file load_dotenv() @@ -187,7 +195,7 @@ agent = Agent( ) ) -def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf')): +def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), force_until_flag=False): """ Run a simple interactive CLI loop for CAI. @@ -283,14 +291,17 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf')): except Exception as e: console.print(f"[red]Error switching agent: {str(e)}[/red]") - # Get user input with command completion and history - user_input = get_user_input( - command_completer, - kb, - history_file, - get_toolbar_with_refresh, - current_text - ) + if not force_until_flag: + # Get user input with command completion and history + user_input = get_user_input( + command_completer, + kb, + history_file, + get_toolbar_with_refresh, + current_text + ) + else: + user_input = messages_ctf idle_time += time.time() - idle_start_time # Stop measuring user idle time and start measuring active time @@ -393,6 +404,11 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf')): # Prevent duplicate cost display from the COST_TRACKER exit handler os.environ["CAI_COST_DISPLAYED"] = "true" + if (is_pentestperf_available() and os.getenv('CTF_NAME', None)): + ctf.stop_ctf() + return False + + except Exception: pass break diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 1846d56b..836a41fb 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -708,6 +708,9 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg # Use the active timer during tool execution stop_idle_timer() start_active_timer() + + from cai.cli import ctf_global + ctf = ctf_global # Parse command into standard parts to ensure consistent naming parts = command.strip().split(' ', 1) diff --git a/src/cai/util.py b/src/cai/util.py index d576202f..82081f95 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -29,6 +29,9 @@ from rich.box import ROUNDED from rich.table import Table import re import uuid +from cai import is_pentestperf_available +if is_pentestperf_available(): + import pentestperf as ptt # Global timing variables for tracking active and idle time _active_timer_start = None @@ -2747,3 +2750,109 @@ def finish_tool_streaming(tool_name, args, output, call_id, execution_info=None, # Mark the streaming session as complete if hasattr(cli_print_tool_output, '_streaming_sessions') and call_id in cli_print_tool_output._streaming_sessions: cli_print_tool_output._streaming_sessions[call_id]['is_complete'] = True + +def check_flag(output, ctf, challenge=None): + """ + Check if the CTF flag is present in the output. + + Args: + output (str): The output to check for the flag. + ctf: The CTF environment object. + challenge (str, optional): The specific challenge to check. + Defaults to None. + + Returns: + tuple: A tuple containing a boolean indicating if the flag was + found and the flag itself if found, otherwise None. + """ + # Get the challenge from the environment variable or default to the first + # challenge + challenge_key = os.getenv("CTF_CHALLENGE") + challenges = list(ctf.get_challenges().keys()) + challenge = ( + challenge_key + if challenge_key in challenges + else (challenges[0] if len(challenges) > 0 else None)) + if ctf: + if ctf.check_flag( + output, challenge + ): # check if the flag is in the output + flag = ctf.flags[challenge] + print( + color( + f"Flag found: {flag}", + fg="green") + + " in output " + + color( + f"{output}", + fg="blue")) + return True, flag + else: + print(color("CTF environment not found or provided", fg="yellow")) + return False, None + +def setup_ctf(): + """Setup CTF environment if CTF_NAME is provided""" + ctf_name = os.getenv('CTF_NAME', None) + if not ctf_name: + print(color("CTF name not provided, necessary to run CTF", fg="white", bg="red")) + sys.exit(1) + + print(color("Setting up CTF: ", fg="black", bg="yellow") + + color(ctf_name, fg="black", bg="yellow")) + + ctf = ptt.ctf( # pylint: disable=I1101 # noqa + ctf_name, + subnet=os.getenv('CTF_SUBNET', "192.168.2.0/24"), + container_name="ctf_target", + ip_address=os.getenv('CTF_IP', "192.168.2.100"), + ) + ctf.start_ctf() + + # Get the challenge from the environment variable or default to the + # first challenge + challenge_key = os.getenv('CTF_CHALLENGE') # TODO: + challenges = list(ctf.get_challenges().keys()) + challenge = challenge_key if challenge_key in challenges else ( + challenges[0] if len(challenges) > 0 else None) + + # Use the user master template + messages = Template( + filename="src/cai/prompts/core/user_master_template.md").render( + ctf=ctf, + challenge=challenge, + ip=ctf.get_ip() if ctf else None, + ) + + + print( color( + "Testing CTF: ", + fg="black", + bg="yellow") + + color( + ctf.name, + fg="black", + bg="yellow")) + if not challenge_key or challenge_key not in challenges: + print( + color( + "No challenge provided or challenge not found. Attempting to use the first challenge.", + fg="white", + bg="blue")) + if challenge: + print( + color( + "Testing challenge: ", + fg="white", + bg="blue") + + color( + "'" + + challenge + + "' (" + + repr( + ctf.flags[challenge]) + + ")", + fg="white", + bg="blue")) + + return ctf, messages From cdafb569c23827bc96af0f594c4f5c4c7a7f8df7 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 13:26:25 +0200 Subject: [PATCH 08/11] inside ctf is working --- src/cai/tools/common.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 836a41fb..11c66772 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -1014,7 +1014,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501 # --- CTF Execution --- - if ctf: + if ctf and os.getenv('CTF_INSIDE', True) == "True": # If streaming is enabled and we have a call_id, show streaming UI for CTF too if stream: # Import the streaming utilities from util @@ -1037,8 +1037,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg call_id = start_tool_streaming(tool_name, tool_args, call_id) target_dir = _get_workspace_dir() - full_command = f"cd '{target_dir}' && {command}" - + #full_command = f"cd '{target_dir}' && {command}" + full_command = command # Update with "executing" status update_tool_streaming( tool_name, From 4c2c979e4831651989a06d5dbdeb77fcae1df04c Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 13:44:53 +0200 Subject: [PATCH 09/11] add true lower --- src/cai/tools/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 11c66772..61bc2a5f 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -1014,7 +1014,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501 # --- CTF Execution --- - if ctf and os.getenv('CTF_INSIDE', True) == "True": + if ctf and os.getenv('CTF_INSIDE', True).lower() == "true": # If streaming is enabled and we have a call_id, show streaming UI for CTF too if stream: # Import the streaming utilities from util From 46e00813bbf3bbbffd2dd195640b7fdc07d70e6d Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 14:36:19 +0200 Subject: [PATCH 10/11] solve issue --- src/cai/tools/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 61bc2a5f..8b9e3e46 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -759,7 +759,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg is_ssh_env = all(os.getenv(var) for var in ['SSH_USER', 'SSH_HOST']) # --- Docker Container Execution --- - if active_container and not ctf and not is_ssh_env: + if active_container and not is_ssh_env: container_id = active_container container_workspace = _get_container_workspace_path() context_msg = f"(docker:{container_id[:12]}:{container_workspace})" From fd60febfbac6c1f68a07ef809834f143fded97a9 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 19 May 2025 14:56:30 +0200 Subject: [PATCH 11/11] fix --- src/cai/cli.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index c1766b59..8a6c3269 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -156,6 +156,7 @@ from cai.internal.components.metrics import process_metrics from cai.repl.commands.parallel import PARALLEL_CONFIGS, ParallelConfig from cai import is_pentestperf_available ctf_global = None +messages_ctf = "" if is_pentestperf_available() and os.getenv('CTF_NAME', None): ctf, messages_ctf = setup_ctf() ctf_global = ctf @@ -299,7 +300,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), history_file, get_toolbar_with_refresh, current_text - ) + ) + messages_ctf else: user_input = messages_ctf idle_time += time.time() - idle_start_time