4. Querying and Retrieving Relevant Context
Purpose:
This retrieves the top 3 most relevant chunks for a given user question:
- Embeds the query the same way the documents were embedded
- Searches the vector DB for closest matches
- Builds a prompt that includes the relevant context
from sentence_transformers import SentenceTransformer
import chromadb
query = input("Ask your question: ")
model = SentenceTransformer('all-MiniLM-L6-v2')
chroma_client = chromadb.Client()
collection = chroma_client.get_collection("rag_chunks")
query_vector = model.encode(query).tolist()
results = collection.query(query_embeddings=[query_vector], n_results=3)
contexts = results['documents'][0]
context_str = "\n---\n".join(contexts)
prompt = f"""
Context:
{context_str}
Question: {query}
Answer:"""
with open("prompt.txt", "w") as f:
f.write(prompt)
print("Prompt saved to prompt.txt")