AGAZO_Final_Assignment / npb_agent.py
agazo's picture
Update npb_agent.py
6b82cb6 verified
raw
history blame
2.71 kB
import re
import constants
import time
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.agents import AgentExecutor, create_tool_calling_agent, create_openai_functions_agent
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import SystemMessage
# --- Custom Tools ---
from internet_search_tool import internet_search
PROMPT = """
You are an agent specialized in answering questions related to Nippon Professional Baseball players. Report your thoughts, and
finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER].
When you search the internet, consider the results of this site as the main reference: https://npb.jp.
The first thing you have to do is to find out the player number and team. To that end, search the web using the player name and
"ippon Professional Baseball". One of the results will have the prefix https://npb.jp/bis/eng/players. Go to that page and extract
the player number and player team. They are contained in the html markups with id "pc_v_no" (the number) and "pc_v_team" (the team) of
the returned page.
Then you can search the web by using the team name and the year of the season you're interested. You will get a page like
this one: https://npb.jp/bis/eng/{year}. Go to that page and extract the names of the players. Then search each player for their numbers.
"""
class NpbAgent:
def __init__(self):
llm = ChatGoogleGenerativeAI(
model=constants.MODEL,
api_key=constants.API_KEY,
temperature=0.4,
timeout=20)
tools = [
internet_search
]
prompt = ChatPromptTemplate.from_messages([
SystemMessage(content=PROMPT),
MessagesPlaceholder(variable_name="chat_history"),
("human", "{input}"),
MessagesPlaceholder(variable_name="agent_scratchpad"),
])
agent = create_tool_calling_agent(llm, tools, prompt=prompt)
self.executor = AgentExecutor(
agent=agent,
tools=tools,
verbose=True,
max_iterations=15)
def __call__(self, question: str) -> str:
print(f"NpbAgent agent received: {question[:50]}...")
result = self.executor.invoke({
"input": question,
"chat_history": []
})
output = result.get("output", "No answer returned.")
print(f"Agent response: {output}")
match = re.search(r"FINAL ANSWER:\s*(.*)", output)
if match:
return match.group(1).strip()
else:
return output