Files
povtornyy-ekzamen-faq-bot-c…/src/agent.py
T

46 lines
1.8 KiB
Python

import os
from typing import List, Dict, Any
from langchain_community.llms import Ollama
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from langchain_core.tools import BaseTool
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain.schema import HumanMessage, SystemMessage
from .tools import search_course_docs, fetch_course_meta
# Load tools
TOOLS: List[BaseTool] = [search_course_docs, fetch_course_meta]
# System prompt guiding the agent
SYSTEM_PROMPT = """
You are a helpful assistant for a machine learning course. Your job is to answer user questions.
- If the question is about course materials, lecture slides, assignments, or any content that can be found in the FAQ documents, use the tool `search_course_docs`.
- If the question is about course schedule, instructor information, or other metadata, use the tool `fetch_course_meta`.
- Do not use both tools unless absolutely necessary.
- In your answer, always include a source tag: `source: chroma` if you used the FAQ tool, or `source: mcp_meta` if you used the metadata tool.
"""
def build_agent() -> AgentExecutor:
"""
Build and return a LangChain AgentExecutor with the defined tools and system prompt.
"""
llm = Ollama(model="llama3", temperature=0.0)
# Prompt template
prompt = ChatPromptTemplate.from_messages(
[
SystemMessage(content=SYSTEM_PROMPT),
MessagesPlaceholder(variable_name="history"),
HumanMessage(content="{input}"),
]
)
# Create the agent
agent = create_openai_tools_agent(llm=llm, tools=TOOLS, prompt=prompt)
# Wrap with AgentExecutor
agent_executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True, handle_parsing_errors=True)
return agent_executor