65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
from vectorstore import create_vectorstore, load_documents, collection_exists
|
|
from agent import create_agent
|
|
|
|
def main():
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# Configuration
|
|
qdrant_path = os.getenv("QDRANT_PATH", "./qdrant_db")
|
|
embedding_model = os.getenv("EMBEDDING_MODEL", "nomic-embed-text")
|
|
llm_model = os.getenv("LLM_MODEL", "llama3")
|
|
tavily_api_key = os.getenv("TAVILY_API_KEY")
|
|
if not tavily_api_key:
|
|
print("Error: TAVILY_API_KEY not set in .env")
|
|
sys.exit(1)
|
|
|
|
# Create vector store
|
|
vectorstore = create_vectorstore(
|
|
persist_directory=qdrant_path,
|
|
collection_name="documents",
|
|
embedding_model=embedding_model,
|
|
)
|
|
|
|
# Load documents if collection is empty
|
|
if not collection_exists(vectorstore):
|
|
print("Loading documents into Qdrant...")
|
|
docs_dir = Path("documents")
|
|
if not docs_dir.exists():
|
|
print(f"Documents directory '{docs_dir}' not found.")
|
|
sys.exit(1)
|
|
load_documents(str(docs_dir), vectorstore)
|
|
print("Documents loaded.")
|
|
else:
|
|
print("Qdrant collection already exists. Skipping document load.")
|
|
|
|
# Create agent
|
|
agent = create_agent(vectorstore, tavily_api_key, llm_model=llm_model)
|
|
|
|
print("Chat agent ready. Type 'exit' to quit.")
|
|
while True:
|
|
try:
|
|
user_input = input("\nYou: ")
|
|
except (KeyboardInterrupt, EOFError):
|
|
print("\nExiting.")
|
|
break
|
|
|
|
if user_input.strip().lower() in {"exit", "quit"}:
|
|
print("Goodbye!")
|
|
break
|
|
|
|
# Run agent
|
|
try:
|
|
response = agent.run(user_input)
|
|
print(f"\nAssistant: {response}")
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
main() |