72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
import os
|
|
from typing import Dict, List
|
|
|
|
from langchain.agents import AgentExecutor, create_agent
|
|
from langchain_ollama import Ollama
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
from langchain_core.messages import HumanMessage
|
|
from langchain.tools import Tool
|
|
|
|
from src.utils import search_course_docs
|
|
from src.mcp_tool import fetch_course_meta
|
|
|
|
# Define tool functions
|
|
|
|
def chroma_tool(query: str) -> List[str]:
|
|
"""Search FAQ documents in Chroma. Use when question about course content."""
|
|
return search_course_docs(query)
|
|
|
|
|
|
def mcp_tool(query: str) -> Dict:
|
|
"""Get course metadata such as schedule. Use when question about schedule or instructor."""
|
|
return fetch_course_meta(query)
|
|
|
|
# Create Tool objects
|
|
TOOL_DEFINITIONS = [
|
|
Tool.from_function(
|
|
func=chroma_tool,
|
|
name="search_course_docs",
|
|
description="Search FAQ documents in Chroma. Use when question about course content.",
|
|
),
|
|
Tool.from_function(
|
|
func=mcp_tool,
|
|
name="fetch_course_meta",
|
|
description="Get course metadata such as schedule. Use when question about schedule or instructor.",
|
|
),
|
|
]
|
|
|
|
# Prompt template
|
|
SYSTEM_PROMPT = (
|
|
"You are a FAQ bot for the course. Use only the tools provided. Do not call both tools unless necessary."
|
|
" Indicate source in your answer: 'source: chroma' or 'source: mcp_meta'."
|
|
)
|
|
|
|
PROMPT = ChatPromptTemplate.from_messages([
|
|
("system", SYSTEM_PROMPT),
|
|
("human", "{input}"),
|
|
])
|
|
|
|
# Create agent
|
|
llm = Ollama(model="llama2")
|
|
agent = create_agent(llm, TOOL_DEFINITIONS, PROMPT)
|
|
executor = AgentExecutor(agent=agent, tools=TOOL_DEFINITIONS, verbose=True)
|
|
|
|
# CLI entry point
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
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:
|
|
result = executor.invoke({"input": args.question})
|
|
print(result["output"])
|
|
else:
|
|
print("Enter questions (Ctrl-D to exit):")
|
|
try:
|
|
while True:
|
|
q = input("Q: ")
|
|
res = executor.invoke({"input": q})
|
|
print("A:", res["output"])
|
|
except EOFError:
|
|
pass
|