Harrison/llm saving (#331)

Co-authored-by: Akash Samant <70665700+asamant21@users.noreply.github.com>

tests/integration_tests/llms/test_ai21.py

@@ -1,6 +1,9 @@
 """Test AI21 API wrapper."""
 
+from pathlib import Path
+
 from langchain.llms.ai21 import AI21
+from langchain.llms.loading import load_llm
 
 
 def test_ai21_call() -> None:
@@ -8,3 +11,11 @@ def test_ai21_call() -> None:
     llm = AI21(maxTokens=10)
     output = llm("Say foo:")
     assert isinstance(output, str)
+
+
+def test_saving_loading_llm(tmp_path: Path) -> None:
+    """Test saving/loading an AI21 LLM."""
+    llm = AI21(maxTokens=10)
+    llm.save(file_path=tmp_path / "ai21.yaml")
+    loaded_llm = load_llm(tmp_path / "ai21.yaml")
+    assert llm == loaded_llm

tests/integration_tests/llms/test_cohere.py

@@ -1,6 +1,10 @@
 """Test Cohere API wrapper."""
 
+from pathlib import Path
+
 from langchain.llms.cohere import Cohere
+from langchain.llms.loading import load_llm
+from tests.integration_tests.llms.utils import assert_llm_equality
 
 
 def test_cohere_call() -> None:
@@ -8,3 +12,11 @@ def test_cohere_call() -> None:
     llm = Cohere(max_tokens=10)
     output = llm("Say foo:")
     assert isinstance(output, str)
+
+
+def test_saving_loading_llm(tmp_path: Path) -> None:
+    """Test saving/loading an Cohere LLM."""
+    llm = Cohere(max_tokens=10)
+    llm.save(file_path=tmp_path / "cohere.yaml")
+    loaded_llm = load_llm(tmp_path / "cohere.yaml")
+    assert_llm_equality(llm, loaded_llm)

tests/integration_tests/llms/test_huggingface_hub.py

@@ -1,8 +1,12 @@
 """Test HuggingFace API wrapper."""
 
+from pathlib import Path
+
 import pytest
 
 from langchain.llms.huggingface_hub import HuggingFaceHub
+from langchain.llms.loading import load_llm
+from tests.integration_tests.llms.utils import assert_llm_equality
 
 
 def test_huggingface_text_generation() -> None:
@@ -24,3 +28,11 @@ def test_huggingface_call_error() -> None:
     llm = HuggingFaceHub(model_kwargs={"max_new_tokens": -1})
     with pytest.raises(ValueError):
         llm("Say foo:")
+
+
+def test_saving_loading_llm(tmp_path: Path) -> None:
+    """Test saving/loading an HuggingFaceHub LLM."""
+    llm = HuggingFaceHub(repo_id="gpt2", model_kwargs={"max_new_tokens": 10})
+    llm.save(file_path=tmp_path / "hf.yaml")
+    loaded_llm = load_llm(tmp_path / "hf.yaml")
+    assert_llm_equality(llm, loaded_llm)

tests/integration_tests/llms/test_nlpcloud.py

@@ -1,6 +1,10 @@
 """Test NLPCloud API wrapper."""
 
+from pathlib import Path
+
+from langchain.llms.loading import load_llm
 from langchain.llms.nlpcloud import NLPCloud
+from tests.integration_tests.llms.utils import assert_llm_equality
 
 
 def test_nlpcloud_call() -> None:
@@ -8,3 +12,11 @@ def test_nlpcloud_call() -> None:
     llm = NLPCloud(max_length=10)
     output = llm("Say foo:")
     assert isinstance(output, str)
+
+
+def test_saving_loading_llm(tmp_path: Path) -> None:
+    """Test saving/loading an NLPCloud LLM."""
+    llm = NLPCloud(max_length=10)
+    llm.save(file_path=tmp_path / "nlpcloud.yaml")
+    loaded_llm = load_llm(tmp_path / "nlpcloud.yaml")
+    assert_llm_equality(llm, loaded_llm)

tests/integration_tests/llms/test_openai.py

@@ -1,7 +1,10 @@
 """Test OpenAI API wrapper."""
 
+from pathlib import Path
+
 import pytest
 
+from langchain.llms.loading import load_llm
 from langchain.llms.openai import OpenAI
 
 
@@ -44,3 +47,11 @@ def test_openai_stop_error() -> None:
     llm = OpenAI(stop="3", temperature=0)
     with pytest.raises(ValueError):
         llm("write an ordered list of five items", stop=["\n"])
+
+
+def test_saving_loading_llm(tmp_path: Path) -> None:
+    """Test saving/loading an OpenAPI LLM."""
+    llm = OpenAI(max_tokens=10)
+    llm.save(file_path=tmp_path / "openai.yaml")
+    loaded_llm = load_llm(tmp_path / "openai.yaml")
+    assert loaded_llm == llm

tests/integration_tests/llms/utils.py

@@ -0,0 +1,16 @@
+"""Utils for LLM Tests."""
+
+from langchain.llms.base import LLM
+
+
+def assert_llm_equality(llm: LLM, loaded_llm: LLM) -> None:
+    """Assert LLM Equality for tests."""
+    # Check that they are the same type.
+    assert type(llm) == type(loaded_llm)
+    # Client field can be session based, so hash is different despite
+    # all other values being the same, so just assess all other fields
+    for field in llm.__fields__.keys():
+        if field != "client":
+            val = getattr(llm, field)
+            new_val = getattr(loaded_llm, field)
+            assert new_val == val

tests/unit_tests/agents/test_agent.py

@@ -2,17 +2,17 @@
 
 from typing import Any, List, Mapping, Optional
 
+from pydantic import BaseModel
+
 from langchain.agents import Tool, initialize_agent
 from langchain.llms.base import LLM
 
 
-class FakeListLLM(LLM):
+class FakeListLLM(LLM, BaseModel):
     """Fake LLM for testing that outputs elements of a list."""
 
-    def __init__(self, responses: List[str]):
-        """Initialize with list of responses."""
-        self.responses = responses
-        self.i = -1
+    responses: List[str]
+    i: int = -1
 
     def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         """Increment counter, and then return response in that index."""
@@ -25,6 +25,11 @@ class FakeListLLM(LLM):
     def _identifying_params(self) -> Mapping[str, Any]:
         return {}
 
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "fake_list"
+
 
 def test_agent_bad_action() -> None:
     """Test react chain when bad action given."""
@@ -33,7 +38,7 @@ def test_agent_bad_action() -> None:
         f"I'm turning evil\nAction: {bad_action_name}\nAction Input: misalignment",
         "Oh well\nAction: Final Answer\nAction Input: curses foiled again",
     ]
-    fake_llm = FakeListLLM(responses)
+    fake_llm = FakeListLLM(responses=responses)
     tools = [
         Tool("Search", lambda x: x, "Useful for searching"),
         Tool("Lookup", lambda x: x, "Useful for looking up things in a table"),

tests/unit_tests/agents/test_react.py

@@ -2,6 +2,8 @@
 
 from typing import Any, List, Mapping, Optional, Union
 
+from pydantic import BaseModel
+
 from langchain.agents.react.base import ReActChain, ReActDocstoreAgent
 from langchain.agents.tools import Tool
 from langchain.docstore.base import Docstore
@@ -20,13 +22,16 @@ Made in 2022."""
 _FAKE_PROMPT = PromptTemplate(input_variables=["input"], template="{input}")
 
 
-class FakeListLLM(LLM):
+class FakeListLLM(LLM, BaseModel):
     """Fake LLM for testing that outputs elements of a list."""
 
-    def __init__(self, responses: List[str]):
-        """Initialize with list of responses."""
-        self.responses = responses
-        self.i = -1
+    responses: List[str]
+    i: int = -1
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "fake_list"
 
     def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         """Increment counter, and then return response in that index."""
@@ -50,7 +55,7 @@ class FakeDocstore(Docstore):
 def test_predict_until_observation_normal() -> None:
     """Test predict_until_observation when observation is made normally."""
     outputs = ["foo\nAction 1: Search[foo]"]
-    fake_llm = FakeListLLM(outputs)
+    fake_llm = FakeListLLM(responses=outputs)
     tools = [
         Tool("Search", lambda x: x),
         Tool("Lookup", lambda x: x),
@@ -65,7 +70,7 @@ def test_predict_until_observation_normal() -> None:
 def test_predict_until_observation_repeat() -> None:
     """Test when no action is generated initially."""
     outputs = ["foo", " Search[foo]"]
-    fake_llm = FakeListLLM(outputs)
+    fake_llm = FakeListLLM(responses=outputs)
     tools = [
         Tool("Search", lambda x: x),
         Tool("Lookup", lambda x: x),
@@ -84,7 +89,7 @@ def test_react_chain() -> None:
         "I should probably lookup\nAction 2: Lookup[made]",
         "Ah okay now I know the answer\nAction 3: Finish[2022]",
     ]
-    fake_llm = FakeListLLM(responses)
+    fake_llm = FakeListLLM(responses=responses)
     react_chain = ReActChain(llm=fake_llm, docstore=FakeDocstore())
     output = react_chain.run("when was langchain made")
     assert output == "2022"
@@ -97,7 +102,7 @@ def test_react_chain_bad_action() -> None:
         f"I'm turning evil\nAction 1: {bad_action_name}[langchain]",
         "Oh well\nAction 2: Finish[curses foiled again]",
     ]
-    fake_llm = FakeListLLM(responses)
+    fake_llm = FakeListLLM(responses=responses)
     react_chain = ReActChain(llm=fake_llm, docstore=FakeDocstore())
     output = react_chain.run("when was langchain made")
     assert output == "curses foiled again"

tests/unit_tests/chains/test_natbot.py

@@ -2,11 +2,13 @@
 
 from typing import Any, List, Mapping, Optional
 
+from pydantic import BaseModel
+
 from langchain.chains.natbot.base import NatBotChain
 from langchain.llms.base import LLM
 
 
-class FakeLLM(LLM):
+class FakeLLM(LLM, BaseModel):
     """Fake LLM wrapper for testing purposes."""
 
     def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str:
@@ -16,6 +18,11 @@ class FakeLLM(LLM):
         else:
             return "bar"
 
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "fake"
+
     @property
     def _identifying_params(self) -> Mapping[str, Any]:
         return {}

tests/unit_tests/llms/fake_llm.py

@@ -1,20 +1,25 @@
 """Fake LLM wrapper for testing purposes."""
 from typing import Any, List, Mapping, Optional
 
+from pydantic import BaseModel
+
 from langchain.llms.base import LLM
 
 
-class FakeLLM(LLM):
+class FakeLLM(LLM, BaseModel):
     """Fake LLM wrapper for testing purposes."""
 
-    def __init__(self, queries: Optional[Mapping] = None):
-        """Initialize with optional lookup of queries."""
-        self._queries = queries
+    queries: Optional[Mapping] = None
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of 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._queries is not None:
-            return self._queries[prompt]
+        if self.queries is not None:
+            return self.queries[prompt]
         if stop is None:
             return "foo"
         else:

tests/unit_tests/llms/test_loading.py

@@ -0,0 +1,15 @@
+"""Test LLM saving and loading functions."""
+from pathlib import Path
+from unittest.mock import patch
+
+from langchain.llms.loading import load_llm
+from tests.unit_tests.llms.fake_llm import FakeLLM
+
+
+@patch("langchain.llms.loading.type_to_cls_dict", {"fake": FakeLLM})
+def test_saving_loading_round_trip(tmp_path: Path) -> None:
+    """Test saving/loading a Fake LLM."""
+    fake_llm = FakeLLM()
+    fake_llm.save(file_path=tmp_path / "fake_llm.yaml")
+    loaded_llm = load_llm(tmp_path / "fake_llm.yaml")
+    assert loaded_llm == fake_llm