File size: 1,633 Bytes
4416e3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | from rag_model import get_qa_chain, HuggingFaceEmbeddings, Chroma
import os
def main():
PERSIST_DIRECTORY = "chroma_db"
if not os.path.exists(PERSIST_DIRECTORY):
print("Knowledge base not built. Building now...")
from rag_model import build_rag_system
vectorstore = build_rag_system()
else:
print("Loading existing knowledge base...")
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
vectorstore = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings)
qa_chain = get_qa_chain(vectorstore)
if not qa_chain:
print("Failed to initialize QA chain. Check your GOOGLE_API_KEY in .env")
return
print("\n--- Retail Product Knowledge Assistant ---")
print("Type 'exit' to quit.")
while True:
query = input("\nAsk a question about our products: ")
if query.lower() == 'exit':
break
print("Thinking...")
try:
result = qa_chain({"query": query})
print(f"\nAnswer: {result['result']}")
print("\nSources (Top Reviews):")
seen_sources = set()
for doc in result["source_documents"]:
source_info = f"{doc.metadata['name']} (Brand: {doc.metadata['brand']})"
if source_info not in seen_sources:
print(f"- {source_info}")
seen_sources.add(source_info)
except Exception as e:
print(f"An error occurred: {e}")
if __name__ == "__main__":
main()
|