| 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() | |