interface

tests/unit_tests/chains/test_natbot.py

@@ -3,18 +3,18 @@
 from typing import List, Optional
 
 from langchain.chains.natbot.base import NatBotChain
-from langchain.llms.base import LLM
+from langchain.llms.base import LLM, CompletionOutput
 
 
 class FakeLLM(LLM):
     """Fake LLM wrapper for testing purposes."""
 
-    def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+    def generate(self, prompt: str, stop: Optional[List[str]] = None) -> List[CompletionOutput]:
         """Return `foo` if longer than 10000 words, else `bar`."""
         if len(prompt) > 10000:
-            return "foo"
+            return [CompletionOutput("foo")]
         else:
-            return "bar"
+            return [CompletionOutput("bar")]
 
 
 def test_proper_inputs() -> None:

tests/unit_tests/chains/test_react.py

@@ -8,7 +8,7 @@ from langchain.chains.llm import LLMChain
 from langchain.chains.react.base import ReActChain, predict_until_observation
 from langchain.docstore.base import Docstore
 from langchain.docstore.document import Document
-from langchain.llms.base import LLM
+from langchain.llms.base import LLM, CompletionOutput
 from langchain.prompt import Prompt
 
 _PAGE_CONTENT = """This is a page about LangChain.
@@ -30,10 +30,10 @@ class FakeListLLM(LLM):
         self.responses = responses
         self.i = -1
 
-    def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+    def generate(self, prompt: str, stop: Optional[List[str]] = None) -> List[CompletionOutput]:
         """Increment counter, and then return response in that index."""
         self.i += 1
-        return self.responses[self.i]
+        return [CompletionOutput(self.responses[self.i])]
 
 
 class FakeDocstore(Docstore):

tests/unit_tests/llms/fake_llm.py

@@ -1,7 +1,7 @@
 """Fake LLM wrapper for testing purposes."""
 from typing import List, Mapping, Optional
 
-from langchain.llms.base import LLM
+from langchain.llms.base import LLM, CompletionOutput
 
 
 class FakeLLM(LLM):
@@ -11,11 +11,11 @@ class FakeLLM(LLM):
         """Initialize with optional lookup of queries."""
         self._queries = queries
 
-    def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+    def generate(self, prompt: str, stop: Optional[List[str]] = None) -> List[CompletionOutput]:
         """First try to lookup in queries, else return 'foo' or 'bar'."""
         if self._queries is not None:
             return self._queries[prompt]
         if stop is None:
-            return "foo"
+            return [CompletionOutput("foo")]
         else:
-            return "bar"
+            return [CompletionOutput("bar")]