test_load.py — langchain Source File

Source Code

"""Test for Serializable base class."""

import pytest
from langchain_core.load.dump import dumpd, dumps
from langchain_core.load.load import load, loads
from langchain_core.prompts.prompt import PromptTemplate

from langchain_classic.chains.llm import LLMChain

pytest.importorskip("langchain_openai", reason="langchain_openai not installed")
pytest.importorskip("langchain_community", reason="langchain_community not installed")

from langchain_community.llms.openai import (  # ignore: community-import
    OpenAI as CommunityOpenAI,
)


class NotSerializable:
    pass


@pytest.mark.requires("openai", "langchain_openai")
def test_loads_openai_llm() -> None:
    from langchain_openai import OpenAI

    llm = CommunityOpenAI(
        model="davinci",
        temperature=0.5,
        openai_api_key="hello",
        top_p=0.8,
    )
    llm_string = dumps(llm)
    llm2 = loads(llm_string, secrets_map={"OPENAI_API_KEY": "hello"})

    assert llm2 == llm
    llm_string_2 = dumps(llm2)
    assert llm_string_2 == llm_string
    assert isinstance(llm2, OpenAI)


@pytest.mark.requires("openai", "langchain_openai")
def test_loads_llmchain() -> None:
    from langchain_openai import OpenAI

    llm = CommunityOpenAI(
        model="davinci",
        temperature=0.5,
        openai_api_key="hello",
        top_p=0.8,
    )
    prompt = PromptTemplate.from_template("hello {name}!")
    chain = LLMChain(llm=llm, prompt=prompt)
    chain_string = dumps(chain)
    chain2 = loads(chain_string, secrets_map={"OPENAI_API_KEY": "hello"})

    assert chain2 == chain
    assert dumps(chain2) == chain_string
    assert isinstance(chain2, LLMChain)
    assert isinstance(chain2.llm, OpenAI)
    assert isinstance(chain2.prompt, PromptTemplate)
// ... (137 more lines)