""" Chat Manager for LLM-based annotator assistance. Provides a singleton ChatManager that handles multi-turn conversations between annotators and an LLM, with context about the current annotation task. """ import logging import time from typing import Any, Dict, List, Optional from potato.ai.ai_endpoint import AIEndpointFactory, BaseAIEndpoint, AIEndpointConfigError logger = logging.getLogger(__name__) # Singleton instance _chat_manager: Optional["ChatManager"] = None DEFAULT_SYSTEM_PROMPT = """You are an annotation assistant for the task "{task_name}". Task description: {task_description} Available labels: {annotation_labels} The annotator is currently looking at this text: --- {instance_text} --- Help the annotator think through how to annotate this instance. You may: - Highlight relevant aspects of the text - Explain what to look for given the task description - Clarify label definitions if asked Do NOT tell the annotator which label to choose. Your role is to help them reason through the decision themselves.""" class ChatManager: """Manages LLM chat interactions for annotator assistance.""" def __init__(self, config: Dict[str, Any]): chat_config = config.get("chat_support", {}) self.enabled = chat_config.get("enabled", False) # UI settings ui_config = chat_config.get("ui", {}) self.title = ui_config.get("title", "Ask AI") self.placeholder = ui_config.get("placeholder", "Ask about this annotation...") self.sidebar_width = ui_config.get("sidebar_width", 380) self.max_history_per_instance = ui_config.get("max_history_per_instance", 50) # System prompt template prompt_config = chat_config.get("system_prompt", {}) self.system_prompt_template = prompt_config.get("template", DEFAULT_SYSTEM_PROMPT) # Extract task-level info for system prompt self.task_name = config.get("annotation_task_name", "Annotation Task") self.task_description = config.get("annotation_task_description", "") # Build label summary from annotation schemes labels = [] for scheme in config.get("annotation_schemes", []): name = scheme.get("name", "") scheme_labels = scheme.get("labels", []) if scheme_labels: labels.append(f"{name}: {', '.join(str(l) for l in scheme_labels)}") elif scheme.get("description"): labels.append(f"{name}: {scheme['description']}") self.annotation_labels = "; ".join(labels) if labels else "See task description" # Create the LLM endpoint self.endpoint: Optional[BaseAIEndpoint] = None if self.enabled: self._init_endpoint(config) def _init_endpoint(self, config: Dict[str, Any]): """Initialize the LLM endpoint from chat_support config.""" chat_config = config["chat_support"] # Build a config dict that AIEndpointFactory expects # (it looks for ai_support.enabled, ai_support.endpoint_type, ai_support.ai_config) endpoint_config = { "ai_support": { "enabled": True, "endpoint_type": chat_config.get("endpoint_type"), "ai_config": chat_config.get("ai_config", {}), } } try: self.endpoint = AIEndpointFactory.create_endpoint(endpoint_config) if self.endpoint: logger.info(f"Chat endpoint initialized: {chat_config.get('endpoint_type')}") else: logger.error("Chat endpoint creation returned None") self.enabled = False except AIEndpointConfigError as e: logger.error(f"Failed to initialize chat endpoint: {e}") self.enabled = False def build_system_prompt(self, instance_text: str, instance_id: str) -> str: """Build the system prompt with current instance context.""" try: return self.system_prompt_template.format( task_name=self.task_name, task_description=self.task_description, annotation_labels=self.annotation_labels, instance_text=instance_text, instance_id=instance_id, ) except KeyError as e: logger.warning(f"System prompt template variable not found: {e}, using default") return DEFAULT_SYSTEM_PROMPT.format( task_name=self.task_name, task_description=self.task_description, annotation_labels=self.annotation_labels, instance_text=instance_text, instance_id=instance_id, ) def send_message( self, user_message: str, instance_text: str, instance_id: str, history: List[Dict[str, str]], ) -> Dict[str, Any]: """ Send a user message and get an LLM response. Args: user_message: The annotator's message instance_text: Current instance text for context instance_id: Current instance ID history: Previous messages as list of {role, content} dicts Returns: Dict with 'content' (str) and 'response_time_ms' (int) """ if not self.endpoint: return {"content": "Chat support is not configured.", "response_time_ms": 0} system_prompt = self.build_system_prompt(instance_text, instance_id) # Build messages array: system + history + new user message messages = [{"role": "system", "content": system_prompt}] messages.extend(history) messages.append({"role": "user", "content": user_message}) start_time = time.time() try: response_text = self.endpoint.chat_query(messages) elapsed_ms = int((time.time() - start_time) * 1000) return {"content": response_text, "response_time_ms": elapsed_ms} except Exception as e: elapsed_ms = int((time.time() - start_time) * 1000) logger.error(f"Chat query failed: {e}") return { "content": "Sorry, I encountered an error. Please try again.", "response_time_ms": elapsed_ms, } def get_ui_config(self) -> Dict[str, Any]: """Return UI configuration for the frontend.""" return { "enabled": self.enabled, "title": self.title, "placeholder": self.placeholder, "sidebar_width": self.sidebar_width, "max_history_per_instance": self.max_history_per_instance, } def init_chat_manager(config: Dict[str, Any]) -> ChatManager: """Initialize the global ChatManager singleton.""" global _chat_manager _chat_manager = ChatManager(config) return _chat_manager def get_chat_manager() -> Optional[ChatManager]: """Get the ChatManager singleton. Returns None if not initialized.""" return _chat_manager def clear_chat_manager(): """Clear the ChatManager singleton. Used for testing.""" global _chat_manager _chat_manager = None