diff --git a/main.py b/main.py new file mode 100644 index 0000000..bb8c6e6 --- /dev/null +++ b/main.py @@ -0,0 +1,51 @@ +import argparse + +from dotenv import load_dotenv + +from agent import agent +from vectorstore import create_vectorstore, load_documents + + +def ingest_documents(directory: str) -> None: + vectorstore = create_vectorstore() + chunk_count = load_documents(directory, vectorstore) + print(f"Loaded {chunk_count} chunks from {directory} into ChromaDB.") + + +def run_chat() -> None: + print("RAG agent ready. Type 'exit' to quit.") + while True: + user_input = input("You: ").strip() + if user_input.lower() == "exit": + print("Goodbye!") + return + if not user_input: + continue + + result = agent.invoke( + {"messages": [{"role": "user", "content": user_input}]} + ) + final_message = result["messages"][-1] + print(f"Assistant: {final_message.content}") + + +def main() -> None: + load_dotenv() + + parser = argparse.ArgumentParser() + parser.add_argument( + "--ingest", + metavar="DIRECTORY", + help="Load .txt and .md files from DIRECTORY into ChromaDB before the demo.", + ) + args = parser.parse_args() + + if args.ingest: + ingest_documents(args.ingest) + return + + run_chat() + + +if __name__ == "__main__": + main()