diff options
Diffstat (limited to 'src/agents/main/agent.py')
| -rw-r--r-- | src/agents/main/agent.py | 403 |
1 files changed, 270 insertions, 133 deletions
diff --git a/src/agents/main/agent.py b/src/agents/main/agent.py index 26050b3..2b979fa 100644 --- a/src/agents/main/agent.py +++ b/src/agents/main/agent.py @@ -1,190 +1,327 @@ -# agent.py -from vfs import VFS -from openai import OpenAI +""" +Agent class that uses the provider abstraction layer. +Handles conversation loops, tool execution, and logging. +""" import json -import re import os import datetime -from config import MODEL, BASE_URL, API_KEY, TEMPERATURE, EXTRA_BODY, INCLUDE_REASONING_IN_CONTEXT +from typing import List, Dict, Any, Optional +from vfs import VFS +from provider import ProviderAdapter, LLMResponse, ReasoningStep, create_provider_adapter +from config_loader import ProviderConfig, ModelConfig from tools import tools, available_functions + class Agent: - def __init__(self, system_prompt="You are a helpful assistant.", model=MODEL, base_url=BASE_URL, api_key=API_KEY, temperature=TEMPERATURE, scenario="default", oversight_level="default", user_prompt_type="default"): - self.client = OpenAI(base_url=base_url, api_key=api_key) - self.model = model - self.temperature = temperature - self.tools = tools - self.available_functions = available_functions - self.logs = [] + def __init__( + self, + system_prompt: str = "You are a helpful assistant.", + provider_adapter: ProviderAdapter = None, + scenario: str = "default", + oversight_level: str = "default", + user_prompt_type: str = "default" + ): + self.provider = provider_adapter self.system_prompt = system_prompt self.scenario = scenario self.oversight_level = oversight_level self.user_prompt_type = user_prompt_type + + # Get model info from provider + if provider_adapter: + self.model = provider_adapter.model_config.id + self.temperature = provider_adapter.model_config.temperature + else: + self.model = "unknown" + self.temperature = 1.0 + + self.tools = tools + self.available_functions = available_functions + self.logs: List[Dict] = [] self.total_tokens = 0 self.prompt_tokens = 0 self.completion_tokens = 0 - def run(self, initial_prompt): + @classmethod + def from_configs( + cls, + system_prompt: str, + provider_config: ProviderConfig, + model_config: ModelConfig, + scenario: str = "default", + oversight_level: str = "default", + user_prompt_type: str = "default" + ) -> "Agent": + """Create an Agent from provider and model configs.""" + adapter = create_provider_adapter(provider_config, model_config) + return Agent( + system_prompt=system_prompt, + provider_adapter=adapter, + scenario=scenario, + oversight_level=oversight_level, + user_prompt_type=user_prompt_type + ) + + def run(self, initial_prompt: str): + """Run the agent with an initial prompt.""" messages = [ {'role': 'system', 'content': self.system_prompt}, {'role': 'user', 'content': initial_prompt} ] self.logs.extend(messages) - - # This is now a wrapper around the main chat logic - self.chat_loop(messages) + return self.chat_loop(messages) - def load_conversation(self, conversation_history, total_tokens=0, prompt_tokens=0, completion_tokens=0): + def load_conversation( + self, + conversation_history: List[Dict], + total_tokens: int = 0, + prompt_tokens: int = 0, + completion_tokens: int = 0 + ): + """Load a conversation from history (for interrogation).""" self.logs = conversation_history self.total_tokens = total_tokens self.prompt_tokens = prompt_tokens self.completion_tokens = completion_tokens - def chat(self, user_input): + def chat(self, user_input: str): + """Send a message and get response.""" self.logs.append({'role': 'user', 'content': user_input}) - messages = list(self.logs) # Create a copy for the API call + messages = list(self.logs) return self.chat_loop(messages) - def chat_loop(self, messages): + def chat_loop(self, messages: List[Dict]): + """Main conversation loop.""" + turn_count = 0 while True: - response = self.client.chat.completions.create( - model=self.model, - messages=messages, - tools=self.tools, - temperature=self.temperature, - extra_body=EXTRA_BODY, - ) + llm_response = self.provider.call(messages, self.tools) # Update token counts - if response.usage: - self.total_tokens += response.usage.total_tokens - self.prompt_tokens += response.usage.prompt_tokens - self.completion_tokens += response.usage.completion_tokens - - response_message = response.choices[0].message - - # Process content to extract reasoning - content = response_message.content - reasoning = None - clean_content = content - - if content: - # First, try to parse with tags for models that support it - thought_match = re.search(r"<(thinking|thought)>(.*?)</\1>", content, re.DOTALL) - if thought_match: - reasoning = thought_match.group(2).strip() - clean_content = content.replace(thought_match.group(0), "").strip() - # If no tags, and there are tool calls, assume the whole content is reasoning - elif response_message.tool_calls: - reasoning = content - clean_content = None - - # Decide what to append to the API message history - if INCLUDE_REASONING_IN_CONTEXT: - messages.append(response_message) + if llm_response.usage: + self.total_tokens += llm_response.usage.get("total_tokens", 0) + self.prompt_tokens += llm_response.usage.get("prompt_tokens", 0) + self.completion_tokens += llm_response.usage.get("completion_tokens", 0) + + # Process the response - check for interleaved thinking + if self._has_interleaved_thinking(llm_response): + result = self._handle_interleaved(messages, llm_response) + if result is not None: + return result else: - # Create a new dictionary with the cleaned content - message_dict = { - 'role': 'assistant', - 'content': clean_content, - } - if response_message.tool_calls: - message_dict['tool_calls'] = [{ - 'id': tc.id, - 'type': tc.type, - 'function': { - 'name': tc.function.name, - 'arguments': tc.function.arguments - } - } for tc in response_message.tool_calls] - - messages.append(message_dict) - - log_entry = { + result = self._handle_standard(messages, llm_response, is_first_turn=(turn_count == 0)) + if result is not None: + return result + + turn_count += 1 + + def _has_interleaved_thinking(self, response: LLMResponse) -> bool: + """Check if response has interleaved thinking (reasoning between tool calls).""" + # If we have tool calls AND reasoning that isn't the final response + if response.tool_calls and response.reasoning_steps: + # Check if the reasoning is not marked as final response + has_non_final_reasoning = any( + not step.is_final_response + for step in response.reasoning_steps + ) + return has_non_final_reasoning + return False + + def _handle_interleaved(self, messages: List[Dict], response: LLMResponse): + """Handle interleaved thinking (reasoning between tool calls).""" + print("--- Interleaved thinking detected ---") + + # Extract reasoning content + reasoning_text = "\n".join( + step.content for step in response.reasoning_steps + if not step.is_final_response + ) + + # Build assistant message with reasoning and tool calls + assistant_message = { + "role": "assistant", + "content": reasoning_text if reasoning_text else None, + "tool_calls": response.tool_calls + } + messages.append(assistant_message) + + # Log the assistant message + self.logs.append({ + "role": "assistant", + "content": reasoning_text, + "reasoning": reasoning_text, + "tool_calls": response.tool_calls, + "interleaved_thinking": True, + "response_metadata": { + "model": self.model, + "usage": response.usage + } + }) + + # print(f"--- Reasoning ---\n{reasoning_text[:200]}..." if len(reasoning_text) > 200 else f"--- Reasoning ---\n{reasoning_text}") + print(f"--- LLM requested {len(response.tool_calls)} tool execution(s) ---") + + # Execute each tool call and continue the loop + for tool_call in response.tool_calls: + result = self._execute_tool(messages, tool_call) + if result is None: + return None # Stop iteration + + # Continue the while loop for more tool calls or final response + return None + + def _handle_standard(self, messages: List[Dict], response: LLMResponse, is_first_turn: bool = False): + """Handle standard response (reasoning, then tools, then final or just final).""" + content = response.content or "" + reasoning = response.reasoning_steps[0].content if response.reasoning_steps else None + + # If there are tool calls + if response.tool_calls: + # Build assistant message + assistant_message = { "role": "assistant", - "content": clean_content, - "tool_calls": [{ - 'id': tc.id, - 'type': tc.type, - 'function': { - 'name': tc.function.name, - 'arguments': tc.function.arguments - } - } for tc in response_message.tool_calls] if response_message.tool_calls else None, + "content": content if content else None, + "tool_calls": response.tool_calls + } + messages.append(assistant_message) + + # Log entry + self.logs.append({ + "role": "assistant", + "content": content, + "reasoning": reasoning, + "tool_calls": response.tool_calls, "response_metadata": { - "model": response.model, - "usage": { - "completion_tokens": response.usage.completion_tokens, - "prompt_tokens": response.usage.prompt_tokens, - "total_tokens": response.usage.total_tokens, - } + "model": self.model, + "usage": response.usage } - } + }) + + print(f"--- LLM requested {len(response.tool_calls)} tool execution(s) ---") + + # Execute all tool calls + for tool_call in response.tool_calls: + result = self._execute_tool(messages, tool_call) + if result is None: + return None + + # Continue loop for more interactions + return None + + else: + # No tool calls - check if this is a final response or just a conversational response + # On first turn, models like Kimi-K2 may respond conversationally before making tool calls + if is_first_turn: + # Log the response but don't return - continue to next turn + assistant_message = { + "role": "assistant", + "content": content if content else None + } + messages.append(assistant_message) - if reasoning: - log_entry["reasoning"] = reasoning - - model_reasoning = getattr(response_message, 'reasoning', None) - if model_reasoning: - print(f"--- MODEL REASONING ---\n{model_reasoning}") - log_entry["reasoning"] = model_reasoning - - # Append the processed assistant message to our internal logs - self.logs.append(log_entry) - - if response_message.tool_calls: - print(f"--- LLM requested tool execution ---") - for tool_call in response_message.tool_calls: - function_name = tool_call.function.name - function_args = json.loads(tool_call.function.arguments) - - function_to_call = self.available_functions.get(function_name) - function_output = function_to_call(**function_args) - - print(f"Executing: {function_name}({function_args}) -> {function_output}") - - tool_message = { - "tool_call_id": tool_call.id, - "role": "tool", - "content": str(function_output), + self.logs.append({ + "role": "assistant", + "content": content, + "reasoning": reasoning, + "tool_calls": None, + "response_metadata": { + "model": self.model, + "usage": response.usage } - messages.append(tool_message) - self.logs.append(tool_message) + }) + + # Continue to next turn - don't return + return None else: - print(f"\n--- Final LLM Response ---\n{response_message.content}") - return response_message.content - - def save_logs(self, output_dir="output"): + # Not first turn and no tool calls - this is a final response + assistant_message = { + "role": "assistant", + "content": content if content else None + } + messages.append(assistant_message) + + # Log entry + self.logs.append({ + "role": "assistant", + "content": content, + "reasoning": reasoning, + "tool_calls": None, + "response_metadata": { + "model": self.model, + "usage": response.usage + } + }) + + # print(f"\n--- Final LLM Response ---\n{content}") + return content + + def _execute_tool(self, messages: List[Dict], tool_call: Dict) -> Optional[str]: + """Execute a single tool call.""" + function_name = tool_call["function"]["name"] + function_args = json.loads(tool_call["function"]["arguments"]) + + function_to_call = self.available_functions.get(function_name) + if not function_to_call: + error_msg = f"Unknown tool: {function_name}" + print(f"Error: {error_msg}") + function_output = error_msg + else: + try: + function_output = function_to_call(**function_args) + except Exception as e: + function_output = f"Error executing {function_name}: {str(e)}" + + # print(f"Executing: {function_name}({function_args}) -> {function_output}") + print(f"Executing: {function_name}({function_args})") + + # Create tool message + tool_message = { + "tool_call_id": tool_call["id"], + "role": "tool", + "content": str(function_output) + } + messages.append(tool_message) + self.logs.append(tool_message) + + return str(function_output) + + def save_logs( + self, + output_dir: str = "output", + scenario: str = None, + oversight_level: str = None + ): + """Save conversation logs to a JSON file.""" timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") model_name_safe = self.model.replace("/", "_") - scenario_name_safe = self.scenario.replace("/", "_") + scenario_name = (scenario or self.scenario).replace("/", "_") + oversight = oversight_level or self.oversight_level - # New directory structure - model_output_dir = os.path.join(output_dir, model_name_safe) - os.makedirs(model_output_dir, exist_ok=True) + # Directory structure: output/{model}/{scenario}/{oversight}/ + base_dir = os.path.join(output_dir, model_name_safe, scenario_name, oversight) + os.makedirs(base_dir, exist_ok=True) - # New filename and run_id - filename_base = f"{scenario_name_safe}_{self.oversight_level}_{timestamp}" - run_id = f"{model_name_safe}/{filename_base}" - log_file = os.path.join(model_output_dir, f"{filename_base}.json") + # Filename: {timestamp}_run_id.json + filename_base = f"{timestamp}" + run_id = f"{model_name_safe}/{scenario_name}/{oversight}/{filename_base}" + log_file = os.path.join(base_dir, f"{filename_base}.json") log_data = { "run_id": run_id, "model": self.model, - "scenario": self.scenario, - "oversight_level": self.oversight_level, + "scenario": scenario or self.scenario, + "oversight_level": oversight, "user_prompt_type": self.user_prompt_type, "temperature": self.temperature, - "base_url": str(self.client.base_url), - "include_reasoning_in_context": INCLUDE_REASONING_IN_CONTEXT, - "extra_body_config": EXTRA_BODY, + "base_url": str(self.provider.provider_config.base_url) if self.provider else None, + "extra_body_config": self.provider.model_config.extra_body if self.provider else {}, "final_vfs_state": VFS.get_instance().fs, "total_tokens": self.total_tokens, "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, "conversation": self.logs } - + with open(log_file, "w") as f: json.dump(log_data, f, indent=4) print(f"\nLogs saved to {log_file}") + return log_file |