67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
"""
|
||
CLI for interacting with the RAG agent.
|
||
|
||
This script provides three explicit commands:
|
||
|
||
* ``/add <text>`` – add a piece of text to the knowledge base.
|
||
* ``/search <query>`` – search the knowledge base and display results.
|
||
* ``/quit`` – exit the program.
|
||
|
||
Any other input is forwarded to the underlying agent as a normal question.
|
||
"""
|
||
import asyncio
|
||
from langchain_core.messages import HumanMessage
|
||
from agent import run_agent
|
||
from tools import add_to_knowledge_base, search_knowledge_base
|
||
|
||
async def main() -> None:
|
||
print("RAG Agent CLI. Type /quit to exit.")
|
||
while True:
|
||
user_input = input("You: ")
|
||
stripped = user_input.strip()
|
||
if not stripped:
|
||
continue
|
||
# Handle explicit commands
|
||
if stripped.startswith("/add "):
|
||
text = stripped[5:].strip()
|
||
if not text:
|
||
print("Error: /add requires a non‑empty text.")
|
||
continue
|
||
try:
|
||
await add_to_knowledge_base(text)
|
||
print("✅ Added to knowledge base.")
|
||
except Exception as exc: # pragma: no cover - defensive
|
||
print(f"❌ Failed to add: {exc}")
|
||
elif stripped.startswith("/search "):
|
||
query = stripped[8:].strip()
|
||
if not query:
|
||
print("Error: /search requires a non‑empty query.")
|
||
continue
|
||
try:
|
||
results = await search_knowledge_base(query)
|
||
if not results:
|
||
print("No results found.")
|
||
else:
|
||
for idx, res in enumerate(results, start=1):
|
||
print(f"{idx}. {res}")
|
||
except Exception as exc: # pragma: no cover - defensive
|
||
print(f"❌ Search error: {exc}")
|
||
elif stripped == "/quit":
|
||
print("Goodbye!")
|
||
break
|
||
else:
|
||
# Forward to agent
|
||
msg = HumanMessage(content=user_input)
|
||
try:
|
||
response = await run_agent([msg])
|
||
# ``run_agent`` returns the last assistant message as a string.
|
||
if isinstance(response, str):
|
||
print(f"Agent: {response}")
|
||
else:
|
||
print("Agent returned unexpected format.")
|
||
except Exception as exc: # pragma: no cover - defensive
|
||
print(f"❌ Agent error: {exc}")
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|