add main.py
This commit is contained in:
@@ -0,0 +1,77 @@
|
|||||||
|
"""CLI for the RAG agent with ChromaDB and Tavily web search.
|
||||||
|
|
||||||
|
Commands:
|
||||||
|
/add <text> - add text to local knowledge base
|
||||||
|
/load <file> - load a .txt file into the knowledge base
|
||||||
|
/help - show available commands
|
||||||
|
/quit - exit
|
||||||
|
<question> - ask the agent (auto-selects local KB or web)
|
||||||
|
"""
|
||||||
|
import sys
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
from vectorstore import add_documents # noqa: E402
|
||||||
|
from agent import run_agent # noqa: E402
|
||||||
|
|
||||||
|
HELP_TEXT = """
|
||||||
|
Commands:
|
||||||
|
/add <text> - add text to local ChromaDB knowledge base
|
||||||
|
/load <file> - load a .txt file into the knowledge base
|
||||||
|
/help - show this help
|
||||||
|
/quit - exit
|
||||||
|
<question> - ask the agent (uses Local KB or Web as appropriate)
|
||||||
|
"""
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
print("RAG Agent: ChromaDB + Tavily Web Search")
|
||||||
|
print("Type /help for commands, /quit to exit.\n")
|
||||||
|
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
line = input("You: ").strip()
|
||||||
|
except (EOFError, KeyboardInterrupt):
|
||||||
|
print("\nExiting.")
|
||||||
|
sys.exit(0)
|
||||||
|
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
|
||||||
|
if line == "/quit":
|
||||||
|
print("Goodbye!")
|
||||||
|
break
|
||||||
|
|
||||||
|
elif line == "/help":
|
||||||
|
print(HELP_TEXT)
|
||||||
|
|
||||||
|
elif line.startswith("/add "):
|
||||||
|
text = line[5:].strip()
|
||||||
|
if not text:
|
||||||
|
print("Usage: /add <text>")
|
||||||
|
continue
|
||||||
|
n = add_documents([text])
|
||||||
|
print(f"Added {n} chunk(s) to knowledge base.")
|
||||||
|
|
||||||
|
elif line.startswith("/load "):
|
||||||
|
path = line[6:].strip()
|
||||||
|
try:
|
||||||
|
with open(path, encoding="utf-8") as fh:
|
||||||
|
content = fh.read()
|
||||||
|
n = add_documents([content])
|
||||||
|
print(f"Loaded '{path}': {n} chunk(s) added.")
|
||||||
|
except FileNotFoundError:
|
||||||
|
print(f"File not found: {path}")
|
||||||
|
except Exception as exc:
|
||||||
|
print(f"Error: {exc}")
|
||||||
|
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
response = run_agent(line)
|
||||||
|
print(f"Agent: {response}\n")
|
||||||
|
except Exception as exc:
|
||||||
|
print(f"Agent error: {exc}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user