Add main.py
This commit is contained in:
@@ -0,0 +1,82 @@
|
|||||||
|
"""Entry point for the FAQ bot.
|
||||||
|
|
||||||
|
The script performs the following steps:
|
||||||
|
1. Loads FAQ markdown files into a persistent Chroma vector store.
|
||||||
|
2. Creates an agent with a system prompt that routes queries to the correct tool.
|
||||||
|
3. Provides a simple CLI with three preset questions (two for the FAQ store and one for metadata).
|
||||||
|
4. Allows interactive querying until the user types /quit.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
from langchain_ollama import ChatOllama
|
||||||
|
from langchain.agents import create_agent
|
||||||
|
from langchain.tools import BaseTool
|
||||||
|
|
||||||
|
from config import CHROMA_PERSIST_DIR
|
||||||
|
from tools import search_course_docs_tool, fetch_course_meta_tool
|
||||||
|
from vector_store import load_faq_to_chroma
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# 1. Load data into Chroma
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
print("Loading FAQ data into Chroma...")
|
||||||
|
load_faq_to_chroma()
|
||||||
|
print("Data loaded.")
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# 2. Define system prompt
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
SYSTEM_PROMPT = (
|
||||||
|
"You are an assistant that answers questions about the course. "
|
||||||
|
"If the question is about lecture materials, use the tool "
|
||||||
|
"search_course_docs. If the question is about course schedule or "
|
||||||
|
"metadata, use fetch_course_meta. After using a tool, answer the user and "
|
||||||
|
"include a source tag: 'source: chroma' or 'source: mcp_meta'. "
|
||||||
|
"Do not use both tools unless the question explicitly requires it."
|
||||||
|
)
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# 3. Create LLM and agent
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
llm = ChatOllama(model="llama3", temperature=0.2)
|
||||||
|
|
||||||
|
# Gather tools
|
||||||
|
TOOLS: List[BaseTool] = [search_course_docs_tool, fetch_course_meta_tool]
|
||||||
|
|
||||||
|
agent = create_agent(model=llm, tools=TOOLS, system_prompt=SYSTEM_PROMPT)
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# 4. CLI with preset questions
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
PRESET_QUESTIONS = [
|
||||||
|
"What topics are covered in Lecture 5?", # Should use chroma
|
||||||
|
"Explain the concept of recursion as described in the notes.", # chroma
|
||||||
|
"What is the schedule for the next week?", # mcp_meta
|
||||||
|
]
|
||||||
|
|
||||||
|
print("\nPreset questions: (type /quit to exit)\n")
|
||||||
|
for i, q in enumerate(PRESET_QUESTIONS, 1):
|
||||||
|
print(f"{i}. {q}")
|
||||||
|
print()
|
||||||
|
|
||||||
|
while True:
|
||||||
|
user_input = input("You: ")
|
||||||
|
if user_input.strip().lower() in {"/quit", "exit", "q"}:
|
||||||
|
print("Goodbye!")
|
||||||
|
break
|
||||||
|
# If user types a number, use preset
|
||||||
|
if user_input.isdigit() and 1 <= int(user_input) <= len(PRESET_QUESTIONS):
|
||||||
|
query = PRESET_QUESTIONS[int(user_input) - 1]
|
||||||
|
else:
|
||||||
|
query = user_input
|
||||||
|
# Invoke agent
|
||||||
|
try:
|
||||||
|
response = agent.invoke({"messages": [{"role": "user", "content": query}]})
|
||||||
|
# The response is a dict with 'messages' list
|
||||||
|
assistant_msg = response["messages"][-1]["content"]
|
||||||
|
print(f"Assistant: {assistant_msg}\n")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Error: {e}\n")
|
||||||
|
""
|
||||||
Reference in New Issue
Block a user