|
|
import os |
|
|
import re |
|
|
import uuid |
|
|
from dotenv import load_dotenv |
|
|
from langchain.agents import create_agent |
|
|
from langgraph.checkpoint.memory import InMemorySaver |
|
|
from langchain_community.tools import DuckDuckGoSearchRun |
|
|
from langchain.agents.middleware import ( |
|
|
ModelCallLimitMiddleware, |
|
|
ToolCallLimitMiddleware, |
|
|
) |
|
|
from langchain.messages import HumanMessage, AIMessage |
|
|
|
|
|
load_dotenv() |
|
|
|
|
|
|
|
|
class LangChainAgent: |
|
|
def __init__(self): |
|
|
os.environ["GOOGLE_API_KEY"] = os.getenv("GEMINI_API_KEY") |
|
|
|
|
|
system_prompt = """finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. |
|
|
YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. |
|
|
If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. |
|
|
If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. |
|
|
If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.""" |
|
|
|
|
|
self.agent = create_agent( |
|
|
model="google_genai:gemini-2.5-flash", |
|
|
tools=[DuckDuckGoSearchRun()], |
|
|
system_prompt=system_prompt, |
|
|
checkpointer=InMemorySaver(), |
|
|
middleware=[ |
|
|
ModelCallLimitMiddleware(run_limit=10, exit_behavior="end"), |
|
|
ToolCallLimitMiddleware(run_limit=20, exit_behavior="end"), |
|
|
], |
|
|
) |
|
|
|
|
|
def __call__(self, question: str) -> str: |
|
|
print(f"Agent received question (first 50 chars): {question[:50]}...") |
|
|
|
|
|
thread_id = str(uuid.uuid4()) |
|
|
response = self.agent.invoke( |
|
|
{"messages": [HumanMessage(content=question)]}, |
|
|
{"configurable": {"thread_id": thread_id}, "recursion_limit": 50}, |
|
|
) |
|
|
answer = response["messages"][-1].text |
|
|
final_answer = self.extract_final_answer(answer) |
|
|
print(f"Agent returning answer: {final_answer}") |
|
|
return final_answer |
|
|
|
|
|
def extract_final_answer(self, text): |
|
|
""" |
|
|
Extracts the text following 'FINAL ANSWER:' using regex. |
|
|
""" |
|
|
|
|
|
pattern = re.compile(r"FINAL ANSWER:\s*(.*)", re.DOTALL) |
|
|
|
|
|
match = pattern.search(text) |
|
|
|
|
|
if match: |
|
|
|
|
|
|
|
|
return match.group(1).strip() |
|
|
else: |
|
|
return "No final answer provided, here is the complete text : " + text |
|
|
|