Adding retrieval augmented generation (#12)

* Add RAG module under benchmarking

* Imported RAG - first edit to MultiAgentTester

* changes to files rag

* RAG model working draft

* improved format

* HTTPclient change - unfixed

* more prelim changes to the rag model

* initial working prototype revised

* added in functions.json and embeddings.json

* extracted from scib lib

* a better working draft of the rag

* more changes to rag

* small change

* finally no errors!

* more changes made

* added new extractor class

* trying to get scanpy to work

* scanpy works

* working proto for scanpy and scib-metrics

* just added type annotations

* trying to resolve json error

* decent working ver before signing off

* Deleted embeddings.json and functions.json

* remove json

* making sure file is safe

* fixed issue by changing to jsonl

* experimenting around with wikipedia lib

* fixeS?

* added umap with new prompts

* fixed umap

* reverted back to the technique with urls

* fixed error

* added some visualization in umap and heatmap

* added fixes to rag file

* file finalized for 2day

* added variations to chunking

* created a series of images to test

* varying wiki +description contents

* diff sizes of wiki page

* more variations

* trying with a bigger embedding model

* trying wiki api

* working implementation -not helpful

* trying to aggresively remove stuff from the wiki result

* trying to aggresively remove more

* fixing bugs

* trying to make a new ver work

* small syntax error

* draft

* switched out wiki lib for beautiful soup extraction + switched model

* working version yayy

* rag file

* working version of the rag class system

* new file for user purposes!

* changes to user and skeleton ver

* added a new function to extract wiki content

* attempt to fix

* fixes and clean up

* fixes and clean ups ongoing

* file improvisations with request failsafe

* resolved critical errors

* testing

* new fixes

* removed embedding and functions file

* fixing shennanigans

* moved files

* changes to file in type annotations

* made more fixes to rag class by adding 1 function for the entire pipeline

* took back some fixes that introduced errors

* changed text to string type in func_Def

* regex for improvements to extracting html

* introduced more aggressive regex for cleaning func def

* added support for dict objects

* type annotations changed + dict incorporated

* errors fixed

* quick fixes to rag file

* more changes for error correction

* deleted folder from wrong location

* moved locations for rag folder

* improper folder placement

* moved locations

* making changes to the runner.py file system

* moved rag

* rag + changes to runner

* one working version of implementation of rag - using agents

* dylan's proposed version of the implementation

* attempts at rag implementation

* Added rag support to agent system

* query function from database by function signature search

* changed function definition search to function signature search

* fixed rag implementation

* may have fixed import, need to consult and fix file

* potentially error fix

* fixed imports

* working with dylans location of rag folder

* userrag file deemed unnecessary and deleted

* change to file names and locations - more clean up, fixed imports

* user rag file changed

* Update system_blueprint.json

* changes to imports and file structures

* finally fixed import situation

* tested code

* trivial errors

* trivial errors

* Update RetrievalAugmentedGeneration.py

* Fixed rag implementation

* small UX fixes

* added in new embeddings

* embeddings and functions file created - however the search results from wikipedia are not accurate

* fixed embedding file structure

* need to fix wikipedia

* restructured embeddings.jsonl to signature and embedding + added in query relevance

* Synced New Embeddings and Functions

---------

Co-authored-by: djriffle <djriffle1@gmail.com>

Author: varshinivij

cli/extra_tools/RetrievalAugmentedEmbedder.py

@@ -22,7 +22,9 @@ dependencies = [
   "jupyter-client",   # NOTE: PyPI name has a hyphen
   "nbformat",
   "typer",
-  "platformdirs"
+  "platformdirs",
+  "sentence_transformers",
+  "tf_keras"
 ]
 
 # If you want a command like `olaf …`

cli/extra_tools/embeddings.jsonl

@@ -12,7 +12,6 @@
 # 2. The package-internal directory (for default samples), found relative to this file
 PACKAGE_CODE_SAMPLES_DIR = Path(__file__).resolve().parent.parent / "code_samples"
 
-
 class Command:
     """Represents a command an agent can issue to a neighboring agent."""
     def __init__(self, name: str, target_agent: str, description: str):
@@ -23,19 +22,18 @@ def __init__(self, name: str, target_agent: str, description: str):
     def __repr__(self) -> str:
         return (f"Command(name='{self.name}', target='{self.target_agent}', "
                 f"desc='{self.description[:30]}...')")
-
-
 class Agent:
     """Represents a single agent in the system."""
-    def __init__(self, name: str, prompt: str, commands: Dict[str, Command], code_samples: Dict[str, str]):
+    def __init__(self, name: str, prompt: str, commands: Dict[str, Command], code_samples: Dict[str, str], is_rag_enabled: bool = False):
         self.name = name
         self.prompt = prompt
         self.commands = commands
         self.code_samples = code_samples
+        self.is_rag_enabled = is_rag_enabled
 
     def __repr__(self) -> str:
         sample_keys = list(self.code_samples.keys())
-        return f"Agent(name='{self.name}', commands={list(self.commands.keys())}, samples={sample_keys})"
+        return f"Agent(name='{self.name}', commands={list(self.commands.keys())}, samples={sample_keys}, rag_enabled={self.is_rag_enabled})"
 
     def get_full_prompt(self, global_policy=None) -> str:
         """Constructs the full prompt including the global policy and command descriptions."""
@@ -53,8 +51,14 @@ def get_full_prompt(self, global_policy=None) -> str:
                 full_prompt += f"\n  - Target Agent: {command.target_agent}"
             full_prompt += "\n\n**YOU MUST USE THESE EXACT COMMANDS TO DELEGATE TASKS. NO OTHER FORMATTING OR COMMANDS ARE ALLOWED.**"
 
+        if self.is_rag_enabled:
+            full_prompt += "\n\nYou can query your specialized knowledge base for more context with the following command:"
+            full_prompt += f"\n- Command: `query_rag_<function>`"
+            full_prompt += f"\n  - Description: Retrieves relevant information about a specific <function> from your knowledge base. Replace <function> with a concise, descriptive search query (e.g., function names, task you are trying to complete)."
+            full_prompt += f"\n  - Example: `query_rag_<scvi model setup>`"
+            
         if self.code_samples:
-            full_prompt += "\n  - Code Samples Available:"
+            full_prompt += "\n\n  - Code Samples Available:"
             for sample_name in self.code_samples.keys():
                 full_prompt += f"\n    - `{sample_name}`"
 
@@ -101,7 +105,6 @@ def load_from_json(cls, file_path: str) -> 'AgentSystem':
                     user_path = USER_CODE_SAMPLES_DIR / filename
                     package_path = PACKAGE_CODE_SAMPLES_DIR / filename
 
-                    # Default to package path, but overwrite if user path exists
                     path_to_load = None
                     source_label = ""
                     if user_path.exists():
@@ -120,11 +123,15 @@ def load_from_json(cls, file_path: str) -> 'AgentSystem':
                     else:
                         print(f"    ❌ WARNING: Code sample file '{filename}' not found in any location.")
 
+            rag_config = agent_data.get("rag", {})
+            is_rag_enabled = rag_config.get("enabled", False)
+
             agent = Agent(
                 name=agent_name,
                 prompt=agent_data['prompt'],
                 commands=commands,
-                code_samples=loaded_samples
+                code_samples=loaded_samples,
+                is_rag_enabled=is_rag_enabled
             )
             agents[agent_name] = agent
 

cli/extra_tools/functions.jsonl

@@ -2,11 +2,11 @@
 import os
 from typing import Dict, Any
 from pathlib import Path
-from platformdirs import PlatformDirs   # pip install platformdirs
+from platformdirs import PlatformDirs
 import tempfile
 
 APP_NAME   = "olaf"
-APP_AUTHOR = "OpenTechBio"   # or your org
+APP_AUTHOR = "OpenTechBio"
 dirs = PlatformDirs(APP_NAME, APP_AUTHOR)
 
 # Root for user-specific OLAF files. Precedence: env -> platformdirs.
@@ -72,9 +72,24 @@ def define_agents() -> Dict[str, Dict[str, Any]]:
         if agent_name in agents:
             print(f"{Colors.FAIL}Agent '{agent_name}' already exists. Please use a unique name.{Colors.ENDC}")
             continue
+        
         prompt = input(f"{Colors.WARNING}Enter the system prompt for '{Colors.OKCYAN}{agent_name}{Colors.WARNING}': {Colors.ENDC}").strip()
-        agents[agent_name] = {"prompt": prompt, "neighbors": {}, "code_samples": []}
-        print(f"{Colors.OKGREEN}Agent '{Colors.OKCYAN}{agent_name}{Colors.OKGREEN}' added successfully.{Colors.ENDC}")
+        
+        # --- New RAG Configuration Section ---
+        rag_enabled_input = input(f"{Colors.WARNING}Enable Retrieval-Augmented Generation (RAG) for '{Colors.OKCYAN}{agent_name}{Colors.WARNING}'? (y/n): {Colors.ENDC}").strip().lower()
+        is_rag_enabled = rag_enabled_input == 'y'
+        
+        # Add the new 'rag' key to the agent's data structure
+        agents[agent_name] = {
+            "prompt": prompt,
+            "neighbors": {},
+            "code_samples": [],
+            "rag": {"enabled": is_rag_enabled}
+        }
+        
+        rag_status = f"{Colors.OKGREEN}enabled" if is_rag_enabled else f"{Colors.FAIL}disabled"
+        print(f"{Colors.OKGREEN}Agent '{Colors.OKCYAN}{agent_name}{Colors.OKGREEN}' added successfully (RAG: {rag_status}{Colors.OKGREEN}).{Colors.ENDC}")
+
     print(f"\n{Colors.OKBLUE}--- All Agents Defined ---{Colors.ENDC}")
     for name in agents:
         print(f"- {Colors.OKCYAN}{name}{Colors.ENDC}")
@@ -162,7 +177,7 @@ def _atomic_write_json(obj: Any, path: Path) -> None:
     with tempfile.NamedTemporaryFile("w", delete=False, dir=str(path.parent), prefix=path.stem, suffix=".tmp") as tmp:
         json.dump(obj, tmp, indent=2)
         tmp_path = Path(tmp.name)
-    tmp_path.replace(path)  # atomic on POSIX; safe on Windows
+    tmp_path.replace(path)
 
 def save_configuration(global_policy: str, agents_config: Dict[str, Any], output_dir: str) -> None:
     if not agents_config:

cli/olaf/pyproject.toml

@@ -2,6 +2,9 @@
   "global_policy": "Always be concise, professional, and helpful. Do not refuse to answer a request unless it is harmful.",
   "agents": {
     "master_agent": {
+      "rag": {
+        "enabled": true
+      },
       "prompt": "You are the master agent. Analyze every user request and delegate the task to the appropriate expert: the general coder for standard single-cell analysis or the integration expert for batch correction and data integration tasks. Respond ONLY with a delegation command.",
       "neighbors": {
         "delegate_to_general": {
@@ -16,6 +19,9 @@
     },
     "general_coder": {
       "prompt": "You are the *general scRNA-seq coder*. You handle standard single-cell analysis tasks like data loading, QC, filtering, normalization, and basic plotting using scanpy. You are not an expert in data integration.\n\nExample of a task you would perform:\n```python\nimport scanpy as sc\n\n# Assume 'adata' is a loaded AnnData object\n# Basic QC and filtering\nsc.pp.filter_cells(adata, min_genes=200)\nsc.pp.filter_genes(adata, min_cells=3)\nadata.var['mt'] = adata.var_names.str.startswith('MT-')\nsc.pp.calculate_qc_metrics(adata, qc_vars=['mt'], inplace=True)\n\n# Normalize and find highly variable genes\nsc.pp.normalize_total(adata, target_sum=1e4)\nsc.pp.log1p(adata)\nsc.pp.highly_variable_genes(adata, min_mean=0.0125, max_mean=3, min_disp=0.5)\n\n# Run PCA\nsc.tl.pca(adata, svd_solver='arpack')\n\nprint('Standard analysis complete. PCA is in adata.obsm[\"X_pca\"].')\n```",
+      "rag": {
+        "enabled": true
+      },
       "neighbors": {
         "delegate_to_master": {
           "target_agent": "master_agent",
@@ -29,6 +35,9 @@
     },
     "integration_expert": {
       "prompt": "You are the *integration expert*. You specialize in combining multiple single-cell datasets and correcting for batch effects using scvi-tools.\n\nExample of a task you would perform:\n```python\nimport scvi\nimport scanpy as sc\n\n# Assume 'adata' is loaded and preprocessed with a 'batch' column\n# Find highly variable genes across batches for integration\nsc.pp.highly_variable_genes(\n    adata,\n    n_top_genes=2000,\n    subset=True,\n    layer='counts',\n    flavor='seurat_v3',\n    batch_key='batch'\n)\n\n# Set up the AnnData object for the scVI model\nscvi.model.SCVI.setup_anndata(adata, layer='counts', batch_key='batch')\n\n# Create and train the scVI model\nmodel = scvi.model.SCVI(adata, n_layers=2, n_latent=30)\nmodel.train()\n\n# Store the integrated latent representation in the AnnData object\nadata.obsm['X_scVI'] = model.get_latent_representation()\n\nprint('Integration complete. Integrated embedding is in adata.obsm[\"X_scVI\"].')\n``` you remeber to wrap your code in triple backticks and python. Please only include one code block per response. Remeber to keep responses short and to the point.",
+      "rag": {
+        "enabled": true
+      },
       "neighbors": {
         "delegate_to_master": {
           "target_agent": "master_agent",

cli/olaf/src/olaf/agents/AgentSystem.py

@@ -3,6 +3,9 @@
   "agents": {
     "master_agent": {
       "prompt": "You are the master agent. Your primary role is to analyze incoming user requests and delegate them to the appropriate specialist agent. You do not perform tasks yourself.",
+      "rag": {
+        "enabled": true
+      },
       "neighbors": {
         "delegate_to_coder": {
           "target_agent": "coder_agent",
@@ -16,14 +19,20 @@
     },
     "coder_agent": {
       "prompt": "You are a specialist single cell RNA coder agent. Your job is to write high-quality, executable code based on the user's request. You do not delegate tasks. The machine you run on has write disabled. You should never save to disk or modify files. Prioritize small step responses and avoid large code dumps.",
+      "rag": {
+        "enabled": true
+      },
       "neighbors": {},
       "code_samples": [
         "load_adata.py"
       ]
     },
     "research_agent": {
       "prompt": "You are a specialist research agent. You fulfill user requests by finding and synthesizing information from reliable sources. You do not write code or delegate tasks.",
+      "rag": {
+        "enabled": true
+      },
       "neighbors": {}
     }
   }
-}
+}

cli/olaf/src/olaf/agents/create_agent_system.py

@@ -17,6 +17,7 @@
     from olaf.config import OLAF_HOME
     from olaf.agents.AgentSystem import Agent, AgentSystem
     from olaf.core.io_helpers import display, extract_python_code, format_execute_response
+    from olaf.rag.RetrievalAugmentedGeneration import RetrievalAugmentedGeneration
 except ImportError as e:
     print(f"Failed to import a required OLAF module: {e}", file=sys.stderr)
     sys.exit(1)
@@ -39,6 +40,9 @@ def exec_code(self, code: str, timeout: int) -> dict:
 _OUTPUTS_DIR = OLAF_HOME / "runs"
 _SNIPPET_DIR = _OUTPUTS_DIR / "snippets"
 _LEDGER_PATH = _OUTPUTS_DIR / f"benchmark_history_{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}.jsonl"
+_RAG_RE = re.compile(r"query_rag_<([^>]+)>")
+RAG = RetrievalAugmentedGeneration()
+
 
 def _init_paths():
     """Ensure output directories exist before writing."""
@@ -51,6 +55,11 @@ def detect_delegation(msg: str) -> Optional[str]:
     m = _DELEG_RE.search(msg)
     return f"delegate_to_{m.group(1)}" if m else None
 
+def detect_rag(msg: str) -> Optional[str]:
+    """Return the *partial* RAG command if present."""
+    m = _RAG_RE.search(msg)
+    return f"{m.group(1)}" if m else None
+
 def _dump_code_snippet(run_id: str, code: str) -> str:
     """Write <run_id>.py under outputs/snippets/ and return the relative path."""
     snippet_path = _SNIPPET_DIR / f"{run_id}.py"
@@ -69,7 +78,8 @@ def _save_benchmark_record(*, run_id: str, results: dict, meta: dict, code: str
         record["code_path"] = _dump_code_snippet(run_id, code)
     with _LEDGER_PATH.open("a") as fh:
         fh.write(json.dumps(record) + "\n")
-
+    
+        
 # --- Core Runner Functions ---
 def run_benchmark(
     console: Console,
@@ -187,7 +197,21 @@ def run_agent_session(
             break
 
         history.append({"role": "assistant", "content": msg})
-        display(console, f"assistant ({current_agent.name})", msg)
+        display(console, f"assistant ({current_agent.name})", msg)  
+
+        # --- RAG handling ---
+        query_from_re = detect_rag(msg)
+        if query_from_re and current_agent.is_rag_enabled:
+            console.print(f"[yellow]🔍 Triggering RAG query: {query_from_re}[/yellow]")
+            retrieved_docs = RAG.query(query_from_re)
+            if retrieved_docs:
+                console.print(f"[green] RAG query successful. [/green]")
+                feedback = retrieved_docs
+                console.print(feedback)
+                history.append({"role": "system", "content": feedback}) 
+            else:
+                console.print(f"[red] RAG query unsuccessful. [/red]")
+            
 
         cmd = detect_delegation(msg)
         if cmd and cmd in current_agent.commands:
@@ -211,8 +235,39 @@ def run_agent_session(
             console.print("[cyan]Executing code in sandbox…[/cyan]")
             exec_result = sandbox_manager.exec_code(code, timeout=300)
             feedback = format_execute_response(exec_result, _OUTPUTS_DIR)
-            history.append({"role": "user", "content": feedback})
-            display(console, "user", feedback)
+            history.append({"role": "assistant", "content": feedback})
+            display(console, "assistant", feedback)
+
+            stderr = exec_result.get('stderr', '')
+            if stderr and current_agent.is_rag_enabled:
+                func_error_patterns = [
+                    r"missing \d+ required positional argument",  # TypeError: missing argument
+                    r"NameError: name '(\w+)' is not defined",    # NameError
+                    r"AttributeError: '.*' object has no attribute '(\w+)'",  # missing attribute
+                    r"got an unexpected keyword argument"         # wrong keyword argument
+                ]
+                function_name = ""
+                retrieved_docs = ""
+                if any(re.search(pat, stderr) for pat in func_error_patterns):
+                    lines = stderr.strip().splitlines()
+                    if len(lines) >= 2:
+                        code_line = lines[-2].strip()  # second-to-last line: code that failed
+                        match = re.search(r'(\w+)\s*\(', code_line)
+                        if match:
+                            function_name = match.group(1)
+            
+                if function_name:
+                    retrieved_docs = RAG.retrieve_function(function_name)
+                    console.print(f"[yellow]🔍 Missing function detected: {function_name}, function database search...[/yellow]")
+                if retrieved_docs:
+                    console.print(f"[green] Query successful - Function signature found. [/green]")
+                    feedback += f"\n {function_name} produced an error. The correct function signature for {function_name} is:\n{retrieved_docs}"
+                    console.print(feedback)
+                    history.append({"role": "system", "content": feedback})
+                    continue
+                else:
+                    print(f"RAG Error Query unsuccessful - Function signature does not exist in the current database.")
+             
 
         if is_auto:
             if benchmark_modules:

cli/olaf/src/olaf/agents/integration_system.json

@@ -0,0 +1,84 @@
+import json
+import sys
+from pathlib import Path
+from typing import List, Dict, Optional
+from contextlib import redirect_stdout, redirect_stderr
+
+import os
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
+# ── Dependencies ─────────────────────────────────────────────
+try: 
+    import re 
+    from sentence_transformers import SentenceTransformer
+    from rich.console import Console
+    import matplotlib.pyplot as plt
+    import numpy as np
+    
+except ImportError as e:
+    print(f"Missing dependency: {e}", file=sys.stderr)
+    sys.exit(1) 
+
+# ── Paths and Constants ─────────────────────────────────────────────
+console = Console()
+
+RAG_DIR = Path(__file__).resolve().parent.parent / "rag"
+EMBEDDING_FILE = RAG_DIR / "embeddings.jsonl"
+FUNCTIONS_FILE = RAG_DIR / "functions.jsonl"
+
+class RetrievalAugmentedGeneration():
+    model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B")
+
+    def __init__(self):
+        self.embeddings = self.load_embeddings()
+        self.functions = self.load_functions()
+        self.queries = []
+
+    def load_embeddings(self) -> List[np.ndarray]:
+        try:
+            with open(EMBEDDING_FILE, "r", encoding="utf-8") as f:
+                return [np.array(json.loads(line)) for line in f if line.strip()]
+        except FileNotFoundError:
+            console.log("[red]Embeddings file not found.")
+            return []
+        except json.JSONDecodeError:
+            console.log("[red]Embeddings file is not valid JSONL.")
+            return []
+    
+    def load_functions(self) -> List[Dict[str, str]]:
+        try:
+            with open(FUNCTIONS_FILE, "r", encoding="utf-8") as f:
+                return [json.loads(line) for line in f if line.strip()]
+        except FileNotFoundError:
+            console.log("[red]Functions file not found.")
+            return []
+        except json.JSONDecodeError:
+            console.log("[red]Functions file is not valid JSONL.")
+            return []
+
+    @staticmethod
+    def cosine_similarity(A: np.ndarray, B: List[np.ndarray]) -> List[float]:
+        sims = [np.dot(A, emb) / (np.linalg.norm(A) * np.linalg.norm(emb)) for emb in B]
+        return sims
+
+    def retrieve_function(self, name:str) -> Optional[str]:
+        for function in self.functions:
+            if name in function["signature"]:
+                return function["signature"]
+        return None
+
+    def query(self, text_query: str) -> Optional[np.ndarray]:
+        self.queries.append(text_query)
+        if not self.embeddings:
+            console.log("[yellow]No embeddings to compare.")
+            return None
+        query_embedding = self.model.encode([text_query])[0]
+        sims = self.cosine_similarity(query_embedding, self.embeddings)
+        idx = np.argmax(sims)
+        return self.functions[idx]["signature"]
+
+ # ──────Implementation──────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    rag = RetrievalAugmentedGeneration()
+    print(rag.query("What is pca"))