Updated agent to use modern LangChain API and fixed tool definitions.: update src/agent.py

This commit is contained in:
2026-06-17 19:38:14 +00:00
parent f215acf785
commit f97f72f34d
+27 -20
View File
@@ -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__":