Jake/dungeon_masters_vault
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: ✨ AI Powered enhanced queries to get better results

Browse Source
This commit is contained in:
Jake Pullen
2026-03-07 11:08:21 +00:00
parent 58f20856fd
commit 26c0049fd8
6 changed files with 73 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/embedding.py
+2 -1
View File
@@ -1,5 +1,6 @@
import requests
from langchain_core.embeddings import Embeddings
from config_loader import load_config
CFG = load_config()
@@ -37,7 +38,7 @@ class LocalLMEmbeddings(Embeddings):
for i in range(0, len(texts), self.batch_size):
batch = texts[i : i + self.batch_size]
print(f"🚀 Processing batch {(i // self.batch_size) + 1} (Size: {len(batch)})...")
# print(f"🚀 Processing batch {(i // self.batch_size) + 1} (Size: {len(batch)})...")
batch_vectors = self._post_request(batch)
all_embeddings.extend(batch_vectors)
src/experts/expansion_agent.py
View File
src/experts/retrieval_agent.py
+50 -25
View File
@@ -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)
src/ingest.py
+6 -7
View File
@@ -16,6 +16,7 @@ from experts.ingestion_agent import IngestionAgent
CFG = load_config()
DATA_DIR = CFG["ingestion"]["data_dir"]
DATABASE_PATH = CFG["ingestion"]["db_path"]
DATABASE_NAME = CFG["ingestion"]["db_name"]
MODEL_BASE = CFG["models"]["enrich"]
EMBEDDING_MODEL = CFG["models"]["embedding"]
API_BASE = CFG["api"]["base_url"]
@@ -139,13 +140,10 @@ def embed_chunks(chunks: List[Any], batch_size: int = EMBEDDING_BATCH_SIZE) -> L
# Process chunks in batches
for i in tqdm(range(0, total_chunks, batch_size), desc="Embedding batches"):
batch = chunks[i : i + batch_size]
print(f"🚀 Processing batch {(i // batch_size) + 1} (Size: {len(batch)})...")
batch_content = [chunk.page_content for chunk in batch]
try:
# Use model's batched embedding method
# batch_embeddings = embeddings_model.embed_query(batch_content)
batch_embeddings = embeddings_model.embed_documents(batch_content)
# Process each chunk in the batch
for j, (chunk, embedding) in enumerate(zip(batch, batch_embeddings)):
# Extract metadata
@@ -228,7 +226,7 @@ def save_to_db(chunk_dicts):
Each dict maps to a row in the 'notes' table.
"""
print("connecting to db")
con = turso.connect(DATABASE_PATH)
con = turso.connect(DATABASE_PATH + DATABASE_NAME)
print("opening cursor")
cur = con.cursor()
@@ -267,7 +265,8 @@ def save_to_db(chunk_dicts):
def create_db():
con = turso.connect(DATABASE_PATH)
Path(DATABASE_PATH).mkdir(exist_ok=True)
con = turso.connect(DATABASE_PATH + DATABASE_NAME)
cur = con.cursor()
cur.execute("""
@@ -334,7 +333,7 @@ def delete_from_db(embedded_chunks):
print(f"Deleting existing rows for {len(file_paths)} file(s)")
con = turso.connect(DATABASE_PATH)
con = turso.connect(DATABASE_PATH + DATABASE_NAME)
cur = con.cursor()
# Use a single DELETE statement with IN clause for efficiency