agent.py

from langchain_openai import ChatOpenAI
from langchain_classic.memory import ConversationBufferMemory
from langchain_classic.agents import initialize_agent, AgentType
from langchain_classic.tools import Tool
import pandas as pd


def create_agent(df: pd.DataFrame, openai_api_key: str):

    llm = ChatOpenAI(
        model="gpt-4o-mini",
        temperature=0,
        api_key=openai_api_key
    )

    memory = ConversationBufferMemory(
        memory_key="chat_history",
        return_messages=True
    )

    def get_table_context(_: str) -> str:
        return df.to_markdown()

    tools = [
        Tool(
            name="GetNHMTable",
            func=get_table_context,
            description="Returns the NHM GBIF records table"
        )
    ]

    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        memory=memory,
        verbose=False,
        handle_parsing_errors=True,
        agent_kwargs={
            "system_message": (
                "You are a data analyst. "
                "If the user greets you, reply politely. "
                "If the user asks about the data, use the table returned by tools. "
                "Answer ONLY based on the table."
            )
        }
    )

    return agent