feat: ✨ Working PoC of the Dungeon Masters Vault

src/retrieve.py

@@ -0,0 +1,109 @@
+import sys
+import dspy
+from langchain_community.vectorstores import FAISS
+from embedding import LocalLMEmbeddings
+from pathlib import Path
+
+# --- DSPy Signature ---
+class DnDContextQA(dspy.Signature):
+    """Answer DnD campaign questions using provided snippets and full file context."""
+    context = dspy.InputField(desc="Relevant chunks and full file contents from the campaign notes.")
+    question = dspy.InputField()
+    answer = dspy.OutputField(desc="A detailed answer based on the notes, citing the source file.")
+
+# --- DSPy Module ---
+class DnDRAG(dspy.Module):
+    def __init__(self, db_path="./local_faiss_db", k=3):
+        super().__init__()
+        # 1. Setup Embeddings & Load FAISS
+        self.embeddings = LocalLMEmbeddings(
+            model="text-embedding-qwen3-embedding-8b", 
+            base_url="http://192.168.0.49:1234"
+        )
+        self.vectorstore = FAISS.load_local(
+            db_path, self.embeddings, allow_dangerous_deserialization=True
+        )
+        self.k = k
+        
+        # 2. Setup the Predictor (Chain of Thought for better reasoning)
+        self.generate_answer = dspy.ChainOfThought(DnDContextQA)
+
+    def get_full_file_content(self, file_path):
+        """Helper to read the full source file if it exists."""
+        try:
+            return Path(file_path).read_text(encoding='utf-8')
+        except Exception:
+            return ""
+
+    def forward(self, question):
+        # 1. Search for top-k chunks
+        results = self.vectorstore.similarity_search(question, k=self.k)
+        
+        # 2. Extract unique file paths to load "Full Context"
+        # This prevents the LLM from being 'blind' to the rest of a relevant session note
+        unique_paths = list(set([doc.metadata.get("full_path") for doc in results]))
+        
+        context_parts = []
+        for i, doc in enumerate(results):
+            source = doc.metadata.get("source", "Unknown")
+            context_parts.append(f"--- Chunk {i+1} from {source} ---\n{doc.page_content}")
+
+        # 3. Add the Full Content of the top match (optional, but requested!)
+        # We'll just take the top 1 file to avoid context window explosion
+        if unique_paths:
+            top_file_content = self.get_full_file_content(unique_paths[0])
+            context_parts.append(f"\n=== FULL SOURCE FILE: {Path(unique_paths[0]).name} ===\n{top_file_content[:10000]}")
+
+        # 4. Join everything into one context string
+        context_str = "\n\n".join(context_parts)
+        
+        # 5. Generate Response
+        prediction = self.generate_answer(context=context_str, question=question)
+        return dspy.Prediction(answer=prediction.answer, context=context_str)
+    
+
+
+def main():
+    # 1. Setup the LLM
+    print("🚀 Initializing Qwen-8B via LM Studio...")
+    lm = dspy.LM("lm_studio/qwen/qwen3-8b", api_base="http://192.168.0.49:1234/v1/")
+    dspy.configure(lm=lm)
+
+    # 2. Load the RAG System (only happens once!)
+    print("📚 Loading FAISS index and campaign notes...")
+    try:
+        rag_system = DnDRAG()
+        print("✅ Ready! Ask me anything about the campaign. (Type 'exit' or 'q' to quit)")
+    except Exception as e:
+        print(f"❌ Failed to initialize: {e}")
+        return
+
+    # 3. Interactive Loop
+    while True:
+        try:
+            print("\n" + "─" * 30)
+            query = input("📝 Query: ").strip()
+
+            # Exit conditions
+            if query.lower() in ["exit", "quit", "q"]:
+                print("Farewell, traveler. Good luck on your quest!")
+                break
+            
+            if not query:
+                continue
+
+            print("🔍 Searching and thinking...")
+            response = rag_system(question=query)
+
+            # Print Response
+            print("\n📜 AI RESPONSE:")
+            print(response.answer)
+
+        except KeyboardInterrupt:
+            print("\n\nExiting... See you next session!")
+            sys.exit(0)
+        except Exception as e:
+            print(f"\n⚠️ An error occurred: {e}")
+
+if __name__ == "__main__":
+    main()