agent.py

import os
import time
from datetime import datetime, timezone
from typing import List, Dict, Any, Tuple, Union, Optional
from openai import OpenAI, APIError, RateLimitError, APIConnectionError

from config import Config, build_system_prompt, FALLBACK_MODEL, SECONDARY_FALLBACK_MODEL
from r2_manager import R2MemoryManager

class HermesAgent:
    def __init__(self):
        self.r2 = R2MemoryManager()
        self.state: Dict[str, Any] = {}
        self.memory: Dict[str, Any] = {}
        self.booted: bool = False
        self.boot_status_message: str = ""
        self.client: OpenAI = None
        self._current_key: str = ""
        self.active_model: str = Config.LLM_MODEL
        self._init_openai()

    def _init_openai(self):
        api_key = os.getenv("LLM_API_KEY") or Config.LLM_API_KEY
        if not api_key:
            print("[!] LLM_API_KEY environment variable is missing!")
        
        self._current_key = api_key
        self.client = OpenAI(
            api_key=api_key or "dummy-key",
            base_url=Config.LLM_BASE_URL
        )

    def _ensure_openai_client(self):
        latest_key = os.getenv("LLM_API_KEY") or Config.LLM_API_KEY
        if latest_key != self._current_key or self.client is None:
            self._init_openai()

    def set_active_model(self, model_name: str):
        if model_name:
            self.active_model = model_name

    def boot(self) -> str:
        """Execute boot sequence: reload state/memory from R2 and measure offline duration."""
        self.state = self.r2.load_state()
        self.memory = self.r2.load_memory()

        last_active_str = self.state.get("last_active")
        memory_entries = self.memory.get("entries", [])
        entry_count = len(memory_entries)

        offline_duration_info = "First boot"
        if last_active_str:
            try:
                last_dt = datetime.fromisoformat(last_active_str)
                now_dt = datetime.now(timezone.utc)
                diff = now_dt - last_dt
                hours, remainder = divmod(int(diff.total_seconds()), 3600)
                minutes, seconds = divmod(remainder, 60)
                offline_duration_info = f"{hours}h {minutes}m ago"
            except Exception:
                offline_duration_info = str(last_active_str)

        storage_mode = "Cloudflare R2" if self.r2.is_connected else "Stateless Local Mode"
        
        self.boot_status_message = (
            f"Online [{storage_mode}]. "
            f"Last active: {offline_duration_info}. "
            f"Memory loaded: {entry_count} entries."
        )

        self.state["session_count"] = self.state.get("session_count", 0) + 1
        self.state["status"] = "online"
        self.r2.save_state(self.state)
        self.r2.append_log(
            action="boot_sequence",
            result=f"Booted successfully. Mode: {storage_mode}",
            tokens_used=0
        )
        self.booted = True
        return self.boot_status_message

    def get_system_prompt(self) -> str:
        memory_status = "Connected (Cloudflare R2)" if self.r2.is_connected else "Stateless (Local Fallback)"
        memory_count = len(self.memory.get("entries", []))
        last_timestamp = str(self.state.get("last_active", "None"))
        return build_system_prompt(memory_status, memory_count, last_timestamp, model_name=self.active_model)

    def _single_llm_attempt(self, messages: List[Dict[str, str]], model_name: str) -> Tuple[Optional[str], int, Optional[str]]:
        """Perform a single model call. Returns (content, tokens, error_msg)."""
        try:
            response = self.client.chat.completions.create(
                model=model_name,
                messages=messages,
                temperature=0.7,
                max_tokens=2048
            )
            content = response.choices[0].message.content or ""
            tokens = response.usage.total_tokens if response.usage else 0
            return content, tokens, None
        except Exception as e:
            return None, 0, str(e)

    def call_llm_with_retry(self, messages: List[Dict[str, str]], model_override: str = None, max_retries: int = 3) -> Tuple[str, int]:
        """Call external LLM API with model fallback chain."""
        self._ensure_openai_client()
        active_key = os.getenv("LLM_API_KEY") or Config.LLM_API_KEY
        
        if not active_key:
            return "⚠️ Error: `LLM_API_KEY` is missing. Please add the `LLM_API_KEY` secret in Space Settings.", 0

        # Build fallback model chain
        models_to_try = []
        primary = model_override or self.active_model
        models_to_try.append(primary)
        
        if FALLBACK_MODEL not in models_to_try:
            models_to_try.append(FALLBACK_MODEL)
        if SECONDARY_FALLBACK_MODEL not in models_to_try:
            models_to_try.append(SECONDARY_FALLBACK_MODEL)

        last_error = ""
        for model in models_to_try:
            print(f"[*] Attempting LLM query with model: '{model}'...")
            for attempt in range(1, max_retries + 1):
                content, tokens, err = self._single_llm_attempt(messages, model)
                if err is None:
                    if model != primary:
                        print(f"[+] Fallback successful! Used model '{model}' instead of primary '{primary}'.")
                    return content, tokens

                last_error = err
                print(f"[!] Model '{model}' error (Attempt {attempt}/{max_retries}): {err}")
                if attempt < max_retries:
                    time.sleep(1)
            
            print(f"[!] Model '{model}' failed. Trying next model in fallback chain...")

        return f"LLM API Error across all fallback models: {last_error}", 0

    def chat(self, user_input: str, history: Any = None, model_name: str = None) -> str:
        """Process a user chat message, generate response, and persist state/memory."""
        if not self.booted:
            self.boot()

        if model_name:
            self.set_active_model(model_name)

        # Build messages payload
        messages = [{"role": "system", "content": self.get_system_prompt()}]

        # Inject long-term memory context if available
        memory_entries = self.memory.get("entries", [])
        if memory_entries:
            recent_memories = "\n".join([f"- {m['summary']}" for m in memory_entries[-10:]])
            messages.append({
                "role": "system",
                "content": f"## Long-Term Memory Summary Context:\n{recent_memories}"
            })

        # Append chat history safely handling tuples or list of dicts
        if history:
            for item in history:
                if isinstance(item, (list, tuple)) and len(item) == 2:
                    u_msg, a_msg = item[0], item[1]
                    if u_msg:
                        messages.append({"role": "user", "content": str(u_msg)})
                    if a_msg:
                        messages.append({"role": "assistant", "content": str(a_msg)})
                elif isinstance(item, dict):
                    role = item.get("role")
                    content = item.get("content")
                    if role and content and role in ("user", "assistant"):
                        messages.append({"role": role, "content": str(content)})

        # Append current user prompt
        messages.append({"role": "user", "content": str(user_input)})

        # Get response from LLM
        reply, tokens = self.call_llm_with_retry(messages, model_override=self.active_model)

        # Log & Update State
        self.r2.append_log(
            action="user_chat",
            result=f"Model: {self.active_model} | Query: {str(user_input)[:50]}...",
            tokens_used=tokens
        )
        
        # Save memory & check compaction asynchronously
        self._record_memory_entry(str(user_input), reply)

        return reply

    def _record_memory_entry(self, user_input: str, reply: str):
        """Append concise summary of conversation to memory and run compaction if > 500 KB."""
        now_str = datetime.now(timezone.utc).isoformat()
        
        summary_text = f"User asked: '{user_input[:80]}...'. Agent responded concisely."
        
        entry = {
            "timestamp": now_str,
            "user_query": user_input[:120],
            "summary": summary_text
        }
        
        if "entries" not in self.memory:
            self.memory["entries"] = []
        
        self.memory["entries"].append(entry)

        # Compact memory if exceeding threshold
        self._check_and_compact_memory()

        # Persist memory & state
        self.r2.save_memory(self.memory)
        self.r2.save_state(self.state)

    def _check_and_compact_memory(self):
        """If memory.json exceeds MAX_MEMORY_BYTES (500 KB), compact oldest 50% entries."""
        mem_size = self.r2.get_memory_size_bytes(self.memory)
        if mem_size <= Config.MAX_MEMORY_BYTES:
            return

        entries = self.memory.get("entries", [])
        if len(entries) < 4:
            return

        print(f"[*] Memory payload size ({mem_size} bytes) exceeds limit ({Config.MAX_MEMORY_BYTES} bytes). Compacting...")
        
        half_idx = len(entries) // 2
        old_entries = entries[:half_idx]
        recent_entries = entries[half_idx:]

        # Summarize old entries via LLM
        old_text = "\n".join([f"[{e['timestamp']}] {e['summary']}" for e in old_entries])
        prompt = [
            {"role": "system", "content": "You are a concise memory compaction assistant. Summarize the following historical conversation memories into a single paragraph under 150 words while preserving critical facts."},
            {"role": "user", "content": old_text}
        ]

        summary_para, _ = self.call_llm_with_retry(prompt, model_override=self.active_model, max_retries=2)
        
        compacted_entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "user_query": "COMPACTED_MEMORY_ARCHIVE",
            "summary": f"[COMPACTED ARCHIVE]: {summary_para}"
        }

        self.memory["entries"] = [compacted_entry] + recent_entries
        self.r2.append_log(
            action="memory_compaction",
            result=f"Compacted {len(old_entries)} old entries into single summary.",
            tokens_used=0
        )
        print("[+] Memory compaction complete.")