120 lines
3.7 KiB
Python
120 lines
3.7 KiB
Python
import os
|
|
|
|
import dspy
|
|
import turso
|
|
|
|
from config_loader import load_config
|
|
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, chunk_data,
|
|
vector_distance_cos(embedding, vector32('{embedded_question}')) AS distance
|
|
FROM notes
|
|
ORDER BY distance ASC
|
|
LIMIT {k};
|
|
"""
|
|
con = turso.connect(DATABASE_PATH + DATABASE_NAME)
|
|
cur = con.cursor()
|
|
cur.execute(query)
|
|
rows = cur.fetchall()
|
|
return rows
|
|
|
|
|
|
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 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,
|
|
)
|
|
self.retrieval_lm = dspy.LM(model=CFG["models"]["retrieval"], api_base=API_BASE)
|
|
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):
|
|
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[3])
|
|
if key not in seen:
|
|
seen.add(key)
|
|
unique_results.append(row)
|
|
|
|
context_parts = []
|
|
for i, row in enumerate(unique_results):
|
|
source = row[0]
|
|
synopsis = row[1]
|
|
tags = row[2]
|
|
# entities = row[3]
|
|
content = row[3]
|
|
closeness = row[4]
|
|
|
|
context_parts.append(f"""
|
|
--- Chunk {i + 1} from {source} ---
|
|
synopsis: {synopsis},
|
|
tags: {tags},
|
|
closeness: {closeness},
|
|
{content}
|
|
""")
|
|
# entities: {entities},
|
|
|
|
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
|