[test] Add integration_test for PandasAgent (#4056)

- confirm creation
- confirm functionality with a simple dimension check.

The test now is calling OpenAI API directly, but learning from
@vowelparrot that we’re caching the requests, so that it’s not that
expensive. I also found we’re calling OpenAI api in other integration
tests. Please lmk if there is any concern of real external API calls. I
can alternatively make a fake LLM for this test. Thanks
parallel_dir_loader
Mike Wang 1 year ago committed by GitHub
parent 6567b73e1a
commit c3044b1bf0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -0,0 +1,30 @@
import re
import numpy as np
import pytest
from pandas import DataFrame
from langchain.agents import create_pandas_dataframe_agent
from langchain.agents.agent import AgentExecutor
from langchain.llms import OpenAI
@pytest.fixture(scope="module")
def df() -> DataFrame:
random_data = np.random.rand(4, 4)
df = DataFrame(random_data, columns=["name", "age", "food", "sport"])
return df
def test_pandas_agent_creation(df: DataFrame) -> None:
agent = create_pandas_dataframe_agent(OpenAI(temperature=0), df)
assert isinstance(agent, AgentExecutor)
def test_data_reading(df: DataFrame) -> None:
agent = create_pandas_dataframe_agent(OpenAI(temperature=0), df)
assert isinstance(agent, AgentExecutor)
response = agent.run("how many rows in df? Give me a number.")
result = re.search(rf".*({df.shape[0]}).*", response)
assert result is not None
assert result.group(1) is not None
Loading…
Cancel
Save