diff options
Diffstat (limited to 'src/agents/main/agent.py')
| -rw-r--r-- | src/agents/main/agent.py | 327 |
1 files changed, 0 insertions, 327 deletions
diff --git a/src/agents/main/agent.py b/src/agents/main/agent.py deleted file mode 100644 index 2b979fa..0000000 --- a/src/agents/main/agent.py +++ /dev/null @@ -1,327 +0,0 @@ -""" -Agent class that uses the provider abstraction layer. -Handles conversation loops, tool execution, and logging. -""" -import json -import os -import datetime -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: 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 - - @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) - return self.chat_loop(messages) - - 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: str): - """Send a message and get response.""" - self.logs.append({'role': 'user', 'content': user_input}) - messages = list(self.logs) - return self.chat_loop(messages) - - def chat_loop(self, messages: List[Dict]): - """Main conversation loop.""" - turn_count = 0 - while True: - llm_response = self.provider.call(messages, self.tools) - - # Update token counts - 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: - 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": 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": 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) - - self.logs.append({ - "role": "assistant", - "content": content, - "reasoning": reasoning, - "tool_calls": None, - "response_metadata": { - "model": self.model, - "usage": response.usage - } - }) - - # Continue to next turn - don't return - return None - else: - # 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 = (scenario or self.scenario).replace("/", "_") - oversight = oversight_level or self.oversight_level - - # 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) - - # 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": scenario or self.scenario, - "oversight_level": oversight, - "user_prompt_type": self.user_prompt_type, - "temperature": self.temperature, - "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 |