Harrison/fix create sql agent (#2870)

Co-authored-by: Timothé Pearce <timothe.pearce@gmail.com>

tests/unit_tests/agents/test_sql.py

@@ -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"

tests/unit_tests/llms/fake_llm.py

@@ -1,5 +1,7 @@
 """Fake LLM wrapper for testing purposes."""
-from typing import Any, List, Mapping, Optional
+from typing import Any, List, Mapping, Optional, cast
+
+from pydantic import validator
 
 from langchain.llms.base import LLM
 
@@ -8,6 +10,18 @@ class FakeLLM(LLM):
     """Fake LLM wrapper for testing purposes."""
 
     queries: Optional[Mapping] = None
+    sequential_responses: Optional[bool] = False
+    response_index: int = 0
+
+    @validator("queries", always=True)
+    def check_queries_required(
+        cls, queries: Optional[Mapping], values: Mapping[str, Any]
+    ) -> Optional[Mapping]:
+        if values.get("sequential_response") and not queries:
+            raise ValueError(
+                "queries is required when sequential_response is set to True"
+            )
+        return queries
 
     @property
     def _llm_type(self) -> str:
@@ -15,7 +29,9 @@ class FakeLLM(LLM):
         return "fake"
 
     def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
-        """First try to lookup in queries, else return 'foo' or 'bar'."""
+        if self.sequential_responses:
+            return self._get_next_response_in_sequence
+
         if self.queries is not None:
             return self.queries[prompt]
         if stop is None:
@@ -26,3 +42,10 @@ class FakeLLM(LLM):
     @property
     def _identifying_params(self) -> Mapping[str, Any]:
         return {}
+
+    @property
+    def _get_next_response_in_sequence(self) -> str:
+        queries = cast(Mapping, self.queries)
+        response = queries[list(queries.keys())[self.response_index]]
+        self.response_index = self.response_index + 1
+        return response