rag_optimizer_environment.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Rag Optimizer Environment Implementation.
The agent acts as a Data Engineer to un-block a broken RAG pipeline.
"""

from uuid import uuid4
from typing import Dict, Any, List

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

# Import scikit-learn for our Grader
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

try:
    from models import RagOptimizerAction, RagOptimizerObservation
except ImportError:
    from models import RagOptimizerAction, RagOptimizerObservation


class RagOptimizerEnvironment(Environment):
    """
    RAG Optimizer Engine.
    Maintains a simulated Knowledge Base and grades it using TF-IDF.
    """

    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self):
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self.kb = self._get_initial_kb()
        
        # Hidden test suite for the grader
        self.test_suite = [
            {
                "query": "What is the current 2024 price for standard?",
                "target_concept": "750/mo"
            },
            {
                "query": "What is the refund policy for enterprise?",
                "target_concept": "Refunds are not permitted"
            },
            {
                "query": "UI issues frontend CSS missing button",
                "target_concept": "frontend team fixed the button"
            },
            {
                "query": "How long does shipping take to branch offices?",
                "target_concept": "5-7 business days"
            },
            {
                "query": "What months do parking passes need to be renewed?",
                "target_concept": "March"
            },
            {
                "query": "What holidays are we off in 2024?",
                "target_concept": "July 4"
            }
        ]

    def _get_initial_kb(self) -> Dict[str, Dict]:
        """Returns a fresh copy of the initial messy knowledge base."""
        return {
            "doc_pricing_legacy": {
                "text": "Pricing for 2021: Enterprise tier is $1000/mo. Standard is $500/mo. All plans include 10 users.",
                "metadata": {"type": "pricing"}
            },
            "doc_pricing_current_v2": {
                "text": "Current Pricing 2024: Enterprise is $1500/mo. Standard is $750/mo. Refunds are not permitted on the enterprise tier.",
                "metadata": {}
            },
            "doc_shipping_policy": {
                "text": "All internal shipments to remote branch offices take 5-7 business days. Overnight shipping is only available for C-suite.",
                "metadata": {"department": "logistics"}
            },
            "doc_messy_support_ticket_1": {
                "text": "User complained the button disappeared on the frontend. Another user said the database latency was high. The frontend team fixed the button by updating CSS.",
                "metadata": {}
            },
            "doc_messy_support_ticket_2": {
                "text": "Email integration is failing with error 401 Unauthorized. The API key was rotated on Tuesday.",
                "metadata": {}
            },
            "doc_monolithic_onboarding": {
                "text": "Welcome to the company! Here are some rules. 1) VPN access requires DUO. 2) The cafetaria opens at 8 AM. 3) For HR issues, email hr@company.com. 4) The 2024 holiday schedule includes Dec 25, Jan 1, and July 4. 5) Parking passes must be renewed annually in March.",
                "metadata": {}
            },
            # Add distractor files
            **{f"doc_distractor_hr_{i}": {"text": f"This is an old HR policy document regarding {['pto', 'sick leave', 'travel', 'expenses'][i%4]} from 201{i%10}.", "metadata":{}} for i in range(10)},
            **{f"doc_distractor_eng_{i}": {"text": f"Engineering architecture decision record {i}. We decided to use {['React', 'Postgres', 'Redis', 'Kafka'][i%4]} because of scaling concerns.", "metadata":{}} for i in range(10)},
            **{f"doc_distractor_random_{i}": {"text": f"Weekly team update notes. Nothing important here, just discussed the weather and the upcoming launch {i}.", "metadata":{}} for i in range(10)},
        }

    def _get_kb_summary(self) -> Dict[str, Dict]:
        """Returns a summary of the KB for the observation."""
        summary = {}
        for k, v in self.kb.items():
            summary[k] = {"metadata": v.get("metadata", {}), "length": len(v.get("text", ""))}
        return summary

    def reset(self) -> RagOptimizerObservation:
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self.kb = self._get_initial_kb()
        return RagOptimizerObservation(
            message="RagOptimizerEnv Initialized. You have messy chunks in the KB. Resolve conflicts, add metadata tags to short tickets, and splinter monolithic files to win.",
            current_docs=self._get_kb_summary(),
            done=False,
            reward=self._evaluate_kb()
        )

    def _evaluate_kb(self) -> float:
        """The Grader: Evaluates the agent's current KB using TF-IDF."""
        if not self.kb:
            return 0.01
            
        doc_texts = [doc["text"] for doc in self.kb.values()]
        
        vectorizer = TfidfVectorizer(stop_words='english')
        try:
            doc_vectors = vectorizer.fit_transform(doc_texts)
        except ValueError:
            return 0.01
            
        score = 0.0
        
        for case in self.test_suite:
            query_vec = vectorizer.transform([case["query"]])
            similarities = cosine_similarity(query_vec, doc_vectors)[0]
            
            # Get top 3
            top_k_indices = similarities.argsort()[-3:][::-1]
            
            found = False
            for idx in top_k_indices:
                if similarities[idx] > 0.01:
                    if case["target_concept"].lower() in doc_texts[idx].lower():
                        found = True
                        break
            if found:
                score += 1.0
                
        return max(0.01, min(0.99, float(score / len(self.test_suite))))

    def step(self, action: RagOptimizerAction) -> RagOptimizerObservation:  # type: ignore[override]
        self._state.step_count += 1
        
        msg = ""
        done = False
        reward = 0.01
        
        try:
            if action.action_type == "read_document":
                if action.doc_id in self.kb:
                    msg = f"Content of {action.doc_id}: {self.kb[action.doc_id]['text']}"
                else:
                    msg = f"Error: doc_id {action.doc_id} not found."
            
            elif action.action_type == "delete_document":
                if action.doc_id in self.kb:
                    del self.kb[action.doc_id]
                    msg = f"Deleted {action.doc_id}."
                else:
                    msg = f"Error: doc_id {action.doc_id} not found."
                    
            elif action.action_type == "update_document":
                if not action.doc_id or not action.text:
                    msg = "Error: doc_id and text required for update_document."
                else:
                    if action.doc_id not in self.kb:
                        self.kb[action.doc_id] = {"text": "", "metadata": {}}
                    self.kb[action.doc_id]["text"] = action.text
                    msg = f"Updated text for {action.doc_id}."
                    
            elif action.action_type == "add_metadata":
                if not action.doc_id or not action.metadata_key or not action.metadata_value:
                    msg = "Error: doc_id, metadata_key, and metadata_value required."
                else:
                    if action.doc_id not in self.kb:
                        msg = f"Error: doc_id {action.doc_id} not found."
                    else:
                        self.kb[action.doc_id]["metadata"][action.metadata_key] = action.metadata_value
                        msg = f"Added metadata to {action.doc_id}."
                        
            elif action.action_type == "submit":
                done = True
                reward = self._evaluate_kb()
                msg = f"Evaluation complete. Final reward: {reward:.2f}"
                
        except Exception as e:
            msg = f"Action failed: {str(e)}"

        if not done:
            reward = self._evaluate_kb()

        return RagOptimizerObservation(
            message=msg,
            current_docs=self._get_kb_summary(),
            done=done,
            reward=reward,
        )

    @property
    def state(self) -> State:
        return self._state