From ccb719708b1b59f86fcb7a94388abdd1833060dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Fri, 5 Jun 2026 11:49:36 +0000 Subject: [PATCH] Add main.py --- main.py | 82 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 82 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..f703664 --- /dev/null +++ b/main.py @@ -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") +"" \ No newline at end of file