""" Agent Wrapper for Web Interface =============================== Wraps the LangChain agent for WebSocket streaming. """ import os import sys import asyncio import logging from pathlib import Path from typing import Optional, Callable, Any, List, Dict from queue import Queue # Add src directory to path for eurus package PROJECT_ROOT = Path(__file__).parent.parent sys.path.insert(0, str(PROJECT_ROOT)) sys.path.insert(0, str(PROJECT_ROOT / "src")) from dotenv import load_dotenv load_dotenv() from langchain_openai import ChatOpenAI from langchain.agents import create_agent # IMPORT FROM EURUS PACKAGE - SINGLE SOURCE OF TRUTH from eurus.config import CONFIG, AGENT_SYSTEM_PROMPT from eurus.retrieval import _arraylake_snippet from eurus.tools.era5 import _auto_detect_query_type from eurus.memory import get_memory, SmartConversationMemory # Singleton for datasets, per-session for chat from eurus.tools import get_all_tools from eurus.tools.repl import PythonREPLTool logger = logging.getLogger(__name__) class AgentSession: """ Manages a single agent session with streaming support. """ # Available models for the selector AVAILABLE_MODELS = [ {"id": "gpt-5.4", "label": "GPT-5.4", "provider": "openai"}, {"id": "gpt-4.1", "label": "GPT-4.1", "provider": "openai"}, {"id": "o3", "label": "o3", "provider": "openai"}, {"id": "gemini-3.1-pro-preview", "label": "Gemini 3.1 Pro", "provider": "google"}, ] def __init__(self, api_keys: Optional[Dict[str, str]] = None): self._agent = None self._repl_tool: Optional[PythonREPLTool] = None self._messages: List[Dict] = [] self._initialized = False self._api_keys = api_keys or {} self._current_model = CONFIG.model_name # Global singleton keeps the dataset cache (shared across sessions) self._memory = get_memory() # Per-session conversation memory — never touches other sessions self._conversation = SmartConversationMemory() # Queue for captured plots (thread-safe) self._plot_queue: Queue = Queue() self._initialize() def _initialize(self): """Initialize the agent and tools.""" logger.info("Initializing agent session...") # Resolve API keys: user-provided take priority over env vars openai_key = self._api_keys.get("openai_api_key") or os.environ.get("OPENAI_API_KEY") arraylake_key = self._api_keys.get("arraylake_api_key") or os.environ.get("ARRAYLAKE_API_KEY") hf_token = self._api_keys.get("hf_token") or os.environ.get("HF_TOKEN") if not arraylake_key: logger.warning("ARRAYLAKE_API_KEY not found") # SECURITY: Do NOT write user-provided keys to os.environ! # os.environ is process-global — leaks keys to other sessions on shared hosts (e.g. HF Spaces). # Instead, store in self and pass directly to tools that need them. self._resolved_keys = { "ARRAYLAKE_API_KEY": arraylake_key or "", "HF_TOKEN": hf_token or "", } if not openai_key: logger.error("OPENAI_API_KEY not found") return try: # Initialize REPL tool with working directory logger.info("Starting Python kernel...") self._repl_tool = PythonREPLTool(working_dir=os.getcwd()) # Inject session-specific keys into the REPL subprocess # (keeps them isolated from other sessions — no os.environ pollution) self._repl_tool.inject_env(self._resolved_keys) # Set up plot callback using the proper method def on_plot_captured(base64_data: str, filepath: str, code: str = ""): logger.info(f"Plot captured, adding to queue: {filepath}") self._plot_queue.put((base64_data, filepath, code)) self._repl_tool.set_plot_callback(on_plot_captured) logger.info("Plot callback registered") # Get ALL tools from centralized registry (no SCIENCE_TOOLS!) # Pass session-specific Arraylake key for isolation arraylake_key = self._resolved_keys.get("ARRAYLAKE_API_KEY") tools = get_all_tools( enable_routing=True, enable_guide=True, arraylake_api_key=arraylake_key or None, ) # Replace the default REPL with our configured one tools = [t for t in tools if t.name != "python_repl"] + [self._repl_tool] # Initialize LLM with resolved key logger.info("Connecting to LLM...") llm = ChatOpenAI( model=CONFIG.model_name, temperature=CONFIG.temperature, api_key=openai_key, ) # Use session-local memory for datasets (NOT global!) datasets = self._memory.list_datasets() enhanced_prompt = AGENT_SYSTEM_PROMPT if datasets != "No datasets in cache.": enhanced_prompt += f"\n\n## CACHED DATASETS\n{datasets}" # Create agent logger.info("Creating agent...") self._agent = create_agent( model=llm, tools=tools, system_prompt=enhanced_prompt, debug=False ) # FRESH conversation - no old messages! self._messages = [] self._initialized = True logger.info("Agent session initialized successfully") except Exception as e: logger.exception(f"Failed to initialize agent: {e}") self._initialized = False def is_ready(self) -> bool: """Check if the agent is ready.""" return self._initialized and self._agent is not None def get_current_model(self) -> str: """Return the current model name.""" return self._current_model def set_provider(self, model_id: str): """Switch the LLM model. Reinitializes the agent with the new model.""" openai_key = self._api_keys.get("openai_api_key") or os.environ.get("OPENAI_API_KEY") vertex_key = self._api_keys.get("vertex_api_key") or os.environ.get("vertex_api_key") # Determine provider from model id is_gemini = model_id.startswith("gemini") if is_gemini and not vertex_key: logger.error("Cannot switch to Gemini: no vertex_api_key in .env") return if not is_gemini and not openai_key: logger.error("Cannot switch model: no OPENAI_API_KEY") return logger.info(f"Switching model from {self._current_model} to {model_id}") self._current_model = model_id try: if is_gemini: from langchain_google_genai import ChatGoogleGenerativeAI llm = ChatGoogleGenerativeAI( model=model_id, temperature=CONFIG.temperature, api_key=vertex_key, vertexai=True, ) else: llm = ChatOpenAI( model=model_id, temperature=CONFIG.temperature, api_key=openai_key, ) tools = get_all_tools(enable_routing=True, enable_guide=True) tools = [t for t in tools if t.name != "python_repl"] + [self._repl_tool] datasets = self._memory.list_datasets() enhanced_prompt = AGENT_SYSTEM_PROMPT if datasets != "No datasets in cache.": enhanced_prompt += f"\n\n## CACHED DATASETS\n{datasets}" self._agent = create_agent( model=llm, tools=tools, system_prompt=enhanced_prompt, debug=False ) # Keep conversation intact — only reset tool calls self._messages = [] logger.info(f"Model switched to {model_id} successfully") except Exception as e: logger.exception(f"Failed to switch model: {e}") def reinitialize(self): """Retry initialization (e.g., after transient failure).""" logger.warning("Attempting agent reinitialization...") self._initialized = False self._agent = None self._initialize() def clear_messages(self): """Clear conversation messages.""" self._messages = [] def get_pending_plots(self) -> List[tuple]: """Get all pending plots from queue.""" plots = [] while not self._plot_queue.empty(): try: plots.append(self._plot_queue.get_nowait()) except Exception: break return plots async def process_message( self, user_message: str, stream_callback: Callable ) -> str: """ Process a user message and stream the response. """ if not self.is_ready(): # Try to reinitialize once before giving up logger.warning("Agent not ready, attempting reinitialization...") self.reinitialize() if not self.is_ready(): raise RuntimeError("Agent not initialized") # Clear any old plots from queue self.get_pending_plots() # Add user message to history (session-local memory) self._conversation.add_message("user", user_message) self._messages.append({"role": "user", "content": user_message}) try: # Send status: analyzing await stream_callback("status", "🔍 Analyzing your request...") await asyncio.sleep(0.3) # Invoke the agent in executor (20 iterations max to save tokens) config = {"recursion_limit": 20} # Stream status updates while agent is working await stream_callback("status", "🤖 Processing with AI...") # Save message state before invoke (protect against corruption) messages_backup = list(self._messages) result = await asyncio.get_event_loop().run_in_executor( None, lambda: self._agent.invoke({"messages": self._messages}, config=config) ) # Only scan NEW messages from this turn prev_count = len(self._messages) self._messages = result["messages"] new_messages = self._messages[prev_count:] # Parse NEW messages to show tool calls made tool_calls_made = [] for msg in new_messages: if hasattr(msg, 'tool_calls') and msg.tool_calls: for tc in msg.tool_calls: tool_name = tc.get('name', 'unknown') if tool_name not in tool_calls_made: tool_calls_made.append(tool_name) if tool_calls_made: tools_str = ", ".join(tool_calls_made) await stream_callback("status", f"🛠️ Used tools: {tools_str}") await asyncio.sleep(0.5) # Collect Arraylake snippet from NEW messages only # Only emit ONE snippet per unique (variable, region) — skip failed calls arraylake_snippets = [] seen_snippet_keys = set() for i, msg in enumerate(new_messages): if hasattr(msg, 'tool_calls') and msg.tool_calls: for tc in msg.tool_calls: if tc.get('name') == 'retrieve_era5_data': # Check if tool call succeeded by looking at the next message # (ToolMessage with same tool_call_id) tc_id = tc.get('id', '') succeeded = True for later_msg in new_messages[i+1:]: if (hasattr(later_msg, 'tool_call_id') and later_msg.tool_call_id == tc_id): content = getattr(later_msg, 'content', '') or '' if any(kw in content.lower() for kw in ['error', 'failed', 'exception', 'limit', 'exceeded', 'rejected', 'too large']): succeeded = False break if not succeeded: continue args = tc.get('args', {}) # Dedup key: variable + rounded region dedup_key = ( args.get('variable_id', 'sst'), round(args.get('min_latitude', -90)), round(args.get('max_latitude', 90)), round(args.get('min_longitude', 0)), round(args.get('max_longitude', 360)), ) if dedup_key in seen_snippet_keys: continue seen_snippet_keys.add(dedup_key) arraylake_snippets.append(_arraylake_snippet( variable=args.get('variable_id', 'sst'), query_type=_auto_detect_query_type( start_date=args.get('start_date', ''), end_date=args.get('end_date', ''), min_lat=args.get('min_latitude', -90), max_lat=args.get('max_latitude', 90), min_lon=args.get('min_longitude', 0), max_lon=args.get('max_longitude', 360), ), start_date=args.get('start_date', ''), end_date=args.get('end_date', ''), min_lat=args.get('min_latitude', -90), max_lat=args.get('max_latitude', 90), min_lon=args.get('min_longitude', 0), max_lon=args.get('max_longitude', 360), )) # Extract response last_message = self._messages[-1] if hasattr(last_message, 'content') and last_message.content: raw_content = last_message.content # Gemini can return content as a list of content blocks if isinstance(raw_content, list): # Extract text from each block parts = [] for block in raw_content: if isinstance(block, str): parts.append(block) elif isinstance(block, dict) and block.get('text'): parts.append(block['text']) elif hasattr(block, 'text'): parts.append(block.text) response_text = "\n".join(parts) if parts else str(raw_content) else: response_text = str(raw_content) elif isinstance(last_message, dict) and last_message.get('content'): response_text = str(last_message['content']) else: response_text = str(last_message) # Send status: generating response await stream_callback("status", "✍️ Generating response...") await asyncio.sleep(0.2) # Stream the response in chunks chunk_size = 50 for i in range(0, len(response_text), chunk_size): chunk = response_text[i:i + chunk_size] await stream_callback("chunk", chunk) await asyncio.sleep(0.01) # Send any captured media (plots and videos) plots = self.get_pending_plots() # NOTE: Only use session-specific _plot_queue, NOT shared folder scan (privacy!) if plots: await stream_callback("status", f"📊 Rendering {len(plots)} visualization(s)...") await asyncio.sleep(0.3) logger.info(f"Sending {len(plots)} media items to client") for plot_data in plots: base64_data, filepath = plot_data[0], plot_data[1] code = plot_data[2] if len(plot_data) > 2 else "" # Determine if this is a video or image ext = filepath.lower().split('.')[-1] if filepath else '' if ext in ('gif',): await stream_callback("video", "", data=base64_data, path=filepath, mimetype="image/gif") elif ext in ('webm',): await stream_callback("video", "", data=base64_data, path=filepath, mimetype="video/webm") elif ext in ('mp4',): await stream_callback("video", "", data=base64_data, path=filepath, mimetype="video/mp4") else: # Default to plot (png, jpg, etc.) await stream_callback("plot", "", data=base64_data, path=filepath, code=code) # Send Arraylake snippets AFTER response + plots exist in DOM for snippet in arraylake_snippets: await stream_callback("arraylake_snippet", snippet) # Save to memory self._conversation.add_message("assistant", response_text) return response_text except Exception as e: # Restore clean message state to prevent corruption on next call self._messages = messages_backup logger.exception(f"Error processing message: {e}") raise def close(self): """Clean up resources.""" logger.info("Closing agent session...") if self._repl_tool: try: self._repl_tool.close() except Exception as e: logger.error(f"Error closing REPL: {e}") # Per-connection sessions (NOT global singleton!) # Key: unique connection ID, Value: AgentSession _sessions: Dict[str, AgentSession] = {} def create_session(connection_id: str, api_keys: Optional[Dict[str, str]] = None) -> AgentSession: """Create a new session for a connection (reuses if already ready).""" if connection_id in _sessions: existing = _sessions[connection_id] if existing.is_ready(): logger.info(f"Reusing existing ready session for: {connection_id}") return existing # Close broken session before replacing existing.close() session = AgentSession(api_keys=api_keys) _sessions[connection_id] = session logger.info(f"Created session for connection: {connection_id}") return session def get_session(connection_id: str) -> Optional[AgentSession]: """Get session for a connection.""" return _sessions.get(connection_id) def close_session(connection_id: str): """Close and remove session for a connection.""" if connection_id in _sessions: _sessions[connection_id].close() del _sessions[connection_id] logger.info(f"Closed session for connection: {connection_id}") # DEPRECATED: Keep for backward compatibility during migration def get_agent_session() -> AgentSession: """DEPRECATED: Use create_session/get_session with connection_id instead.""" logger.warning("get_agent_session() is deprecated - use create_session(connection_id)") # Create default session for CLI/testing if "_default" not in _sessions: _sessions["_default"] = AgentSession() return _sessions["_default"] def shutdown_agent_session(): """Shutdown all agent sessions.""" count = len(_sessions) for conn_id in list(_sessions.keys()): close_session(conn_id) logger.info(f"Shutdown {count} sessions")