38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
# rag_query.py
|
|
import os, sys, chromadb
|
|
from openai import OpenAI
|
|
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
|
|
query = " ".join(sys.argv[1:]) or input("Ask Lyra-Archive: ")
|
|
|
|
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
|
|
chroma = chromadb.PersistentClient(path="./chromadb")
|
|
collection = chroma.get_or_create_collection("lyra_chats")
|
|
|
|
# embed the question
|
|
q_emb = client.embeddings.create(
|
|
model="text-embedding-3-small",
|
|
input=query
|
|
).data[0].embedding
|
|
|
|
# search the collection
|
|
results = collection.query(query_embeddings=[q_emb], n_results=5)
|
|
|
|
print("\n🔍 Top related excerpts:\n")
|
|
for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
|
|
print(f"📄 {meta['source']} ({meta['role']}) — {meta['title']}")
|
|
print(doc[:300].strip(), "\n---")
|
|
|
|
# synthesize an answer
|
|
context = "\n\n".join(results["documents"][0])
|
|
answer = client.chat.completions.create(
|
|
model="gpt-4o-mini",
|
|
messages=[
|
|
{"role": "system", "content": "Answer based only on the context below. Be concise and practical."},
|
|
{"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"}
|
|
]
|
|
).choices[0].message.content
|
|
|
|
print("\n💡 Lyra-Archive Answer:\n", answer)
|