Harrison/llm saving (#331)

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

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