mirror of https://github.com/hwchase17/langchain
[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. Thankspull/4198/head
parent
6567b73e1a
commit
c3044b1bf0
@ -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…
Reference in New Issue