58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
import argparse
|
|
from langchain.agents.openai_functions import create_openai_functions_agent
|
|
from langchain.agents import AgentExecutor
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
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
|
|
from langchain.tools import tool
|
|
|
|
@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])
|
|
|
|
@tool
|
|
def fetch_course_meta_tool(query: str) -> str:
|
|
"""Fetch course metadata via MCP-style tool."""
|
|
results = fetch_course_meta(query)
|
|
return str(results)
|
|
|
|
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__":
|
|
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"])
|