feat: ✨ AI Powered enhanced queries to get better results

src/experts/retrieval_agent.py

@@ -1,7 +1,7 @@
 import os
-import turso
-import dspy
 
+import dspy
+import turso
 
 from config_loader import load_config
 from embedding import LocalLMEmbeddings
@@ -9,27 +9,28 @@ from embedding import LocalLMEmbeddings
 CFG = load_config()
 
 DATABASE_PATH = CFG["ingestion"]["db_path"]
+DATABASE_NAME = CFG["ingestion"]["db_name"]
 EMBEDDING_MODEL = CFG["models"]["embedding"]
 API_BASE = CFG["api"]["base_url"]
 RETRIEVAL_CONFIG = CFG["retrieval_agent"]
+EXPANSION_CONFIG = CFG["expansion_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
+    vector_distance_cos(embedding, vector32('{embedded_question}')) AS distance
     FROM notes
     ORDER BY distance ASC
     LIMIT {k};
     """
-    con = turso.connect(DATABASE_PATH)
+    con = turso.connect(DATABASE_PATH + DATABASE_NAME)
     cur = con.cursor()
     cur.execute(query)
     rows = cur.fetchall()
     return rows
 
 
-# --- DSPy Signature ---
 class DnDContextQA(dspy.Signature):
     f"{RETRIEVAL_CONFIG['retrieval_signature']}"
 
@@ -38,47 +39,71 @@ class DnDContextQA(dspy.Signature):
     answer = dspy.OutputField(desc="A detailed answer based on the notes, citing the source file.")
 
 
+class ExpansionSignature(dspy.Signature):
+    f"{EXPANSION_CONFIG['expansion_signature']}"
+    question = dspy.InputField()
+    answer = dspy.OutputField(
+        desc="A list of questions that will be used to vector search the database."
+    )
+
+
 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.
+            # batch_size=1,
         )
-        # Tools exposed to the ReAct loop
+        self.retrieval_lm = dspy.LM(
+            model=CFG["models"]["retrieval"], api_base=API_BASE + CFG["api"]["api_version"]
+        )
+        with dspy.context(lm=self.retrieval_lm, signature=ExpansionSignature):
+            self.query_expander = dspy.Predict("question -> queries:list[str]")
+
         self.tools = [self.load_file]
         self.generate_answer = dspy.ReAct(signature=DnDContextQA, tools=self.tools)
 
     def forward(self, question):
-        # TODO: Add step here to LLM Expand
-        # given the current question, generate 3-5 distinct search queries.
-        # embed all the questions
-        embedded_question = self.embeddings_model._post_request(question)
-        # store the 5 from all 3-5 questions (15 - 25 results)
-        results = retrieve_from_turso(embedded_question, k=5)  # k is limit to return
+        print("Enhancing Question")
+        with dspy.context(lm=self.retrieval_lm):
+            expanded_queries = self.query_expander(question=question).queries
+        print("Enhanced Queries:")
+        for q in expanded_queries:
+            print("    ", q)
+        all_embeddings = self.embeddings_model.embed_documents([question] + expanded_queries)
+        # print(all_embeddings)
+        all_results = []
+        for embedded_question in all_embeddings:
+            results = retrieve_from_turso(embedded_question, k=5)
+            all_results.extend(results)
+
+        seen = set()
+        unique_results = []
+        for row in all_results:
+            key = (row[0], row[4])
+            if key not in seen:
+                seen.add(key)
+                unique_results.append(row)
 
-        # 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
+        for i, row in enumerate(unique_results):
+            source = row[0]
+            synopsis = row[1]
+            tags = row[2]
+            entities = row[3]
+            content = row[4]
+            closeness = row[5]
 
             context_parts.append(f"""
 --- Chunk {i + 1} from {source} ---
 synopsis: {synopsis},
 tags: {tags},
-entities: {entities}
+entities: {entities},
+closeness: {closeness},
 {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)