feat: ✨ AI Read File Tool, Configurable system prompts and loading lots of llms

src/experts/retrieval_agent.py

@@ -0,0 +1,97 @@
+import os
+import turso
+import dspy
+
+
+from config_loader import load_config
+from embedding import LocalLMEmbeddings
+
+CFG = load_config()
+
+DATABASE_PATH = CFG["ingestion"]["db_path"]
+EMBEDDING_MODEL = CFG["models"]["embedding"]
+API_BASE = CFG["api"]["base_url"]
+RETRIEVAL_CONFIG = CFG["retrieval_agent"]
+
+
+def retrieve_from_turso(embedded_question, k=5):
+    query = f"""
+    SELECT file_path, synopsis, tags, entities, chunk_data,
+    vector_distance_cos(embedding, vector32('{embedded_question[0]}')) AS distance
+    FROM notes
+    ORDER BY distance ASC
+    LIMIT {k};
+    """
+    con = turso.connect(DATABASE_PATH)
+    cur = con.cursor()
+    cur.execute(query)
+    rows = cur.fetchall()
+    return rows
+
+# --- DSPy Signature ---
+class DnDContextQA(dspy.Signature):
+    f"{RETRIEVAL_CONFIG["retrieval_signature"]}"
+
+    context = dspy.InputField(
+        desc="Relevant chunks and metadata from the campaign notes."
+    )
+    question = dspy.InputField()
+    answer = dspy.OutputField(desc="A detailed answer based on the notes, citing the source file.")
+
+
+class DnDRAG(dspy.Module):
+    def __init__(self):
+        super().__init__()
+        self.embeddings_model = LocalLMEmbeddings(
+            model=EMBEDDING_MODEL,
+            base_url=API_BASE,
+            batch_size=1, # we only send 1 question at a time.
+            )
+                # Tools exposed to the ReAct loop
+        self.tools = [
+            self.load_file
+        ]
+        self.generate_answer = dspy.ReAct(signature=DnDContextQA,tools=self.tools)
+
+    def forward(self, question):
+        # Use Turso to retrieve relevant notes
+        embedded_question = self.embeddings_model._post_request(question)
+        results = retrieve_from_turso(embedded_question, k=5)  # k is limit to return
+
+        # Format context as before
+        context_parts = []
+        for i, row in enumerate(results):
+            source = row[0]  # file_path
+            synopsis = row[1]  # synopsis
+            tags = row[2]    # tags
+            entities = row[3]  # entities
+            content = row[4]  # chunk_data
+
+
+            context_parts.append(f"""
+--- Chunk {i+1} from {source} ---
+synopsis: {synopsis},
+tags: {tags},
+entities: {entities}
+{content}
+""")
+
+        # print('Closest embedding hits')
+        # for part in context_parts:
+        #     print(part)
+        
+        context = "\n\n".join(context_parts)
+
+        prediction = self.generate_answer(context=context, question=question)
+        return dspy.Prediction(answer=prediction.answer, context=context)
+
+    def load_file(self, file_path) -> str | None:
+        """Load and return specified file."""
+        if os.path.exists(file_path):
+            try:
+                with open(file_path) as file:
+                    return file.read()
+            except Exception:
+                return None
+        else:
+            return None