diff --git a/src/agent.py b/src/agent.py index 1759f93..e9e0b8e 100644 --- a/src/agent.py +++ b/src/agent.py @@ -1,38 +1,45 @@ import os -from typing import Dict +from typing import Dict, List -from langchain.agents import AgentExecutor, create_openai_tools_agent +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 tools +# Define tool functions -def chroma_tool(query: str): +def chroma_tool(query: str) -> List[str]: + """Search FAQ documents in Chroma. Use when question about course content.""" return search_course_docs(query) -# Wrap MCP tool -def mcp_tool(query: str): +def mcp_tool(query: str) -> Dict: + """Get course metadata such as schedule. Use when question about schedule or instructor.""" return fetch_course_meta(query) -# Simple tool registry -TOOL_DEFINITIONS = { - "search_course_docs": { - "func": chroma_tool, - "description": "Search FAQ documents in Chroma. Use when question about course content.", - }, - "fetch_course_meta": { - "func": mcp_tool, - "description": "Get course metadata such as schedule. Use when question about schedule or instructor.", - }, -} +# 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'.""" +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), @@ -41,8 +48,8 @@ PROMPT = ChatPromptTemplate.from_messages([ # Create agent llm = Ollama(model="llama2") -agent = create_openai_tools_agent(llm, TOOL_DEFINITIONS, PROMPT) -executor = AgentExecutor(agent=agent, tools=list(TOOL_DEFINITIONS.values()), verbose=True) +agent = create_agent(llm, TOOL_DEFINITIONS, PROMPT) +executor = AgentExecutor(agent=agent, tools=TOOL_DEFINITIONS, verbose=True) # CLI entry point if __name__ == "__main__":