mirror of https://github.com/hwchase17/langchain
Harrison/fix create sql agent (#2870)
Co-authored-by: Timothé Pearce <timothe.pearce@gmail.com>pull/2871/head
parent
8a98e5b50b
commit
705596b46a
@ -0,0 +1,18 @@
|
||||
from langchain.agents import create_sql_agent
|
||||
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
|
||||
from langchain.sql_database import SQLDatabase
|
||||
from tests.unit_tests.llms.fake_llm import FakeLLM
|
||||
|
||||
|
||||
def test_create_sql_agent() -> None:
|
||||
db = SQLDatabase.from_uri("sqlite:///:memory:")
|
||||
queries = {"foo": "Final Answer: baz"}
|
||||
llm = FakeLLM(queries=queries, sequential_responses=True)
|
||||
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
|
||||
|
||||
agent_executor = create_sql_agent(
|
||||
llm=llm,
|
||||
toolkit=toolkit,
|
||||
)
|
||||
|
||||
assert agent_executor.run("hello") == "baz"
|
Loading…
Reference in New Issue