diff --git a/src/agent.py b/src/agent.py index 3e6ad73..c0b50ba 100644 --- a/src/agent.py +++ b/src/agent.py @@ -1,57 +1,57 @@ -import os -from typing import Dict, Any -from langchain_ollama import OllamaLLM -from langchain.agents import initialize_agent, Tool, AgentType -from langchain.memory import ConversationBufferMemory -from src.utils import search_course_docs, fetch_course_meta +import argparse +from langchain.agents import create_openai_functions_agent, AgentExecutor +from langchain_core.prompts import ChatPromptTemplate +from langchain_core.messages import HumanMessage +from langchain_ollama import Ollama +from src.utils import load_faq_to_chroma, search_course_docs, fetch_course_meta + +# Initialize embeddings and LLM +llm = Ollama(model="llama3.1") + +# Load or create Chroma collection +try: + chroma = load_faq_to_chroma() +except Exception: + chroma = None # Define tools -search_tool = Tool( - name="search_course_docs", - func=search_course_docs, - description="Search local FAQ docs in ChromaDB. Use when question about course content. Returns list of relevant documents." -) -meta_tool = Tool( - name="fetch_course_meta", - func=fetch_course_meta, - description="Fetch course metadata (schedule, exams) from MCP-style tool. Use when question about schedule or metadata. Returns list of matching items." -) +from langchain.tools import tool -# System prompt to guide routing -SYSTEM_PROMPT = ( - "You are an FAQ bot for the course. Use search_course_docs for content questions and fetch_course_meta for schedule or metadata questions.\n" - "When answering, include a field 'source' with value 'chroma' or 'mcp_meta' to indicate which tool was used." -) +@tool +def search_course_docs_tool(query: str, k: int = 3) -> str: + """Search local FAQ docs in ChromaDB.""" + docs = search_course_docs(query, k) + return "\n".join([doc.page_content for doc in docs]) -# LLM and agent setup -llm = OllamaLLM(model="llama3") -memory = ConversationBufferMemory(memory_key="chat_history") -agent = initialize_agent( - tools=[search_tool, meta_tool], - llm=llm, - agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION, - memory=memory, - verbose=True, - system_prompt=SYSTEM_PROMPT, -) +@tool +def fetch_course_meta_tool(query: str) -> str: + """Fetch course metadata via MCP-style tool.""" + results = fetch_course_meta(query) + return str(results) -def ask(question: str) -> Dict[str, Any]: - response = agent.run(question) - # Parse response to extract source if present - source = "unknown" - if "source:" in response.lower(): - parts = response.lower().split("source:") - source = parts[1].strip().split()[0] - return {"answer": response, "source": source} +tools = [search_course_docs_tool, fetch_course_meta_tool] + +# Prompt template with source hint +prompt = ChatPromptTemplate.from_messages([ + ("system", "You are a helpful FAQ assistant. Use the tools only when necessary. In your answer, include a line like 'source: chroma' or 'source: mcp_meta' to indicate which tool was used.") +]) + +agent = create_openai_functions_agent(llm=llm, tools=tools, prompt=prompt) +executor = AgentExecutor(agent=agent, tools=tools, verbose=True) if __name__ == "__main__": - # Simple CLI with 3 preset questions - questions = [ - "Как подключить ChromaDB?", - "Что такое MCP‑tool?", - "Когда проходят экзамены?" - ] - for q in questions: - print("Q:", q) - print("A:", ask(q)["answer"], "(source:", ask(q)["source"], ")") - print() + parser = argparse.ArgumentParser(description="FAQ bot CLI") + parser.add_argument("--question", type=str, help="Question to ask the bot") + args = parser.parse_args() + if args.question: + response = executor.invoke({"input": args.question}) + print(response["output"]) + else: + # Interactive mode + print("FAQ Bot. Type 'exit' to quit.") + while True: + q = input("> ") + if q.lower() in ("exit", "quit"): + break + resp = executor.invoke({"input": q}) + print(resp["output"])