64 lines
1.7 KiB
Python
64 lines
1.7 KiB
Python
import os
|
|
from dotenv import load_dotenv
|
|
from agent import create_agent
|
|
from vectorstore import create_vectorstore, load_documents
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
|
|
def initialize_knowledge_base(documents_dir: str = "./documents"):
|
|
"""Initialize the vectorstore and load documents if not already loaded."""
|
|
vectorstore = create_vectorstore()
|
|
|
|
# Check if vectorstore is empty
|
|
collection_count = vectorstore._collection.count()
|
|
if collection_count == 0 and os.path.exists(documents_dir):
|
|
print(f"Loading documents from {documents_dir}...")
|
|
load_documents(documents_dir, vectorstore)
|
|
else:
|
|
print(f"Vectorstore already contains {collection_count} documents")
|
|
|
|
return vectorstore
|
|
|
|
|
|
def main():
|
|
"""Main CLI chat loop."""
|
|
print("=" * 50)
|
|
print("RAG Agent with ChromaDB and Web Search")
|
|
print("=" * 50)
|
|
|
|
# Initialize knowledge base
|
|
initialize_knowledge_base()
|
|
|
|
# Create agent
|
|
agent = create_agent()
|
|
|
|
print("\nAgent ready. Type 'exit' to quit.\n")
|
|
|
|
while True:
|
|
try:
|
|
query = input("Запрос: ").strip()
|
|
|
|
if query.lower() == "exit":
|
|
print("Goodbye!")
|
|
break
|
|
|
|
if not query:
|
|
continue
|
|
|
|
# Run agent
|
|
result = agent.invoke({"input": query})
|
|
answer = result.get("output", "No answer generated")
|
|
|
|
print(f"\n{answer}\n")
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nGoodbye!")
|
|
break
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |