""" User Simulator Simulates a user with specific preferences who: 1. Presents problems to the agent 2. Checks if agent responses satisfy their preferences 3. Enforces preferences when violated 4. Tracks draft answer and decides when to terminate """ import json import os from dataclasses import dataclass, field from typing import List, Dict, Any, Optional from ..profiles.generator import UserProfile from ..preference_bank.schemas import PreferenceItem # User simulator system prompt template USER_SYSTEM_PROMPT = """You are simulating a user who is collaborating with an AI assistant to solve a problem. You have specific preferences about how the assistant should respond. # Problem to Solve {task_description} {problem} Note: The assistant cannot see this problem description directly. You need to communicate with them. # Your Persona {persona} # Your Preferences (Grouped by Topic) {preferences_grouped} # Preference Enforcement Rules - For each assistant response, check which of YOUR preferences are RELEVANT to the current context - A preference is relevant if the assistant's response touches on that topic/condition - If a relevant preference is VIOLATED, you MUST enforce it before proceeding - Do NOT update your draft answer or proceed until violated preferences are fixed - Only check preferences that apply to the current response (e.g., coding preferences for code responses) # Draft Answer Management - Maintain a working draft answer to the problem - Start with "I don't know" - Update it based on helpful information from the assistant - Do NOT update if you're enforcing preferences # Conversation Guidelines - Be somewhat vague initially, let the assistant ask clarifying questions - Respond naturally like a real user - Do not copy the problem description directly # Termination Terminate when: - Your draft answer seems correct and complete - The assistant cannot help further When ready to terminate, include "TERMINATE" in your response. # Output Format (JSON) {{ "preference_checks": [ {{ "preference_id": str, "topic": str, "relevant": bool, "satisfied": bool or null, "violation_detail": str }} ], "any_violation": bool, "enforcement_needed": bool, "reasoning": str, "draft_answer": str, "should_terminate": bool, "response": str }} IMPORTANT: Only include preferences that are RELEVANT to the current assistant response in preference_checks. Output valid JSON only, no other text.""" @dataclass class PreferenceCheck: """Result of checking one preference.""" preference_id: str topic: str relevant: bool satisfied: Optional[bool] # None if not relevant violation_detail: str = "" @dataclass class UserSimulatorResponse: """Response from the user simulator.""" response: str # Text response to agent preference_checks: List[PreferenceCheck] # Checked preferences any_violation: bool # Any preference violated? enforcement_needed: bool # Need to enforce? draft_answer: str # Current draft answer should_terminate: bool # Should end conversation? reasoning: str # Internal reasoning raw_output: Dict[str, Any] = field(default_factory=dict) class UserSimulator: """ Simulates a user with preferences interacting with an agent. """ def __init__( self, model_name: str = "Llama-3.3-70B-Instruct", api_base: Optional[str] = None, api_key: Optional[str] = None, temperature: float = 0.8, max_tokens: int = 2048, ): self.model_name = model_name self.api_base = api_base or os.getenv("USER_SIM_API_BASE", "http://localhost:8004/v1") self.api_key = api_key or os.getenv("USER_SIM_API_KEY", "EMPTY") self.temperature = temperature self.max_tokens = max_tokens # Current session state self._profile: Optional[UserProfile] = None self._task_description: str = "" self._problem: str = "" self._solution: str = "" self._init_client() def _init_client(self): """Initialize OpenAI client.""" try: import openai self.client = openai.OpenAI( base_url=self.api_base, api_key=self.api_key, ) except Exception as e: print(f"Warning: Could not initialize OpenAI client for user simulator: {e}") self.client = None def setup( self, profile: UserProfile, task_description: str, problem: str, solution: str = "", ): """ Set up the simulator for a new task. Args: profile: User profile with preferences task_description: Description of the task type problem: The specific problem to solve solution: Ground truth solution (for evaluation) """ self._profile = profile self._task_description = task_description self._problem = problem self._solution = solution def _build_system_prompt(self) -> str: """Build the system prompt with user profile and task.""" if self._profile is None: raise ValueError("User profile not set. Call setup() first.") return USER_SYSTEM_PROMPT.format( task_description=self._task_description, problem=self._problem, persona=self._profile.persona, preferences_grouped=self._profile.format_preferences_grouped(), ) def _parse_response(self, raw_text: str) -> UserSimulatorResponse: """Parse LLM output into structured response.""" try: # Try to extract JSON from response text = raw_text.strip() # Handle markdown code blocks if "```json" in text: text = text.split("```json")[1].split("```")[0] elif "```" in text: text = text.split("```")[1].split("```")[0] data = json.loads(text) # Parse preference checks pref_checks = [] for check in data.get("preference_checks", []): pref_checks.append(PreferenceCheck( preference_id=check.get("preference_id", ""), topic=check.get("topic", ""), relevant=check.get("relevant", False), satisfied=check.get("satisfied"), violation_detail=check.get("violation_detail", ""), )) return UserSimulatorResponse( response=data.get("response", ""), preference_checks=pref_checks, any_violation=data.get("any_violation", False), enforcement_needed=data.get("enforcement_needed", False), draft_answer=data.get("draft_answer", "I don't know"), should_terminate=data.get("should_terminate", False), reasoning=data.get("reasoning", ""), raw_output=data, ) except Exception as e: print(f"Error parsing user simulator response: {e}") print(f"Raw text: {raw_text[:500]}...") # Return a basic response return UserSimulatorResponse( response=raw_text if len(raw_text) < 500 else "Could you please continue?", preference_checks=[], any_violation=False, enforcement_needed=False, draft_answer="I don't know", should_terminate=False, reasoning="Parse error", raw_output={"error": str(e), "raw": raw_text}, ) def respond( self, conversation_history: List[Dict[str, str]], ) -> UserSimulatorResponse: """ Generate user response based on conversation. Args: conversation_history: List of {"role": "user/assistant", "content": "..."} Returns: UserSimulatorResponse with user's reply and preference status """ if self._profile is None: raise ValueError("User profile not set. Call setup() first.") system_prompt = self._build_system_prompt() # Build messages - reverse roles (user simulator sees itself as user) messages = [{"role": "system", "content": system_prompt}] for msg in conversation_history: # Flip roles: agent's messages become user input to simulator if msg["role"] == "assistant": messages.append({"role": "user", "content": msg["content"]}) else: messages.append({"role": "assistant", "content": msg["content"]}) if self.client is None: # Fallback for testing return self._fallback_response(conversation_history) try: response = self.client.chat.completions.create( model=self.model_name, messages=messages, temperature=self.temperature, max_tokens=self.max_tokens, ) raw_text = response.choices[0].message.content return self._parse_response(raw_text) except Exception as e: print(f"Error calling user simulator LLM: {e}") return self._fallback_response(conversation_history) def _fallback_response( self, conversation_history: List[Dict[str, str]], ) -> UserSimulatorResponse: """Generate a simple fallback response for testing.""" num_turns = len([m for m in conversation_history if m["role"] == "assistant"]) if num_turns == 0: # First turn - present the problem response = f"Hi, I need help with this: {self._problem[:200]}..." elif num_turns < 3: response = "Thanks, that helps. Can you explain more?" else: response = "Got it, I think I understand now. TERMINATE" return UserSimulatorResponse( response=response, preference_checks=[], any_violation=False, enforcement_needed=False, draft_answer="Draft answer from fallback", should_terminate="TERMINATE" in response, reasoning="Fallback mode", raw_output={}, ) def get_solution(self) -> str: """Get the ground truth solution.""" return self._solution def get_profile(self) -> Optional[UserProfile]: """Get the current user profile.""" return self._profile