Updated agent to use modern LangChain API and fixed tool definitions.: update src/agent.py
This commit is contained in:
+27
-20
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user