From 9814161edb29654e09bc44cd2a541f3ee774f3e9 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Thu, 27 Oct 2022 06:49:50 -0700 Subject: [PATCH] update tests --- tests/unit_tests/chains/test_natbot.py | 4 ++-- tests/unit_tests/chains/test_react.py | 2 +- tests/unit_tests/llms/fake_llm.py | 6 +++--- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/unit_tests/chains/test_natbot.py b/tests/unit_tests/chains/test_natbot.py index 5d46add2..524ed17f 100644 --- a/tests/unit_tests/chains/test_natbot.py +++ b/tests/unit_tests/chains/test_natbot.py @@ -12,9 +12,9 @@ class FakeLLM(LLM): 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 [CompletionOutput("foo")] + return [CompletionOutput(text="foo")] else: - return [CompletionOutput("bar")] + return [CompletionOutput(text="bar")] def test_proper_inputs() -> None: diff --git a/tests/unit_tests/chains/test_react.py b/tests/unit_tests/chains/test_react.py index d38975ad..3dcd89d4 100644 --- a/tests/unit_tests/chains/test_react.py +++ b/tests/unit_tests/chains/test_react.py @@ -33,7 +33,7 @@ class FakeListLLM(LLM): 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 [CompletionOutput(self.responses[self.i])] + return [CompletionOutput(text=self.responses[self.i])] class FakeDocstore(Docstore): diff --git a/tests/unit_tests/llms/fake_llm.py b/tests/unit_tests/llms/fake_llm.py index d9d2846b..ce5520dc 100644 --- a/tests/unit_tests/llms/fake_llm.py +++ b/tests/unit_tests/llms/fake_llm.py @@ -14,8 +14,8 @@ class FakeLLM(LLM): 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] + return [CompletionOutput(text=self._queries[prompt])] if stop is None: - return [CompletionOutput("foo")] + return [CompletionOutput(text="foo")] else: - return [CompletionOutput("bar")] + return [CompletionOutput(text="bar")]