add main.py
This commit is contained in:
@@ -0,0 +1,64 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user