Upload folder using huggingface_hub

README.md

@@ -43,7 +43,7 @@ RagOptimizerEnv simulates the critical and computationally intensive role of an
 
 <br>
 <p align="center">
-  <img src="assets/Architecture_diagram.png" width="800" alt="System Architecture Diagram" />
+  <img src="assets/Architecture_diagram.jpeg" width="800" alt="System Architecture Diagram" />
 </p>
 <br>
 
@@ -109,7 +109,7 @@ Agents manipulate the embedding search space via a strictly defined remote schem
 Pydantic-typed JSON states returned immediately upon trajectory execution:
 1. `message`: Terminal I/O logs, error stack-traces, or `read_document` raw buffers.
 2. `current_docs`: A hierarchical mapping of the existing document index payload ensuring contextual grounding.
-3. `reward`: The live theoretical model convergence rate formulated between `0.0` and `1.0`.
+3. `reward`: The live theoretical model convergence rate formulated between `0.01` and `0.99`.
 
 ---
 
@@ -125,6 +125,33 @@ Evaluation is robust, empirical, and mathematically bounded:
 
 **Reward Yield:** `(Successful Vectors / Total Vector Payload)` providing high-density intermediate signals mapping continuously toward the `±1.0` upper bound.
 
+### Practical RAG FAQ (Important Clarifications)
+
+These are fantastic questions, and they cut right to the core of how real-world RAG systems operate.
+
+#### 1. If we just send the KB summary, is it just metadata optimization?
+Not exactly. While the agent sees the summary by default, it also has the `read_document` action.
+
+- When the agent uses `read_document("doc_monolithic")`, the environment returns the entire raw text of that document in the `message` field of the next observation.
+- The agent can read the text, identify that it is noisy or overloaded across multiple topics, and then use `update_document` to split or rewrite content into cleaner chunks.
+- This means the benchmark supports true content optimization, not just metadata operations.
+
+#### 2. Doesn't sending the entire KB to the embedding model blow up its context window?
+No. The grader does not send the entire KB as one giant prompt.
+
+- Embedding models operate per document, not as one monolithic concatenated input.
+- The environment encodes each document individually into vectors, then stores those vectors in an index-like structure for retrieval scoring.
+- Query-time evaluation compares a query vector against document vectors. It does not require processing the full KB in one context window.
+
+#### 3. Is the result just fetching the right document, not the exact answer?
+Exactly. This environment evaluates retrieval quality, the "R" in RAG.
+
+- If retrieval is wrong, generation quality collapses and hallucination risk rises.
+- The grader checks whether documents containing the target concept are ranked near the top for each control query.
+- Optimizing document quality, structure, and noise levels pushes the clean source document toward rank 1, enabling downstream generators to answer correctly.
+
+In practice, this environment trains an agent to behave like a retrieval-focused KB operator: maintain high signal quality so the retriever does not fail under noisy enterprise conditions.
+
 ---
 
 ## 4. Curriculum Topologies
@@ -135,9 +162,9 @@ The environment tests agents across three progressively demanding task distribut
 **The Vector Issue:** The base contains heavily overlapping parameters (competing versions of legacy and modern timeline protocols).
 **System Goal:** Autonomously survey the semantic differences, deduce the temporal conflict, and execute `delete_document` sweeps to purge vector hallucination triggers.
 
-### Level II: Ontological Tagging
-**The Vector Issue:** Textual segments representing technical telemetry are too dense and sparse on specific categorical keywords, resulting in low coordinate density for deterministic algorithms.
-**System Goal:** Execute inferential reading, deduce categorical bounds organically, and route exact programmatic `metadata` tags onto corresponding payloads to anchor the search vectors.
+### Level II: Signal Separation
+**The Vector Issue:** The KB contains overlapping incident narratives mixed with distractor engineering notes, causing retrieval ambiguity.
+**System Goal:** Deduplicate noisy/partial incident content so retrieval consistently surfaces the clean resolution document at rank 1.
 
 ### Level III: Syntactic Splintering (The Monolith)
 **The Vector Issue:** Extreme embedding decay caused by disparate conceptual structures compacted under a single referential document. This represents the well-known "PDF chunk wash-out" phenomenon.
@@ -156,7 +183,7 @@ Telemetry formats strictly adhere to high-velocity logging required for large-sc
 [STEP] step=1 action=read('doc_monolithic_onboarding') reward=0.33 done=false error=null
 [STEP] step=2 action=update('doc_vpn_policy') reward=0.50 done=false error=null
 ...
-[END] success=true steps=12 score=1.00 rewards=0.33,0.50,... 
+[END] success=true steps=12 score=0.99 rewards=0.33,0.50,... 
 ```
 
 ### Execution

inference.py

@@ -41,7 +41,7 @@ SYSTEM_PROMPT = """You are an automated Data Engineer managing an AI Knowledge B
 Your goal is to optimize the messy chunks of text in the database so that a TF-IDF Search Algorithm can find answers easily.
 You must resolve contradictions, categorize documents, and delete unnecessary documents.
 
-After each action you will receive a "current_reward" score (0.0 to 1.0) indicating how well the KB currently performs. Use this to guide your strategy.
+After each action you will receive a "current_reward" score (0.01 to 0.99) indicating how well the KB currently performs. Use this to guide your strategy.
 
 You have the following actions:
 - {"action_type": "read_document", "doc_id": "..."}
@@ -65,6 +65,34 @@ def format_action_str(action: RagOptimizerAction) -> str:
         return "submit()"
     return f"{action.action_type}()"
 
+# --- Reflexion (Long Term Memory) ---
+LESSONS_FILE = os.path.join(os.path.dirname(__file__), "memory", "lessons_learned.json")
+
+def load_lessons():
+    if os.path.exists(LESSONS_FILE):
+        try:
+            with open(LESSONS_FILE, "r") as f:
+                return json.load(f)
+        except:
+            pass
+    return []
+
+def save_lesson(lesson_text, task_id):
+    os.makedirs(os.path.dirname(LESSONS_FILE), exist_ok=True)
+    lessons = load_lessons()
+    lessons.append({"task": task_id, "lesson": lesson_text})
+    with open(LESSONS_FILE, "w") as f:
+        json.dump(lessons, f, indent=2)
+
+def get_system_prompt():
+    prompt = SYSTEM_PROMPT
+    lessons = load_lessons()
+    if lessons:
+        prompt += "\n\nPAST LESSONS LEARNED (DO NOT REPEAT MISTAKES):\n"
+        for l in lessons[-5:]:  # Show only top 5 recent
+            prompt += f"- {l['lesson']}\n"
+    return prompt
+
 
 def _safe_reset(env: RagOptimizerEnvClient, task_id: str):
     """Reset env for a specific task with compatibility fallbacks."""
@@ -111,7 +139,7 @@ def run_task_episode(
             print(f"[END] success=false steps=0 score=0.01 rewards=")
             return
 
-    history = [{"role": "system", "content": SYSTEM_PROMPT}]
+    history = [{"role": "system", "content": get_system_prompt()}]
 
     init_obs = {
         "server_feedback": observation.message,
@@ -194,6 +222,23 @@ def run_task_episode(
     done_str = "true" if success else "false"
     print(f"[END] success={done_str} steps={step} score={score:.2f} rewards={rewards_str}")
 
+    # Memory Reflexion Trigger
+    if not success and score < 0.6:
+        # Agent failed, try to reflect
+        hist_str = json.dumps([m["content"] for m in history[-6:]])  # get last few actions/obs
+        ref_prompt = f"The agent failed task '{task_id}' with final reward {score}. Last context: {hist_str}. Write a 1-sentence tactical lesson stating explicitly what data engineering action the agent should have done instead."
+        try:
+            resp = llm_client.chat.completions.create(
+                model=MODEL_NAME, 
+                messages=[{"role": "user", "content": ref_prompt}], 
+                max_tokens=60
+            )
+            lesson = resp.choices[0].message.content.strip()
+            save_lesson(lesson, task_id)
+            print(f"[MEMORY] Learned lesson: {lesson}")
+        except:
+            pass
+
 def main():
     # Setup OpenAI Client
     client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)

models.py

@@ -58,5 +58,5 @@ class RagOptimizerObservation(BaseModel):
     
     # Required OpenEnv standard fields
     done: bool = Field(False, description="Whether the episode has finished.")
-    reward: float = Field(0.0, description="The reward obtained from the last step.")
+    reward: float = Field(0.01, description="The reward obtained from the last step.")
     metadata: Dict = Field(default_factory=dict, description="Additional optional information.")

openenv.yaml

@@ -13,17 +13,23 @@ tasks:
   - id: easy
     name: "Conflict Resolution"
     description: "Resolve overlapping pricing parameters by deleting the legacy pricing format."
-    reward_range: [0.1, 0.99]
+    reward_range: [0.01, 0.99]
     grader: "rag_optimizer.server.grader:grade_easy"
   - id: medium
-    name: "Ontological Tagging"
-    description: "Route exact programmatic metadata tags to support tickets."
-    reward_range: [0.1, 0.99]
+    name: "Signal Separation"
+    description: >
+      Multi-doc deduplication. KB contains 3 overlapping incident reports
+      for the same outage - one raw, one partial, one with the correct
+      resolution - plus 20 engineering distractors. Agent must delete the
+      noisy duplicates so retrieval surfaces the clean resolution doc at
+      rank 1. Reward improves continuously as noise is removed; no binary
+      metadata gate.
+    reward_range: [0.01, 0.99]
     grader: "rag_optimizer.server.grader:grade_medium"
   - id: hard
     name: "Syntactic Splintering"
     description: "Break down the monolithic onboarding blob into multiple granular chunks."
-    reward_range: [0.1, 0.99]
+    reward_range: [0.01, 0.99]
     grader: "rag_optimizer.server.grader:grade_hard"
 agent:
   inference_entrypoint: "inference.py"

rag_optimizer_environment.py

@@ -119,7 +119,7 @@ class RagOptimizerEnvironment(Environment):
     def _evaluate_kb(self) -> float:
         """The Grader: Evaluates the agent's current KB using TF-IDF."""
         if not self.kb:
-            return 0.0
+            return 0.01
             
         doc_texts = [doc["text"] for doc in self.kb.values()]
         
@@ -127,7 +127,7 @@ class RagOptimizerEnvironment(Environment):
         try:
             doc_vectors = vectorizer.fit_transform(doc_texts)
         except ValueError:
-            return 0.0
+            return 0.01
             
         score = 0.0
         
@@ -147,14 +147,14 @@ class RagOptimizerEnvironment(Environment):
             if found:
                 score += 1.0
                 
-        return float(score / len(self.test_suite))
+        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.0
+        reward = 0.01
         
         try:
             if action.action_type == "read_document":

server/app.py

@@ -53,7 +53,7 @@ app = create_app(
 )
 
 
-def main(host: str = "0.0.0.0", port: int = 8002):
+def main(host: str = "0.0.0.0", port: int = 8000):
     """
     Entry point for direct execution via uv run or python -m.
 

server/kb_seed.json

@@ -0,0 +1,161 @@
+{
+    "easy": {
+        "fastapi_1": {
+            "text": "FastAPI is a modern Python web framework designed for building APIs with high performance. It leverages Python type hints and integrates tightly with validation libraries. One of its key strengths is that it defaults to Pydantic v2, enabling efficient data parsing and validation with minimal boilerplate. Developers can define request and response models clearly, ensuring correctness. FastAPI also supports asynchronous programming using async and await, making it ideal for I/O-bound services. Its automatic OpenAPI schema generation provides interactive API documentation out of the box, improving developer productivity.",
+            "metadata": {
+                "topic": "fastapi",
+                "type": "useful"
+            }
+        },
+        "fastapi_2": {
+            "text": "FastAPI includes a powerful dependency injection system that allows developers to manage shared logic such as authentication, database sessions, and configuration in a clean and reusable way. Dependencies can be declared at different levels, including routes and entire applications. This feature reduces duplication and improves maintainability. The framework also supports async database calls, ensuring efficient resource usage. By combining dependency injection with type hints, FastAPI ensures that applications remain scalable and structured, making it a preferred choice for production-grade microservices.",
+            "metadata": {
+                "topic": "fastapi"
+            }
+        },
+        "fastapi_3": {
+            "text": "FastAPI supports WebSocket communication, enabling real-time features such as chat applications, live dashboards, and streaming data pipelines. Unlike traditional HTTP request-response cycles, WebSockets maintain persistent connections. FastAPI provides simple decorators to define WebSocket endpoints and manage connections efficiently. Combined with async support, this allows handling thousands of concurrent connections. Developers can integrate WebSockets alongside REST endpoints in the same application, making FastAPI highly versatile for both synchronous and real-time applications.",
+            "metadata": {
+                "topic": "fastapi"
+            }
+        },
+        "fastapi_4": {
+            "text": "Running FastAPI applications in production typically involves using ASGI servers such as Uvicorn. A common command used is uvicorn main:app --workers 4, which starts multiple worker processes for handling concurrent requests. This setup improves performance and scalability. FastAPI is designed around ASGI, allowing it to fully utilize asynchronous capabilities. With proper configuration, it can handle high throughput while maintaining low latency, making it suitable for APIs serving large-scale applications.",
+            "metadata": {
+                "topic": "fastapi"
+            }
+        },
+        "fastapi_5": {
+            "text": "FastAPI automatically generates interactive documentation using Swagger UI and ReDoc. This feature is powered by OpenAPI standards and provides developers with a convenient way to test endpoints directly from the browser. The documentation reflects the actual codebase, ensuring accuracy. By leveraging type annotations, FastAPI produces clear schemas for requests and responses. This reduces manual documentation effort and improves collaboration between frontend and backend teams.",
+            "metadata": {
+                "topic": "fastapi"
+            }
+        },
+        "fastapi_6": {
+            "text": "FastAPI emphasizes performance comparable to Node.js and Go by utilizing asynchronous programming and efficient serialization. It uses Starlette under the hood, which provides high-speed routing and middleware support. Developers can create middleware for logging, authentication, and caching. FastAPI also integrates seamlessly with ORMs and external services. Its performance benefits make it suitable for modern cloud-native applications.",
+            "metadata": {
+                "topic": "fastapi"
+            }
+        },
+        "fastapi_7": {
+            "text": "FastAPI simplifies request validation using Python type hints. When defining endpoints, developers specify expected input types, and FastAPI automatically validates incoming requests. Errors are returned in a structured format, helping developers debug quickly. This reduces the need for manual validation logic. Combined with automatic serialization, FastAPI ensures consistency and reliability across APIs.",
+            "metadata": {
+                "topic": "fastapi"
+            }
+        },
+        "fastapi_8": {
+            "text": "FastAPI integrates well with modern Python ecosystems, including async frameworks, ORMs, and testing tools. It supports dependency overrides, which are useful during testing. Developers can simulate database connections or authentication systems easily. This makes writing unit and integration tests straightforward. The framework encourages clean architecture and modular design.",
+            "metadata": {
+                "topic": "fastapi"
+            }
+        },
+        "fastapi_9": {
+            "text": "FastAPI allows developers to define background tasks that run after returning a response. This is useful for sending emails, logging events, or processing data asynchronously. The background task system is lightweight and easy to use. It helps improve response times while handling additional processing efficiently.",
+            "metadata": {
+                "topic": "fastapi"
+            }
+        },
+        "fastapi_10": {
+            "text": "FastAPI supports security features such as OAuth2, API keys, and JWT authentication. It provides built-in utilities for implementing authentication flows. These features integrate seamlessly with dependency injection, allowing secure endpoints to be defined easily. This makes FastAPI suitable for building secure APIs.",
+            "metadata": {
+                "topic": "fastapi"
+            }
+        },
+        "distractor_1": {
+            "text": "Django is a high-level Python web framework that follows the Model-View-Template pattern. It provides an ORM for database interactions and includes built-in features such as authentication, admin panels, and form handling. Django is known for its batteries-included philosophy, allowing developers to build full-stack applications quickly. It is widely used for monolithic applications and supports scalability through modular apps.",
+            "metadata": {
+                "topic": "distractor"
+            }
+        },
+        "distractor_2": {
+            "text": "Flask is a lightweight Python web framework designed for simplicity and flexibility. Unlike Django, it does not include many built-in components, giving developers more control. Flask uses extensions to add functionality such as database integration and authentication. It is ideal for small projects and microservices.",
+            "metadata": {
+                "topic": "distractor"
+            }
+        },
+        "distractor_3": {
+            "text": "Express.js is a popular web framework for Node.js that simplifies building server-side applications. It provides routing, middleware support, and integration with databases. Express is widely used for building REST APIs and full-stack applications using JavaScript. Its minimal design allows developers to structure applications as needed.",
+            "metadata": {
+                "topic": "distractor"
+            }
+        },
+        "distractor_4": {
+            "text": "React is a JavaScript library for building user interfaces. It uses a component-based architecture and virtual DOM for efficient rendering. React is commonly used for frontend development and can be integrated with backend APIs. It supports state management and hooks for managing component lifecycle.",
+            "metadata": {
+                "topic": "distractor"
+            }
+        },
+        "distractor_5": {
+            "text": "Ruby on Rails is a web application framework that emphasizes convention over configuration. It provides a full-stack solution with built-in tools for database management, routing, and views. Rails is known for rapid development and clean code structure.",
+            "metadata": {
+                "topic": "distractor"
+            }
+        },
+        "distractor_6": {
+            "text": "Angular is a frontend framework developed by Google. It uses TypeScript and provides tools for building large-scale applications. Angular includes dependency injection, routing, and state management.",
+            "metadata": {
+                "topic": "distractor"
+            }
+        },
+        "distractor_7": {
+            "text": "Laravel is a PHP framework that simplifies web development with features like routing, authentication, and database migrations. It is widely used for building scalable web applications.",
+            "metadata": {
+                "topic": "distractor"
+            }
+        },
+        "distractor_8": {
+            "text": "Spring Boot is a Java-based framework for building enterprise applications. It simplifies configuration and provides tools for building REST APIs and microservices.",
+            "metadata": {
+                "topic": "distractor"
+            }
+        },
+        "distractor_9": {
+            "text": "Vue.js is a progressive JavaScript framework for building user interfaces. It focuses on simplicity and flexibility, making it easy to integrate into projects.",
+            "metadata": {
+                "topic": "distractor"
+            }
+        },
+        "distractor_10": {
+            "text": "ASP.NET Core is a cross-platform framework for building web applications using C#. It supports MVC architecture and integrates with Microsoft technologies.",
+            "metadata": {
+                "topic": "distractor"
+            }
+        }
+    },
+    "medium": {
+        "falcon9": {
+            "text": "Falcon 9 is a reusable rocket developed by SpaceX for transporting payloads to orbit. It uses SpaceX Merlin engines, which are designed for efficiency and reusability. The first stage can land vertically after launch, significantly reducing costs. Falcon 9 has been used for satellite deployments, cargo missions, and crewed flights to the International Space Station. Its reliability and cost-effectiveness have made it a cornerstone of modern space exploration.",
+            "metadata": {}
+        },
+        "dragon": {
+            "text": "The Dragon spacecraft developed by SpaceX is used for cargo and crew transport. The original Dragon capsule was replaced by Dragon 2, which includes advanced life-support systems and autonomous docking capabilities. Dragon has been used extensively for resupply missions to the International Space Station and has enabled private spaceflight.",
+            "metadata": {}
+        },
+        "starship": {
+            "text": "Starship is a fully reusable spacecraft designed by SpaceX for deep space missions. It aims to transport humans to Mars and beyond. The system includes a booster and spacecraft, both reusable. Starship is powered by Raptor engines and is designed for rapid turnaround.",
+            "metadata": {}
+        },
+        "distractor_1": {
+            "text": "Deep sea exploration involves studying ocean depths using submersibles and remote-operated vehicles. Scientists investigate marine ecosystems and geological formations.",
+            "metadata": {}
+        }
+    },
+    "hard": {
+        "crisis_monolith": {
+            "text": "The 2008 financial crisis was one of the most severe economic downturns since the Great Depression. It was triggered by the collapse of the housing market and widespread use of Mortgage-backed securities. Financial institutions had heavily invested in subprime mortgages, which were loans given to low-income homebuyers with poor credit histories. As housing prices declined, these borrowers began defaulting, leading to massive losses. The bankruptcy of Lehman Brothers marked a critical moment, causing panic across global markets. Complex financial instruments such as CDOs amplified risk, spreading it throughout the system. Governments intervened with bailouts, including the TARP program, to stabilize banks. Central banks implemented quantitative easing to inject liquidity into the economy. The crisis exposed weaknesses in financial regulation and led to reforms such as the Dodd-Frank Act. It also highlighted the dangers of excessive leverage and lack of transparency in financial markets.",
+            "metadata": {}
+        },
+        "tarp": {
+            "text": "The Troubled Asset Relief Program was introduced to stabilize financial institutions during the crisis. It allowed the government to purchase toxic assets and inject capital into banks.",
+            "metadata": {}
+        },
+        "qe": {
+            "text": "Quantitative easing was used by central banks to increase liquidity by purchasing government securities and lowering interest rates.",
+            "metadata": {}
+        },
+        "noise_1": {
+            "text": "During the 1920s, financial markets experienced rapid growth followed by a crash. Speculative investments and excessive leverage contributed to economic instability.",
+            "metadata": {}
+        }
+    }
+}

server/rag_optimizer_environment.py

@@ -16,11 +16,15 @@ 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
+# Import scikit-learn for our Grader (fallback/legacy)
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 
+# Hybrid Search imports
+from sentence_transformers import SentenceTransformer
+from rank_bm25 import BM25Okapi
+
 try:
     from models import RagOptimizerAction, RagOptimizerObservation
 except ImportError:
@@ -37,62 +41,81 @@ class RagOptimizerEnvironment(Environment):
 
     def __init__(self):
         self._state = State(episode_id=str(uuid4()), step_count=0)
+        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2') 
         self.kb = {}
         self.test_suite = []
+        self.dense_vectors = {}
+        
+        # Load the Real World Datasets
+        import json
+        import os
+        seed_path = os.path.join(os.path.dirname(__file__), "kb_seed.json")
+        with open(seed_path, "r") as f:
+            self.seed_data = json.load(f)
+            
         self._setup_task("easy")
         
     def _setup_task(self, task_id: str):
+        if task_id not in ["easy", "medium", "hard"]:
+            task_id = "easy"
+            
+        # Clone the fresh dataset from seed so the agent can destroy it
+        self.kb = deepcopy(self.seed_data[task_id])
+        
         if task_id == "easy":
-            self.kb = {
-                "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": {}
-                },
-                **{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)}
-            }
             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": "How do I run the server in production with concurrency?", "target_concept": "uvicorn main:app --workers 4"},
+                {"query": "Which version of Pydantic does FastAPI use by default?", "target_concept": "defaults to Pydantic v2"}
             ]
+            
         elif task_id == "medium":
             self.kb = {
-                "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.",
+                "doc_incident_raw": {
+                    "text": "User reported frontend button missing. Database latency also flagged. No fix yet.",
+                    "metadata": {}
+                },
+                "doc_incident_partial": {
+                    "text": "Button issue and DB latency investigated. CSS fix attempted. Email 401 error also reported.",
                     "metadata": {}
                 },
-                "doc_messy_support_ticket_2": {
-                    "text": "Email integration is failing with error 401 Unauthorized. The API key was rotated on Tuesday.",
+                "doc_incident_resolved": {
+                    "text": "Frontend button restored by updating CSS stylesheet. Email 401 resolved after API key rotation on Tuesday.",
                     "metadata": {}
                 },
-                **{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)}
             }
+            for i in range(20):
+                self.kb[f"doc_distractor_eng_{i}"] = {
+                    "text": f"Architecture decision {i}: chose Postgres for horizontal scaling.",
+                    "metadata": {}
+                }
             self.test_suite = [
-                {"query": "UI issues frontend CSS missing button", "target_concept": "frontend team fixed the button"},
-                {"query": "Email integration 401", "target_concept": "API key was rotated"}
+                {"query": "How was the missing UI element fixed?",
+                 "target_concept": "updating CSS stylesheet"},
+                {"query": "What caused the authentication failure?",
+                 "target_concept": "API key rotation on Tuesday"},
             ]
+            
         elif task_id == "hard":
-            self.kb = {
-                "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_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": {}
-                },
-                **{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)}
-            }
             self.test_suite = [
-                {"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"}
+                {"query": "Which entity bankruptcy was the climax of the disaster?", "target_concept": "bankruptcy of Lehman Brothers"},
+                {"query": "What types of collateralized assets lost their worth?", "target_concept": "Mortgage-backed securities"},
+                {"query": "Who did predatory lenders primarily go after?", "target_concept": "low-income homebuyers"}
             ]
-        else:
-            self._setup_task("easy")
+
+        self._rebuild_cache()
+
+    def _rebuild_cache(self):
+        """Called whenever KB documents are added, removed, or updated."""
+        if not self.kb:
+            self.dense_vectors = {}
+            return
+            
+        doc_ids = list(self.kb.keys())
+        # Append metadata to text for embedding
+        doc_texts = [(self.kb[d]["text"] + " " + " ".join(self.kb[d]["metadata"].values())).strip() for d in doc_ids]
+        
+        vectors = self.embedding_model.encode(doc_texts, convert_to_tensor=False)
+        self.dense_vectors = {doc_id: vectors[i] for i, doc_id in enumerate(doc_ids)}
 
     def _get_kb_summary(self) -> Dict[str, Dict]:
         """Returns a summary of the KB for the observation."""
@@ -114,44 +137,78 @@ class RagOptimizerEnvironment(Environment):
         )
 
     def _evaluate_kb(self) -> float:
-        """The Grader: Evaluates the agent's current KB using TF-IDF."""
-        if not self.kb:
-            return 0.0
+        """The Grader: Evaluates the current KB using Hybrid RRF (BM25 + Semantic MRR)."""
+        if not self.kb or not self.test_suite:
+            return 0.01
             
-        doc_texts = [doc["text"] for doc in self.kb.values()]
+        doc_ids = list(self.kb.keys())
+        doc_texts = [(self.kb[d]["text"] + " " + " ".join(self.kb[d]["metadata"].values())).strip() for d in doc_ids]
         
-        vectorizer = TfidfVectorizer(stop_words='english')
-        try:
-            doc_vectors = vectorizer.fit_transform(doc_texts)
-        except ValueError:
-            return 0.0
-            
-        score = 0.0
+        # 1. BM25 Corpus Preparation
+        tokenized_corpus = [doc.lower().split() for doc in doc_texts]
+        bm25 = BM25Okapi(tokenized_corpus)
+        
+        # 2. Dense Matrix
+        doc_vectors = np.array([self.dense_vectors[d] for d in doc_ids])
+        
+        mrr_sum = 0.0
         
         for case in self.test_suite:
-            query_vec = vectorizer.transform([case["query"]])
-            similarities = cosine_similarity(query_vec, doc_vectors)[0]
+            # BM25 Search
+            tokenized_query = case["query"].lower().split()
+            bm25_scores = bm25.get_scores(tokenized_query)
+            bm25_ranks = bm25_scores.argsort()[::-1]
             
-            # Get top 3
-            top_k_indices = similarities.argsort()[-3:][::-1]
+            # Dense Search
+            query_vec = self.embedding_model.encode(case["query"])
+            dense_scores = cosine_similarity([query_vec], doc_vectors)[0]
+            dense_ranks = dense_scores.argsort()[::-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
+            # Reciprocal Rank Fusion (RRF)
+            rrf_scores = {d: 0.0 for d in doc_ids}
+            k = 60
+            for rank, idx in enumerate(bm25_ranks):
+                rrf_scores[doc_ids[idx]] += 1.0 / (k + rank + 1)
+            for rank, idx in enumerate(dense_ranks):
+                rrf_scores[doc_ids[idx]] += 1.0 / (k + rank + 1)
                 
-        return float(score / len(self.test_suite))
+            # Grade Top-K fused list using MRR
+            ranked_doc_ids = sorted(rrf_scores.keys(), key=lambda x: rrf_scores[x], reverse=True)
+            
+            req_meta = case.get("required_metadata_key")
+            req_meta_val = case.get("required_metadata_value")
+            
+            case_mrr = 0.0
+            for i, doc_id in enumerate(ranked_doc_ids):
+                # Is this the true doc?
+                doc_text = self.kb[doc_id]["text"].lower()
+                if case["target_concept"].lower() in doc_text:
+                    valid = True
+                    
+                    # Medium Task Semantic Check
+                    if req_meta and req_meta_val:
+                        if self.kb[doc_id]["metadata"].get(req_meta) != req_meta_val:
+                            valid = False
+                            
+                    if valid:
+                        case_mrr = 1.0 / (i + 1)  # MRR formula starts at rank 1
+                        break
+            
+            mrr_sum += case_mrr
+            
+        base_reward = float(mrr_sum / len(self.test_suite))
+        
+        # Step Cost Penalty calculation (-0.01 per step)
+        cost_penalty = self._state.step_count * 0.01
+        
+        return max(0.01, min(0.99, base_reward - cost_penalty))
 
     def step(self, action: RagOptimizerAction) -> RagOptimizerObservation:  # type: ignore[override]
         self._state.step_count += 1
         
         msg = ""
         done = False
-        reward = 0.0
+        reward = 0.01
         
         try:
             if action.action_type == "read_document":
@@ -163,6 +220,7 @@ class RagOptimizerEnvironment(Environment):
             elif action.action_type == "delete_document":
                 if action.doc_id in self.kb:
                     del self.kb[action.doc_id]
+                    self._rebuild_cache()
                     msg = f"Deleted {action.doc_id}."
                 else:
                     msg = f"Error: doc_id {action.doc_id} not found."
@@ -174,6 +232,7 @@ class RagOptimizerEnvironment(Environment):
                     if action.doc_id not in self.kb:
                         self.kb[action.doc_id] = {"text": "", "metadata": {}}
                     self.kb[action.doc_id]["text"] = action.text
+                    self._rebuild_cache()
                     msg = f"Updated text for {action.doc_id}."
                     
             elif action.action_type == "add_metadata":
@@ -184,6 +243,7 @@ class RagOptimizerEnvironment(Environment):
                         msg = f"Error: doc_id {action.doc_id} not found."
                     else:
                         self.kb[action.doc_id]["metadata"][action.metadata_key] = action.metadata_value
+                        self._rebuild_cache()
                         msg = f"Added metadata to {action.doc_id}."
                         
             elif action.action_type == "submit":

server/requirements.txt

@@ -2,3 +2,5 @@ openenv[core]>=0.2.0
 fastapi>=0.115.0
 uvicorn>=0.24.0
 scikit-learn>=1.3.0
+sentence-transformers>=3.0.0
+rank_bm25>=0.2.2