path: root/collaborativeagents/adapters/reflection_adapter.py

"""
Reflection Adapter - vLLM-based implementation using original CollaborativeAgents prompts.

This implements the "reflection" baseline from the MULTISESSIONCOLLAB paper:
- After each session, agent reflects on interaction to update memory (agent_notes)
- Memory is provided to agent at start of subsequent sessions
- Uses session-level reflection + persistent memory
- Uses LLM-based retrieval (proper_scaffolding) to prevent context overflow

Now uses vLLM for fast inference instead of local transformers.
Uses EXACT prompts from the original CollaborativeAgents paper for fairness.
"""

import sys
from pathlib import Path
from typing import Optional, List, Dict, Any
from json_repair import repair_json

# Add parent for utils import
sys.path.insert(0, str(Path(__file__).parent.parent))
from utils.vllm_client import VLLMClient, VLLMConfig

# Import ORIGINAL prompts from CollaborativeAgents for fair reproduction
sys.path.insert(0, str(Path(__file__).parent.parent / "collaborativeagents"))
from collaborativeagents.prompts import (
    reflective_agent_system_prompt_no_json,
    update_agent_notes_prompt,
    proper_scaffolding_prompt,
)
from collaborativeagents.utils import get_conversation_string

# Default vLLM URL (agent server on port 8003)
DEFAULT_VLLM_URL = "http://localhost:8003/v1"


class ReflectionAdapter:
    """
    Adapter for the Reflection baseline from MULTISESSIONCOLLAB.

    Uses vLLM for fast inference with:
    - agent_notes: Persistent memory updated via session-level reflection
    - Memory retrieval at each turn via proper_scaffolding (when notes are long)

    Uses ORIGINAL CollaborativeAgents prompts for fair benchmark reproduction.
    """

    def __init__(
        self,
        model_name: str = None,  # Ignored - vLLM auto-discovers model
        device_assignment: dict = None,  # Ignored - vLLM handles GPU
        api_base: str = None,  # vLLM server URL
        api_key: str = None,   # Ignored
        with_scaffolding: bool = True,
        with_proper_scaffolding: bool = True,  # Enable LLM-based retrieval by default
        vllm_url: str = None,  # vLLM server URL
        max_new_tokens: int = 2048,  # Match original CollaborativeAgents setting
    ):
        self.vllm_url = vllm_url or api_base or DEFAULT_VLLM_URL
        self.with_scaffolding = with_scaffolding
        self.with_proper_scaffolding = with_proper_scaffolding
        self.max_new_tokens = max_new_tokens

        # Per-user memory storage
        self._user_notes: Dict[str, str] = {}
        self._current_user_id: Optional[str] = None
        self._conversation_history: List[Dict[str, str]] = []

        # vLLM client (initialized lazily)
        self._client: Optional[VLLMClient] = None
        self._initialized = False

    def initialize(self):
        """Initialize the adapter (connects to vLLM server)."""
        if self._initialized:
            return

        print(f"[ReflectionAdapter] Connecting to vLLM server at {self.vllm_url}...")
        print(f"[ReflectionAdapter] Using proper_scaffolding={self.with_proper_scaffolding}")

        # Retry connection with exponential backoff
        import time
        max_retries = 30
        for attempt in range(max_retries):
            try:
                self._client = VLLMClient(base_url=self.vllm_url)
                if self._client.health_check():
                    break
            except Exception as e:
                pass

            if attempt < max_retries - 1:
                wait_time = min(2 ** attempt * 0.5, 10)  # 0.5, 1, 2, 4, 8, 10, 10...
                time.sleep(wait_time)
        else:
            raise RuntimeError(f"vLLM server not responding at {self.vllm_url} after {max_retries} retries")

        self._initialized = True
        print(f"[ReflectionAdapter] Connected to vLLM (model: {self._client.config.model})")

    def _generate(self, messages: List[Dict[str, str]], max_new_tokens: int = 1024) -> str:
        """Generate response using vLLM server."""
        if not self._initialized:
            self.initialize()

        result = self._client.chat(
            messages=messages,
            max_tokens=max_new_tokens,
            temperature=1.0,  # Match original CollaborativeAgents setting
            top_p=0.9,
        )

        return result["content"]

    def _prepend_notes_to_conversation(self, conversation, notes_text):
        """Prepend notes to the first message of a conversation copy."""
        if conversation:
            conversation = list(conversation)
            conversation[0] = dict(conversation[0])
            conversation[0]["content"] = (
                f"Remember, you have been taking notes throughout past conversations "
                f"about user preferences. Use whatever is relevant in these notes to "
                f"guide your response:\n{notes_text}\n\n" + conversation[0]["content"]
            )
        return conversation

    def _add_scaffolding_to_conversation(
        self,
        conversation: List[Dict[str, str]],
        agent_notes: str
    ) -> List[Dict[str, str]]:
        """
        Add scaffolding (memory notes) to conversation.
        Sequential version - for non-batch use only.
        """
        if not self.with_proper_scaffolding:
            return self._prepend_notes_to_conversation(conversation, agent_notes)
        else:
            prompt = self.get_scaffolding_prompt(conversation, agent_notes)
            if prompt is None:
                return self._prepend_notes_to_conversation(conversation, agent_notes)
            response = self._generate([{"role": "user", "content": prompt}], max_new_tokens=512)
            return self.apply_scaffolding_response(conversation, agent_notes, response)

    def get_scaffolding_prompt(self, conversation, agent_notes):
        """Build the scaffolding prompt for batch processing. Returns None if no scaffolding needed."""
        if not self.with_proper_scaffolding:
            return None
        conversation_str = get_conversation_string(conversation)
        return proper_scaffolding_prompt.format(
            conversation_history=conversation_str,
            complete_agent_notes=agent_notes
        )

    def apply_scaffolding_response(self, conversation, agent_notes, response):
        """Apply a scaffolding LLM response to the conversation."""
        try:
            parsed = repair_json(response, return_objects=True)
            if "relevant_notes" in parsed:
                notes_text = parsed["relevant_notes"]
                if conversation:
                    conversation = list(conversation)
                    conversation[0] = dict(conversation[0])
                    conversation[0]["content"] = (
                        f"Remember, you have been taking notes throughout past conversations "
                        f"about user preferences. Use these notes to guide your response:\n"
                        f"{notes_text}\n\n" + conversation[0]["content"]
                    )
                return conversation
        except Exception:
            pass
        # Fallback: use all notes
        return self._prepend_notes_to_conversation(conversation, agent_notes)

    def get_note_update_prompt(self, user_id=None):
        """Build the note-update prompt for batch processing. Returns (messages, user_id) or None."""
        uid = user_id or self._current_user_id
        if not uid or not self._conversation_history:
            return None
        current_notes = self._user_notes.get(uid, "No notes yet.")
        conversation_str = get_conversation_string(self._conversation_history)
        prompt = update_agent_notes_prompt.format(
            agent_notes=current_notes,
            conversation_str=conversation_str
        )
        return [{"role": "user", "content": prompt}]

    def apply_note_update_response(self, response, user_id=None):
        """Apply a note-update LLM response."""
        uid = user_id or self._current_user_id
        if not uid:
            return
        try:
            # 8B model: use raw response directly
            updated_notes = response.strip()
            if updated_notes:
                old_len = len(self._user_notes.get(uid, ""))
                self._user_notes[uid] = updated_notes
        except Exception:
            pass

    def start_session(self, user_id: str, user_profile: dict = None):
        """Start a new session for a user."""
        if not self._initialized:
            self.initialize()

        self._current_user_id = user_id
        self._conversation_history = []

    def generate_response(
        self,
        query: str,
        conversation_history: List[Dict[str, str]] = None
    ) -> Dict[str, Any]:
        """Generate a response using the reflection agent."""
        if not self._initialized:
            self.initialize()

        # Add user query to history
        self._conversation_history.append({"role": "user", "content": query})

        # Get current notes for this user
        agent_notes = self._user_notes.get(self._current_user_id, "No notes yet about this user.")

        # Build conversation with scaffolding (uses proper_scaffolding if enabled)
        if self.with_scaffolding and agent_notes != "No notes yet about this user.":
            conversation_with_notes = self._add_scaffolding_to_conversation(
                self._conversation_history, agent_notes
            )
        else:
            conversation_with_notes = self._conversation_history

        # Build system prompt using ORIGINAL CollaborativeAgents prompt
        # Note: For no_json mode, we don't include agent_notes in system prompt
        # because they're added via scaffolding to the conversation
        system_prompt = reflective_agent_system_prompt_no_json.format(
            agent_notes=agent_notes if not self.with_scaffolding else "See notes in conversation."
        )

        # Build messages for generation
        messages = [{"role": "system", "content": system_prompt}]
        messages.extend(conversation_with_notes)

        # Generate response
        response_text = self._generate(messages, max_new_tokens=self.max_new_tokens)

        self._conversation_history.append({"role": "assistant", "content": response_text})

        return {
            "response": response_text,
            "reasoning": "",
            "debug": {"agent_notes": agent_notes, "proper_scaffolding": self.with_proper_scaffolding}
        }

    def prepare_prompt(
        self,
        query: str,
        conversation_history: List[Dict[str, str]] = None
    ) -> tuple:
        """
        Prepare prompt for batch processing without calling main LLM.

        Note: This may still call LLM for scaffolding (memory retrieval),
        but the main generation is deferred for batching.

        Args:
            query: Current user query
            conversation_history: Previous conversation

        Returns:
            Tuple of (messages, context) for batch processing
        """
        if not self._initialized:
            self.initialize()

        # Add user query to history
        self._conversation_history.append({"role": "user", "content": query})

        # Get current notes for this user
        agent_notes = self._user_notes.get(self._current_user_id, "No notes yet about this user.")

        # Store for batch scaffolding (prepare_prompt may be called after batch scaffolding)
        self._pending_agent_notes = agent_notes
        self._pending_scaffolded = False

        # Build conversation with scaffolding
        if self.with_scaffolding and agent_notes != "No notes yet about this user.":
            if hasattr(self, '_scaffolding_result') and self._scaffolding_result is not None:
                # Use pre-computed batch scaffolding result
                conversation_with_notes = self.apply_scaffolding_response(
                    list(self._conversation_history), agent_notes, self._scaffolding_result)
                self._scaffolding_result = None
                self._pending_scaffolded = True
            elif not self.with_proper_scaffolding:
                conversation_with_notes = self._prepend_notes_to_conversation(
                    self._conversation_history, agent_notes)
                self._pending_scaffolded = True
            else:
                # Sequential fallback - should not happen in batch mode
                conversation_with_notes = self._add_scaffolding_to_conversation(
                    self._conversation_history, agent_notes)
                self._pending_scaffolded = True
        else:
            conversation_with_notes = self._conversation_history

        # Build system prompt using ORIGINAL CollaborativeAgents prompt
        system_prompt = reflective_agent_system_prompt_no_json.format(
            agent_notes=agent_notes if not self.with_scaffolding else "See notes in conversation."
        )

        # Build messages for generation
        messages = [{"role": "system", "content": system_prompt}]
        messages.extend(conversation_with_notes)

        # Context for post-processing
        ctx = {
            "agent_notes": agent_notes,
        }

        return messages, ctx

    def process_response(
        self,
        response: str,
        context: dict
    ) -> Dict[str, Any]:
        """
        Process LLM response after batch call.

        Args:
            response: LLM response text
            context: Context dict from prepare_prompt()

        Returns:
            Dict with 'response', 'reasoning', and debug info
        """
        self._conversation_history.append({"role": "assistant", "content": response})

        return {
            "response": response,
            "reasoning": "",
            "debug": {
                "agent_notes": context["agent_notes"],
                "proper_scaffolding": self.with_proper_scaffolding
            }
        }

    def _update_agent_notes(self, agent_notes: str, conversation: List[Dict[str, str]]) -> Optional[Dict]:
        """
        Update agent notes using ORIGINAL CollaborativeAgents logic.

        For 8B models (no_json): Use raw response directly
        For JSON models: Check for required keys, retry up to num_retries times
        Returns None if all retries fail (keeps old notes)
        """
        conversation_str = get_conversation_string(conversation)
        formatted_prompt = update_agent_notes_prompt.format(
            agent_notes=agent_notes,
            conversation_str=conversation_str
        )

        num_retries = 10
        no_json = True  # 8B model

        for attempt in range(num_retries):
            try:
                messages = [{"role": "user", "content": formatted_prompt}]
                response = self._generate(messages, max_new_tokens=512)

                # For 8B models (no_json=True): use raw response directly
                if no_json:
                    return {"agent_notes": response}

                # For JSON models: parse and check keys
                processed_response = repair_json(response, return_objects=True)
                missing_keys = [k for k in ["user_preferences_reasoning", "agent_notes"] if k not in processed_response]

                if missing_keys:
                    print(f"[ReflectionAdapter] Missing keys: {missing_keys}, attempt {attempt + 1}")
                    continue

                return processed_response

            except Exception as e:
                print(f"[ReflectionAdapter] Failed to update agent notes: {e}")

        return None  # All retries failed, keep old notes

    def end_session(self, task_success: bool = False, skip_note_update: bool = False) -> Dict[str, Any]:
        """
        End session and update agent notes via reflection.

        Args:
            task_success: Whether the task was completed successfully
            skip_note_update: If True, skip note update (already done via batch)
        """
        if not self._current_user_id:
            return {}

        # Update notes via session-level reflection (skip if batch already did it)
        if not skip_note_update and len(self._conversation_history) > 0:
            result = self._update_agent_notes(
                self._user_notes.get(self._current_user_id, "No notes yet."),
                self._conversation_history)
            if result is not None and "agent_notes" in result:
                self._user_notes[self._current_user_id] = result["agent_notes"]

        return {
            "turns": len(self._conversation_history),
            "task_success": task_success,
            "notes_updated": True,
        }

    def reset_user(self, user_id: str):
        """Reset all memory for a user."""
        if user_id in self._user_notes:
            del self._user_notes[user_id]

    def __call__(
        self,
        messages: List[Dict[str, str]],
        user_profile: dict = None,
        **kwargs
    ) -> str:
        """Callable interface for ConversationGenerator compatibility."""
        if not messages:
            return "How can I help you?"

        last_user_msg = None
        for msg in reversed(messages):
            if msg["role"] == "user":
                last_user_msg = msg["content"]
                break

        if last_user_msg is None:
            return "How can I help you?"

        result = self.generate_response(last_user_msg, messages)
        return result["response"]