def test_structured_args_decorator_no_infer_schema() -> None:
    """Test functionality with structured arguments parsed as a decorator."""

    @tool(infer_schema=False)
    def structured_tool_input(
        arg1: int, arg2: Union[float, datetime], opt_arg: Optional[dict] = None
    ) -> str:
        """Return the arguments directly."""
        return f"{arg1}, {arg2}, {opt_arg}"

    assert isinstance(structured_tool_input, BaseTool)
    assert structured_tool_input.name == "structured_tool_input"
    args = {"arg1": 1, "arg2": 0.001, "opt_arg": {"foo": "bar"}}
    with pytest.raises(ToolException):
        assert structured_tool_input.run(args)

def test_structured_single_str_decorator_no_infer_schema() -> None:
    """Test functionality with structured arguments parsed as a decorator."""

    @tool(infer_schema=False)
    def unstructured_tool_input(tool_input: str) -> str:
        """Return the arguments directly."""
        assert isinstance(tool_input, str)
        return f"{tool_input}"

    assert isinstance(unstructured_tool_input, BaseTool)
    assert unstructured_tool_input.args_schema is None
    assert unstructured_tool_input.run("foo") == "foo"

def test_structured_tool_types_parsed() -> None:
    """Test the non-primitive types are correctly passed to structured tools."""

    class SomeEnum(Enum):
        A = "a"
        B = "b"

    class SomeBaseModel(BaseModel):
        foo: str

    @tool
    def structured_tool(
        some_enum: SomeEnum,
        some_base_model: SomeBaseModel,
    ) -> dict:
        """Return the arguments directly."""
        return {
            "some_enum": some_enum,
            "some_base_model": some_base_model,
        }

    assert isinstance(structured_tool, StructuredTool)
    args = {
        "some_enum": SomeEnum.A.value,
        "some_base_model": SomeBaseModel(foo="bar").model_dump(),
    }
    result = structured_tool.run(json.loads(json.dumps(args)))
    expected = {
        "some_enum": SomeEnum.A,
        "some_base_model": SomeBaseModel(foo="bar"),
    }
    assert result == expected

def test_structured_tool_types_parsed_pydantic_v1() -> None:
    """Test the non-primitive types are correctly passed to structured tools."""

    class SomeBaseModel(BaseModelV1):
        foo: str

    class AnotherBaseModel(BaseModelV1):
        bar: str

    @tool
    def structured_tool(some_base_model: SomeBaseModel) -> AnotherBaseModel:
        """Return the arguments directly."""
        return AnotherBaseModel(bar=some_base_model.foo)

    assert isinstance(structured_tool, StructuredTool)

    expected = AnotherBaseModel(bar="baz")
    for arg in [
        SomeBaseModel(foo="baz"),
        SomeBaseModel(foo="baz").dict(),
    ]:
        args = {"some_base_model": arg}
        result = structured_tool.run(args)
        assert result == expected

def test_structured_tool_types_parsed_pydantic_mixed() -> None:
    """Test handling of tool with mixed Pydantic version arguments."""

    class SomeBaseModel(BaseModelV1):
        foo: str

    class AnotherBaseModel(BaseModel):
        bar: str

    with pytest.raises(NotImplementedError):

        @tool
        def structured_tool(
            some_base_model: SomeBaseModel, another_base_model: AnotherBaseModel
        ) -> None:
            """Return the arguments directly."""

def test_base_tool_inheritance_base_schema() -> None:
    """Test schema is correctly inferred when inheriting from BaseTool."""

    class _MockSimpleTool(BaseTool):
        name: str = "simple_tool"
        description: str = "A Simple Tool"

        def _run(self, tool_input: str) -> str:
            return f"{tool_input}"

        async def _arun(self, tool_input: str) -> str:
            raise NotImplementedError

    simple_tool = _MockSimpleTool()
    assert simple_tool.args_schema is None
    expected_args = {"tool_input": {"title": "Tool Input", "type": "string"}}
    assert simple_tool.args == expected_args

def test_tool_lambda_args_schema() -> None:
    """Test args schema inference when the tool argument is a lambda function."""

    tool = Tool(
        name="tool",
        description="A tool",
        func=lambda tool_input: tool_input,
    )
    assert tool.args_schema is None
    expected_args = {"tool_input": {"type": "string"}}
    assert tool.args == expected_args

def test_structured_tool_from_function_docstring() -> None:
    """Test that structured tools can be created from functions."""

    def foo(bar: int, baz: str) -> str:
        """Docstring

        Args:
            bar: the bar value
            baz: the baz value
        """
        raise NotImplementedError

    structured_tool = StructuredTool.from_function(foo)
    assert structured_tool.name == "foo"
    assert structured_tool.args == {
        "bar": {"title": "Bar", "type": "integer"},
        "baz": {"title": "Baz", "type": "string"},
    }

    assert _schema(structured_tool.args_schema) == {
        "properties": {
            "bar": {"title": "Bar", "type": "integer"},
            "baz": {"title": "Baz", "type": "string"},
        },
        "description": inspect.getdoc(foo),
        "title": "foo",
        "type": "object",
        "required": ["bar", "baz"],
    }

    assert foo.__doc__ is not None
    assert structured_tool.description == textwrap.dedent(foo.__doc__.strip())

def test_structured_tool_from_function_docstring_complex_args() -> None:
    """Test that structured tools can be created from functions."""

    def foo(bar: int, baz: list[str]) -> str:
        """Docstring

        Args:
            bar: int
            baz: List[str]
        """
        raise NotImplementedError

    structured_tool = StructuredTool.from_function(foo)
    assert structured_tool.name == "foo"
    assert structured_tool.args == {
        "bar": {"title": "Bar", "type": "integer"},
        "baz": {
            "title": "Baz",
            "type": "array",
            "items": {"type": "string"},
        },
    }

    assert _schema(structured_tool.args_schema) == {
        "properties": {
            "bar": {"title": "Bar", "type": "integer"},
            "baz": {
                "title": "Baz",
                "type": "array",
                "items": {"type": "string"},
            },
        },
        "description": inspect.getdoc(foo),
        "title": "foo",
        "type": "object",
        "required": ["bar", "baz"],
    }

    assert foo.__doc__ is not None
    assert structured_tool.description == textwrap.dedent(foo.__doc__).strip()

def test_structured_tool_lambda_multi_args_schema() -> None:
    """Test args schema inference when the tool argument is a lambda function."""
    tool = StructuredTool.from_function(
        name="tool",
        description="A tool",
        func=lambda tool_input, other_arg: f"{tool_input}{other_arg}",  # type: ignore
    )
    assert tool.args_schema is not None
    expected_args = {
        "tool_input": {"title": "Tool Input"},
        "other_arg": {"title": "Other Arg"},
    }
    assert tool.args == expected_args

def test_tool_partial_function_args_schema() -> None:
    """Test args schema inference when the tool argument is a partial function."""

    def func(tool_input: str, other_arg: str) -> str:
        assert isinstance(tool_input, str)
        assert isinstance(other_arg, str)
        return tool_input + other_arg

    tool = Tool(
        name="tool",
        description="A tool",
        func=partial(func, other_arg="foo"),
    )
    assert tool.run("bar") == "barfoo"

def test_empty_args_decorator() -> None:
    """Test inferred schema of decorated fn with no args."""

    @tool
    def empty_tool_input() -> str:
        """Return a constant."""
        return "the empty result"

    assert isinstance(empty_tool_input, BaseTool)
    assert empty_tool_input.name == "empty_tool_input"
    assert empty_tool_input.args == {}
    assert empty_tool_input.run({}) == "the empty result"

def test_tool_from_function_with_run_manager() -> None:
    """Test run of tool when using run_manager."""

    def foo(bar: str, callbacks: Optional[CallbackManagerForToolRun] = None) -> str:
        """Docstring
        Args:
            bar: str
        """
        assert callbacks is not None
        return "foo" + bar

    handler = FakeCallbackHandler()
    tool = Tool.from_function(foo, name="foo", description="Docstring")

    assert tool.run(tool_input={"bar": "bar"}, run_manager=[handler]) == "foobar"
    assert tool.run("baz", run_manager=[handler]) == "foobaz"


@pytest.mark.skipif(PYDANTIC_MAJOR_VERSION != 1, reason="Testing pydantic v1.")
def test__get_all_basemodel_annotations_v1() -> None:
    A = TypeVar("A")

    class ModelA(BaseModel, Generic[A], extra="allow"):
        a: A

    class ModelB(ModelA[str]):
        b: Annotated[ModelA[dict[str, Any]], "foo"]

    class Mixin:
        def foo(self) -> str:
            return "foo"

    class ModelC(Mixin, ModelB):
        c: dict

    expected = {"a": str, "b": Annotated[ModelA[dict[str, Any]], "foo"], "c": dict}
    actual = _get_all_basemodel_annotations(ModelC)
    assert actual == expected

    expected = {"a": str, "b": Annotated[ModelA[dict[str, Any]], "foo"]}
    actual = _get_all_basemodel_annotations(ModelB)
    assert actual == expected

    expected = {"a": Any}
    actual = _get_all_basemodel_annotations(ModelA)
    assert actual == expected

    expected = {"a": int}
    actual = _get_all_basemodel_annotations(ModelA[int])
    assert actual == expected

    D = TypeVar("D", bound=Union[str, int])

    class ModelD(ModelC, Generic[D]):
        d: Optional[D]

    expected = {
        "a": str,
        "b": Annotated[ModelA[dict[str, Any]], "foo"],
        "c": dict,
        "d": Union[str, int, None],
    }
    actual = _get_all_basemodel_annotations(ModelD)
    assert actual == expected

    expected = {
        "a": str,
        "b": Annotated[ModelA[dict[str, Any]], "foo"],
        "c": dict,
        "d": Union[int, None],
    }
    actual = _get_all_basemodel_annotations(ModelD[int])
    assert actual == expected

def test_tool_annotations_preserved() -> None:
    """Test that annotations are preserved when creating a tool."""

    @tool
    def my_tool(val: int, other_val: Annotated[dict, "my annotation"]) -> str:
        """Tool docstring."""
        return "foo"

    schema = my_tool.get_input_schema()  # type: ignore[attr-defined]

    func = my_tool.func  # type: ignore[attr-defined]

    expected_type_hints = {
        name: hint
        for name, hint in func.__annotations__.items()
        if name in inspect.signature(func).parameters
    }
    assert schema.__annotations__ == expected_type_hints

@pytest.mark.skipif(PYDANTIC_MAJOR_VERSION != 2, reason="Testing pydantic v2.")
def test_tool_args_schema_pydantic_v2_with_metadata() -> None:
    from pydantic import BaseModel as BaseModelV2
    from pydantic import Field as FieldV2
    from pydantic import ValidationError as ValidationErrorV2

    class Foo(BaseModelV2):
        x: list[int] = FieldV2(
            description="List of integers", min_length=10, max_length=15
        )

    @tool(args_schema=Foo)
    def foo(x):  # type: ignore[no-untyped-def]
        """foo"""
        return x

    assert foo.tool_call_schema.model_json_schema() == {
        "description": "foo",
        "properties": {
            "x": {
                "description": "List of integers",
                "items": {"type": "integer"},
                "maxItems": 15,
                "minItems": 10,
                "title": "X",
                "type": "array",
            }
        },
        "required": ["x"],
        "title": "foo",
        "type": "object",
    }

    assert foo.invoke({"x": [0] * 10})
    with pytest.raises(ValidationErrorV2):
        foo.invoke({"x": [0] * 9})

def test_imports() -> None:
    expected_all = [
        "FILTERED_ARGS",
        "SchemaAnnotationError",
        "create_schema_from_function",
        "ToolException",
        "BaseTool",
        "Tool",
        "StructuredTool",
        "tool",
        "RetrieverInput",
        "create_retriever_tool",
        "ToolsRenderer",
        "render_text_description",
        "render_text_description_and_args",
        "BaseToolkit",
        "convert_runnable_to_tool",
        "InjectedToolArg",
    ]
    for module_name in expected_all:
        assert hasattr(tools, module_name) and getattr(tools, module_name) is not None

def test_structured_tool_direct_init() -> None:
    def foo(bar: str) -> str:
        return bar

    async def async_foo(bar: str) -> str:
        return bar

    class FooSchema(BaseModel):
        bar: str = Field(..., description="The bar")

    tool = StructuredTool(name="foo", args_schema=FooSchema, coroutine=async_foo)

    with pytest.raises(NotImplementedError):
        assert tool.invoke("hello") == "hello"

def test_injected_arg_with_complex_type() -> None:
    """Test that an injected tool arg can be a complex type."""

    class Foo:
        def __init__(self) -> None:
            self.value = "bar"

    @tool
    def injected_tool(x: int, foo: Annotated[Foo, InjectedToolArg]) -> str:
        """Tool that has an injected tool arg."""
        return foo.value

    assert injected_tool.invoke({"x": 5, "foo": Foo()}) == "bar"  # type: ignore


"""Test for some custom pydantic decorators."""

from typing import Any, Optional

import pytest
from pydantic import ConfigDict

from langchain_core.utils.pydantic import (
    PYDANTIC_MAJOR_VERSION,
    _create_subset_model_v2,
    create_model_v2,
    get_fields,
    is_basemodel_instance,
    is_basemodel_subclass,
    pre_init,
)

def test_pre_init_decorator() -> None:
    from pydantic import BaseModel

    class Foo(BaseModel):
        x: int = 5
        y: int

        @pre_init
        def validator(cls, v: dict[str, Any]) -> dict[str, Any]:
            v["y"] = v["x"] + 1
            return v

    # Type ignore initialization b/c y is marked as required
    foo = Foo()  # type: ignore
    assert foo.y == 6
    foo = Foo(x=10)  # type: ignore
    assert foo.y == 11

def test_pre_init_decorator_with_more_defaults() -> None:
    from pydantic import BaseModel, Field

    class Foo(BaseModel):
        a: int = 1
        b: Optional[int] = None
        c: int = Field(default=2)
        d: int = Field(default_factory=lambda: 3)

        @pre_init
        def validator(cls, v: dict[str, Any]) -> dict[str, Any]:
            assert v["a"] == 1
            assert v["b"] is None
            assert v["c"] == 2
            assert v["d"] == 3
            return v

    # Try to create an instance of Foo
    # nothing is required, but mypy can't track the default for `c`
    Foo()  # type: ignore

def test_with_aliases() -> None:
    from pydantic import BaseModel, Field

    class Foo(BaseModel):
        x: int = Field(default=1, alias="y")
        z: int

        model_config = ConfigDict(
            populate_by_name=True,
        )

        @pre_init
        def validator(cls, v: dict[str, Any]) -> dict[str, Any]:
            v["z"] = v["x"]
            return v

    # Based on defaults
    # z is required
    foo = Foo()  # type: ignore
    assert foo.x == 1
    assert foo.z == 1

    # Based on field name
    # z is required
    foo = Foo(x=2)  # type: ignore
    assert foo.x == 2
    assert foo.z == 2

    # Based on alias
    # z is required
    foo = Foo(y=2)  # type: ignore
    assert foo.x == 2
    assert foo.z == 2

def test_is_basemodel_subclass() -> None:
    """Test pydantic."""
    if PYDANTIC_MAJOR_VERSION == 1:
        from pydantic import BaseModel as BaseModelV1Proper

        assert is_basemodel_subclass(BaseModelV1Proper)
    elif PYDANTIC_MAJOR_VERSION == 2:
        from pydantic import BaseModel as BaseModelV2
        from pydantic.v1 import BaseModel as BaseModelV1

        assert is_basemodel_subclass(BaseModelV2)

        assert is_basemodel_subclass(BaseModelV1)
    else:
        msg = f"Unsupported Pydantic version: {PYDANTIC_MAJOR_VERSION}"
        raise ValueError(msg)

def test_is_basemodel_instance() -> None:
    """Test pydantic."""
    if PYDANTIC_MAJOR_VERSION == 1:
        from pydantic import BaseModel as BaseModelV1Proper

        class FooV1(BaseModelV1Proper):
            x: int

        assert is_basemodel_instance(FooV1(x=5))
    elif PYDANTIC_MAJOR_VERSION == 2:
        from pydantic import BaseModel as BaseModelV2
        from pydantic.v1 import BaseModel as BaseModelV1

        class Foo(BaseModelV2):
            x: int

        assert is_basemodel_instance(Foo(x=5))

        class Bar(BaseModelV1):
            x: int

        assert is_basemodel_instance(Bar(x=5))
    else:
        msg = f"Unsupported Pydantic version: {PYDANTIC_MAJOR_VERSION}"
        raise ValueError(msg)

@pytest.mark.skipif(PYDANTIC_MAJOR_VERSION != 2, reason="Only tests Pydantic v2")
def test_with_field_metadata() -> None:
    """Test pydantic with field metadata"""
    from pydantic import BaseModel as BaseModelV2
    from pydantic import Field as FieldV2

    class Foo(BaseModelV2):
        x: list[int] = FieldV2(
            description="List of integers", min_length=10, max_length=15
        )

    subset_model = _create_subset_model_v2("Foo", Foo, ["x"])
    assert subset_model.model_json_schema() == {
        "properties": {
            "x": {
                "description": "List of integers",
                "items": {"type": "integer"},
                "maxItems": 15,
                "minItems": 10,
                "title": "X",
                "type": "array",
            }
        },
        "required": ["x"],
        "title": "Foo",
        "type": "object",
    }

@pytest.mark.skipif(PYDANTIC_MAJOR_VERSION != 1, reason="Only tests Pydantic v1")
def test_fields_pydantic_v1() -> None:
    from pydantic import BaseModel

    class Foo(BaseModel):
        x: int

    fields = get_fields(Foo)
    assert fields == {"x": Foo.model_fields["x"]}  # type: ignore[index]

@pytest.mark.skipif(PYDANTIC_MAJOR_VERSION != 2, reason="Only tests Pydantic v2")
def test_fields_pydantic_v2_proper() -> None:
    from pydantic import BaseModel

    class Foo(BaseModel):
        x: int

    fields = get_fields(Foo)
    assert fields == {"x": Foo.model_fields["x"]}

@pytest.mark.skipif(PYDANTIC_MAJOR_VERSION != 2, reason="Only tests Pydantic v2")
def test_fields_pydantic_v1_from_2() -> None:
    from pydantic.v1 import BaseModel

    class Foo(BaseModel):
        x: int

    fields = get_fields(Foo)
    assert fields == {"x": Foo.__fields__["x"]}

def test_create_model_v2() -> None:
    """Test that create model v2 works as expected."""

    with pytest.warns(None) as record:  # type: ignore
        foo = create_model_v2("Foo", field_definitions={"a": (int, None)})
        foo.model_json_schema()

    assert list(record) == []

    # schema is used by pydantic, but OK to re-use
    with pytest.warns(None) as record:  # type: ignore
        foo = create_model_v2("Foo", field_definitions={"schema": (int, None)})
        foo.model_json_schema()

    assert list(record) == []

    # From protected namespaces, but definitely OK to use.
    with pytest.warns(None) as record:  # type: ignore
        foo = create_model_v2("Foo", field_definitions={"model_id": (int, None)})
        foo.model_json_schema()

    assert list(record) == []

    with pytest.warns(None) as record:  # type: ignore
        # Verify that we can use non-English characters
        field_name = "もしもし"
        foo = create_model_v2("Foo", field_definitions={field_name: (int, None)})
        foo.model_json_schema()

    assert list(record) == []


# Add a test for schema of runnable assign
    def foo(x: int) -> int:
        return x

    foo_ = RunnableLambda(foo)

    assert foo_.assign(bar=lambda x: "foo").get_output_schema().model_json_schema() == {
        "properties": {"bar": {"title": "Bar"}, "root": {"title": "Root"}},
        "required": ["root", "bar"],
        "title": "RunnableAssignOutput",
        "type": "object",
    }

def test_passthrough_assign_schema() -> None:
    retriever = FakeRetriever()  # str -> List[Document]
    prompt = PromptTemplate.from_template("{context} {question}")
    fake_llm = FakeListLLM(responses=["a"])  # str -> List[List[str]]

    seq_w_assign: Runnable = (
        RunnablePassthrough.assign(context=itemgetter("question") | retriever)
        | prompt
        | fake_llm
    )

    assert seq_w_assign.get_input_jsonschema() == {
        "properties": {"question": {"title": "Question", "type": "string"}},
        "title": "RunnableSequenceInput",
        "type": "object",
        "required": ["question"],
    }
    assert seq_w_assign.get_output_jsonschema() == {
        "title": "FakeListLLMOutput",
        "type": "string",
    }

    invalid_seq_w_assign: Runnable = (
        RunnablePassthrough.assign(context=itemgetter("question") | retriever)
        | fake_llm
    )

    # fallback to RunnableAssign.input_schema if next runnable doesn't have
    # expected dict input_schema
    assert invalid_seq_w_assign.get_input_jsonschema() == {
        "properties": {"question": {"title": "Question"}},
        "title": "RunnableParallel<context>Input",
        "type": "object",
        "required": ["question"],
    }

@pytest.mark.skipif(
    sys.version_info < (3, 9), reason="Requires python version >= 3.9 to run."
)
def test_lambda_schemas(snapshot: SnapshotAssertion) -> None:
    first_lambda = lambda x: x["hello"]  # noqa: E731
    assert RunnableLambda(first_lambda).get_input_jsonschema() == {
        "title": "RunnableLambdaInput",
        "type": "object",
        "properties": {"hello": {"title": "Hello"}},
        "required": ["hello"],
    }

    second_lambda = lambda x, y: (x["hello"], x["bye"], y["bah"])  # noqa: E731
    assert RunnableLambda(second_lambda).get_input_jsonschema() == {  # type: ignore[arg-type]
        "title": "RunnableLambdaInput",
        "type": "object",
        "properties": {"hello": {"title": "Hello"}, "bye": {"title": "Bye"}},
        "required": ["bye", "hello"],
    }

    def get_value(input):  # type: ignore[no-untyped-def]
        return input["variable_name"]

    assert RunnableLambda(get_value).get_input_jsonschema() == {
        "title": "get_value_input",
        "type": "object",
        "properties": {"variable_name": {"title": "Variable Name"}},
        "required": ["variable_name"],
    }

    async def aget_value(input):  # type: ignore[no-untyped-def]
        return (input["variable_name"], input.get("another"))

    assert RunnableLambda(aget_value).get_input_jsonschema() == {
        "title": "aget_value_input",
        "type": "object",
        "properties": {
            "another": {"title": "Another"},
            "variable_name": {"title": "Variable Name"},
        },
        "required": ["another", "variable_name"],
    }

    async def aget_values(input):  # type: ignore[no-untyped-def]
        return {
            "hello": input["variable_name"],
            "bye": input["variable_name"],
            "byebye": input["yo"],
        }

    assert RunnableLambda(aget_values).get_input_jsonschema() == {
        "title": "aget_values_input",
        "type": "object",
        "properties": {
            "variable_name": {"title": "Variable Name"},
            "yo": {"title": "Yo"},
        },
        "required": ["variable_name", "yo"],
    }

    class InputType(TypedDict):
        variable_name: str
        yo: int

    class OutputType(TypedDict):
        hello: str
        bye: str
        byebye: int

    async def aget_values_typed(input: InputType) -> OutputType:
        return {
            "hello": input["variable_name"],
            "bye": input["variable_name"],
            "byebye": input["yo"],
        }

    assert (
        _normalize_schema(
            RunnableLambda(
                aget_values_typed  # type: ignore[arg-type]
            ).get_input_jsonschema()
        )
        == _normalize_schema(
            {
                "$defs": {
                    "InputType": {
                        "properties": {
                            "variable_name": {
                                "title": "Variable " "Name",
                                "type": "string",
                            },
                            "yo": {"title": "Yo", "type": "integer"},
                        },
                        "required": ["variable_name", "yo"],
                        "title": "InputType",
                        "type": "object",
                    }
                },
                "allOf": [{"$ref": "#/$defs/InputType"}],
                "title": "aget_values_typed_input",
            }
        )
    )

    if PYDANTIC_VERSION >= (2, 9):
        assert _normalize_schema(
            RunnableLambda(aget_values_typed).get_output_jsonschema()  # type: ignore
        ) == snapshot(name="schema8")

def test_with_types_with_type_generics() -> None:
    """Verify that with_types works if we use things like List[int]"""

    def foo(x: int) -> None:
        """Add one to the input."""
        raise NotImplementedError

    # Try specifying some
    RunnableLambda(foo).with_types(
        output_type=list[int],  # type: ignore[arg-type]
        input_type=list[int],  # type: ignore[arg-type]
    )
    RunnableLambda(foo).with_types(
        output_type=Sequence[int],  # type: ignore[arg-type]
        input_type=Sequence[int],  # type: ignore[arg-type]
    )

def test_schema_with_itemgetter() -> None:
    """Test runnable with itemgetter."""
    foo = RunnableLambda(itemgetter("hello"))
    assert _schema(foo.input_schema) == {
        "properties": {"hello": {"title": "Hello"}},
        "required": ["hello"],
        "title": "RunnableLambdaInput",
        "type": "object",
    }
    prompt = ChatPromptTemplate.from_template("what is {language}?")
    chain: Runnable = {"language": itemgetter("language")} | prompt
    assert _schema(chain.input_schema) == {
        "properties": {"language": {"title": "Language"}},
        "required": ["language"],
        "title": "RunnableParallel<language>Input",
        "type": "object",
    }

def test_schema_complex_seq() -> None:
    prompt1 = ChatPromptTemplate.from_template("what is the city {person} is from?")
    prompt2 = ChatPromptTemplate.from_template(
        "what country is the city {city} in? respond in {language}"
    )

    model = FakeListChatModel(responses=[""])

    chain1: Runnable = RunnableSequence(
        prompt1, model, StrOutputParser(), name="city_chain"
    )

    assert chain1.name == "city_chain"

    chain2: Runnable = (
        {"city": chain1, "language": itemgetter("language")}
        | prompt2
        | model
        | StrOutputParser()
    )

    assert chain2.get_input_jsonschema() == {
        "title": "RunnableParallel<city,language>Input",
        "type": "object",
        "properties": {
            "person": {"title": "Person", "type": "string"},
            "language": {"title": "Language"},
        },
        "required": ["person", "language"],
    }

    assert chain2.get_output_jsonschema() == {
        "title": "StrOutputParserOutput",
        "type": "string",
    }

    assert chain2.with_types(input_type=str).get_input_jsonschema() == {
        "title": "RunnableSequenceInput",
        "type": "string",
    }

    assert chain2.with_types(input_type=int).get_output_jsonschema() == {
        "title": "StrOutputParserOutput",
        "type": "string",
    }

    class InputType(BaseModel):
        person: str

    assert chain2.with_types(input_type=InputType).get_input_jsonschema() == {
        "title": "InputType",
        "type": "object",
        "properties": {"person": {"title": "Person", "type": "string"}},
        "required": ["person"],
    }


from typing import Any, Callable, NamedTuple, Union

import pytest

from langchain_core.beta.runnables.context import Context
from langchain_core.language_models import FakeListLLM, FakeStreamingListLLM
from langchain_core.output_parsers.string import StrOutputParser
from langchain_core.prompt_values import StringPromptValue
from langchain_core.prompts.prompt import PromptTemplate
from langchain_core.runnables.base import Runnable, RunnableLambda
from langchain_core.runnables.passthrough import RunnablePassthrough
from langchain_core.runnables.utils import aadd, add

class _TestCase(NamedTuple):
    input: Any
    output: Any

def seq_naive_rag() -> Runnable:
    context = [
        "Hi there!",
        "How are you?",
        "What's your name?",
    ]

    retriever = RunnableLambda(lambda x: context)
    prompt = PromptTemplate.from_template("{context} {question}")
    llm = FakeListLLM(responses=["hello"])

    return (
        Context.setter("input")
        | {
            "context": retriever | Context.setter("context"),
            "question": RunnablePassthrough(),
        }
        | prompt
        | llm
        | StrOutputParser()
        | {
            "result": RunnablePassthrough(),
            "context": Context.getter("context"),
            "input": Context.getter("input"),
        }
    )

def seq_naive_rag_alt() -> Runnable:
    context = [
        "Hi there!",
        "How are you?",
        "What's your name?",
    ]

    retriever = RunnableLambda(lambda x: context)
    prompt = PromptTemplate.from_template("{context} {question}")
    llm = FakeListLLM(responses=["hello"])

    return (
        Context.setter("input")
        | {
            "context": retriever | Context.setter("context"),
            "question": RunnablePassthrough(),
        }
        | prompt
        | llm
        | StrOutputParser()
        | Context.setter("result")
        | Context.getter(["context", "input", "result"])
    )

def seq_naive_rag_scoped() -> Runnable:
    context = [
        "Hi there!",
        "How are you?",
        "What's your name?",
    ]

    retriever = RunnableLambda(lambda x: context)
    prompt = PromptTemplate.from_template("{context} {question}")
    llm = FakeListLLM(responses=["hello"])

    scoped = Context.create_scope("a_scope")

    return (
        Context.setter("input")
        | {
            "context": retriever | Context.setter("context"),
            "question": RunnablePassthrough(),
            "scoped": scoped.setter("context") | scoped.getter("context"),
        }
        | prompt
        | llm
        | StrOutputParser()
        | Context.setter("result")
        | Context.getter(["context", "input", "result"])
    )

test_cases = [
    (
        Context.setter("foo") | Context.getter("foo"),
        (
            _TestCase("foo", "foo"),
            _TestCase("bar", "bar"),
        ),
    ),
    (
        Context.setter("input") | {"bar": Context.getter("input")},
        (
            _TestCase("foo", {"bar": "foo"}),
            _TestCase("bar", {"bar": "bar"}),
        ),
    ),
    (
        {"bar": Context.setter("input")} | Context.getter("input"),
        (
            _TestCase("foo", "foo"),
            _TestCase("bar", "bar"),
        ),
    ),
    (
        (
            PromptTemplate.from_template("{foo} {bar}")
            | Context.setter("prompt")
            | FakeListLLM(responses=["hello"])
            | StrOutputParser()
            | {
                "response": RunnablePassthrough(),
                "prompt": Context.getter("prompt"),
            }
        ),
        (
            _TestCase(
                {"foo": "foo", "bar": "bar"},
                {"response": "hello", "prompt": StringPromptValue(text="foo bar")},
            ),
            _TestCase(
                {"foo": "bar", "bar": "foo"},
                {"response": "hello", "prompt": StringPromptValue(text="bar foo")},
            ),
        ),
    ),
    (
        (
            PromptTemplate.from_template("{foo} {bar}")
            | Context.setter("prompt", prompt_str=lambda x: x.to_string())
            | FakeListLLM(responses=["hello"])
            | StrOutputParser()
            | {
                "response": RunnablePassthrough(),
                "prompt": Context.getter("prompt"),
                "prompt_str": Context.getter("prompt_str"),
            }
        ),
        (
            _TestCase(
                {"foo": "foo", "bar": "bar"},
                {
                    "response": "hello",
                    "prompt": StringPromptValue(text="foo bar"),
                    "prompt_str": "foo bar",
                },
            ),
            _TestCase(
                {"foo": "bar", "bar": "foo"},
                {
                    "response": "hello",
                    "prompt": StringPromptValue(text="bar foo"),
                    "prompt_str": "bar foo",
                },
            ),
        ),
    ),
    (
        (
            PromptTemplate.from_template("{foo} {bar}")
            | Context.setter(prompt_str=lambda x: x.to_string())
            | FakeListLLM(responses=["hello"])
            | StrOutputParser()
            | {
                "response": RunnablePassthrough(),
                "prompt_str": Context.getter("prompt_str"),
            }
        ),
        (
            _TestCase(
                {"foo": "foo", "bar": "bar"},
                {"response": "hello", "prompt_str": "foo bar"},
            ),
            _TestCase(
                {"foo": "bar", "bar": "foo"},
                {"response": "hello", "prompt_str": "bar foo"},
            ),
        ),
    ),
    (
        (
            PromptTemplate.from_template("{foo} {bar}")
            | Context.setter("prompt_str", lambda x: x.to_string())
            | FakeListLLM(responses=["hello"])
            | StrOutputParser()
            | {
                "response": RunnablePassthrough(),
                "prompt_str": Context.getter("prompt_str"),
            }
        ),
        (
            _TestCase(
                {"foo": "foo", "bar": "bar"},
                {"response": "hello", "prompt_str": "foo bar"},
            ),
            _TestCase(
                {"foo": "bar", "bar": "foo"},
                {"response": "hello", "prompt_str": "bar foo"},
            ),
        ),
    ),
    (
        (
            PromptTemplate.from_template("{foo} {bar}")
            | Context.setter("prompt")
            | FakeStreamingListLLM(responses=["hello"])
            | StrOutputParser()
            | {
                "response": RunnablePassthrough(),
                "prompt": Context.getter("prompt"),
            }
        ),
        (
            _TestCase(
                {"foo": "foo", "bar": "bar"},
                {"response": "hello", "prompt": StringPromptValue(text="foo bar")},
            ),
            _TestCase(
                {"foo": "bar", "bar": "foo"},
                {"response": "hello", "prompt": StringPromptValue(text="bar foo")},
            ),
        ),
    ),
    (
        seq_naive_rag,
        (
            _TestCase(
                "What up",
                {
                    "result": "hello",
                    "context": [
                        "Hi there!",
                        "How are you?",
                        "What's your name?",
                    ],
                    "input": "What up",
                },
            ),
            _TestCase(
                "Howdy",
                {
                    "result": "hello",
                    "context": [
                        "Hi there!",
                        "How are you?",
                        "What's your name?",
                    ],
                    "input": "Howdy",
                },
            ),
        ),
    ),
    (
        seq_naive_rag_alt,
        (
            _TestCase(
                "What up",
                {
                    "result": "hello",
                    "context": [
                        "Hi there!",
                        "How are you?",
                        "What's your name?",
                    ],
                    "input": "What up",
                },
            ),
            _TestCase(
                "Howdy",
                {
                    "result": "hello",
                    "context": [
                        "Hi there!",
                        "How are you?",
                        "What's your name?",
                    ],
                    "input": "Howdy",
                },
            ),
        ),
    ),
    (
        seq_naive_rag_scoped,
        (
            _TestCase(
                "What up",
                {
                    "result": "hello",
                    "context": [
                        "Hi there!",
                        "How are you?",
                        "What's your name?",
                    ],
                    "input": "What up",
                },
            ),
            _TestCase(
                "Howdy",
                {
                    "result": "hello",
                    "context": [
                        "Hi there!",
                        "How are you?",
                        "What's your name?",
                    ],
                    "input": "Howdy",
                },
            ),
        ),
    ),
]


"langchain_core",
            "language_models",
            "fake",
            "FakeStreamingListLLM"
          ],
          "repr": "FakeStreamingListLLM(responses=['first item, second item, third item'])",
          "name": "FakeStreamingListLLM"
        },
        {
          "lc": 1,
          "type": "constructor",
          "id": [
            "tests",
            "unit_tests",
            "runnables",
            "test_runnable",
            "FakeSplitIntoListParser"
          ],
          "kwargs": {},
          "name": "FakeSplitIntoListParser"
        }
      ],
      "last": {
        "lc": 1,
        "type": "constructor",
        "id": [
          "langchain",
          "schema",
          "runnable",
          "RunnableEach"
        ],
        "kwargs": {
          "bound": {
            "lc": 1,
            "type": "not_implemented",
            "id": [
              "langchain_core",
              "language_models",
              "fake",
              "FakeStreamingListLLM"
            ],
            "repr": "FakeStreamingListLLM(responses=['this', 'is', 'a', 'test'])",
            "name": "FakeStreamingListLLM"
          }
        },
        "name": "RunnableEach<FakeStreamingListLLM>"
      }
    },
    "name": "RunnableSequence"
  }
  '''
# ---
# name: test_higher_order_lambda_runnable
  '''
  {
    "lc": 1,
    "type": "constructor",
    "id": [
      "langchain",
      "schema",
      "runnable",
      "RunnableSequence"
    ],
    "kwargs": {
      "first": {
        "lc": 1,
        "type": "constructor",
        "id": [
          "langchain",
          "schema",
          "runnable",
          "RunnableParallel"
        ],
        "kwargs": {
          "steps__": {
            "key": {
              "lc": 1,
              "type": "not_implemented",
              "id": [
                "langchain_core",
                "runnables",
                "base",
                "RunnableLambda"
              ],
              "repr": "RunnableLambda(lambda x: x['key'])"
            },
            "input": {
              "lc": 1,
              "type": "constructor",
              "id": [
                "langchain",
                "schema",
                "runnable",
                "RunnableParallel"
              ],
              "kwargs": {
                "steps__": {
                  "question": {
                    "lc": 1,
                    "type": "not_implemented",
                    "id": [
                      "langchain_core",
                      "runnables",
                      "base",
                      "RunnableLambda"
                    ],
                    "repr": "RunnableLambda(lambda x: x['question'])"
                  }
                }
              },
              "name": "RunnableParallel<question>"
            }
          }
        },
        "name": "RunnableParallel<key,input>"
      },
      "last": {
        "lc": 1,
        "type": "not_implemented",
        "id": [
          "langchain_core",
          "runnables",
          "base",
          "RunnableLambda"
        ],
        "repr": "RunnableLambda(router)"
      }
    },
    "name": "RunnableSequence"
  }
  '''
# ---
# name: test_lambda_schemas[schema8]
  dict({
    '$defs': dict({
      'OutputType': dict({
        'properties': dict({
          'bye': dict({
            'title': 'Bye',
            'type': 'string',
          }),
          'byebye': dict({
            'title': 'Byebye',
            'type': 'integer',
          }),
          'hello': dict({
            'title': 'Hello',
            'type': 'string',
          }),
        }),
        'required': list([
          'hello',
          'bye',
          'byebye',
        ]),
        'title': 'OutputType',
        'type': 'object',
      }),
    }),
    '$ref': '#/$defs/OutputType',
    'title': 'aget_values_typed_output',
  })
# ---
# name: test_prompt_with_chat_model
  '''
  ChatPromptTemplate(input_variables=['question'], input_types={}, partial_variables={}, messages=[SystemMessagePromptTemplate(prompt=PromptTemplate(input_variables=[], input_types={}, partial_variables={}, template='You are a nice assistant.'), additional_kwargs={}), HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['question'], input_types={}, partial_variables={}, template='{question}'), additional_kwargs={})])
  | FakeListChatModel(responses=['foo'])
  '''
# ---
# name: test_prompt_with_chat_model.1
  '''
  {
    "lc": 1,
    "type": "constructor",
    "id": [
      "langchain",
      "schema",
      "runnable",
      "RunnableSequence"
    ],
    "kwargs": {
      "first": {
        "lc": 1,
        "type": "constructor",
        "id": [
          "langchain",
          "prompts",
          "chat",
          "ChatPromptTemplate"
        ],
        "kwargs": {
          "input_variables": [
            "question"
          ],
          "messages": [
            {
              "lc": 1,
              "type": "constructor",
              "id": [
                "langchain",
                "prompts",
                "chat",
                "SystemMessagePromptTemplate"
              ],
              "kwargs": {
                "prompt": {
                  "lc": 1,
                  "type": "constructor",
                  "id": [
                    "langchain",
                    "prompts",
                    "prompt",
                    "PromptTemplate"
                  ],
                  "kwargs": {
                    "input_variables": [],
                    "template": "You are a nice assistant.",
                    "template_format": "f-string"
                  },
                  "name": "PromptTemplate"
                }
              }
            },
            {
              "lc": 1,
              "type": "constructor",
              "id": [
                "langchain",
                "prompts",
                "chat",
                "HumanMessagePromptTemplate"
              ],
              "kwargs": {
                "prompt": {
                  "lc": 1,
                  "type": "constructor",
                  "id": [
                    "langchain",
                    "prompts",
                    "prompt",
                    "PromptTemplate"
                  ],
                  "kwargs": {
                    "input_variables": [
                      "question"
                    ],
                    "template": "{question}",
                    "template_format": "f-string"
                  },
                  "name": "PromptTemplate"
                }
              }
            }
          ]
        },
        "name": "ChatPromptTemplate"
      },
      "last": {
        "lc": 1,
        "type": "not_implemented",
        "id": [
          "langchain_core",
          "language_models",
          "fake_chat_models",
          "FakeListChatModel"
        ],
        "repr": "FakeListChatModel(responses=['foo'])",
        "name": "FakeListChatModel"
      }
    },
    "name": "RunnableSequence"
  }
  '''
# ---
# name: test_prompt_with_chat_model.2
  list([


"kwargs": {
      "first": {
        "lc": 1,
        "type": "constructor",
        "id": [
          "langchain",
          "prompts",
          "chat",
          "ChatPromptTemplate"
        ],
        "kwargs": {
          "input_variables": [
            "question"
          ],
          "messages": [
            {
              "lc": 1,
              "type": "constructor",
              "id": [
                "langchain",
                "prompts",
                "chat",
                "SystemMessagePromptTemplate"
              ],
              "kwargs": {
                "prompt": {
                  "lc": 1,
                  "type": "constructor",
                  "id": [
                    "langchain",
                    "prompts",
                    "prompt",
                    "PromptTemplate"
                  ],
                  "kwargs": {
                    "input_variables": [],
                    "template": "You are a nice assistant.",
                    "template_format": "f-string"
                  },
                  "name": "PromptTemplate"
                }
              }
            },
            {
              "lc": 1,
              "type": "constructor",
              "id": [
                "langchain",
                "prompts",
                "chat",
                "HumanMessagePromptTemplate"
              ],
              "kwargs": {
                "prompt": {
                  "lc": 1,
                  "type": "constructor",
                  "id": [
                    "langchain",
                    "prompts",
                    "prompt",
                    "PromptTemplate"
                  ],
                  "kwargs": {
                    "input_variables": [
                      "question"
                    ],
                    "template": "{question}",
                    "template_format": "f-string"
                  },
                  "name": "PromptTemplate"
                }
              }
            }
          ]
        },
        "name": "ChatPromptTemplate"
      },
      "last": {
        "lc": 1,
        "type": "not_implemented",
        "id": [
          "langchain_core",
          "language_models",
          "fake",
          "FakeListLLM"
        ],
        "repr": "FakeListLLM(responses=['foo', 'bar'])",
        "name": "FakeListLLM"
      }
    },
    "name": "RunnableSequence"
  }
  '''
# ---
# name: test_prompt_with_llm.1
  list([
    RunTree(id=UUID('00000000-0000-4000-8000-000000000000'), name='RunnableSequence', start_time=FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc), run_type='chain', end_time=FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc), extra={}, error=None, serialized=None, events=[{'name': 'start', 'time': FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc)}, {'name': 'end', 'time': FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc)}], inputs={'question': 'What is your name?'}, outputs={'output': 'foo'}, reference_example_id=None, parent_run_id=None, tags=[], attachments={}, child_runs=[RunTree(id=UUID('00000000-0000-4000-8000-000000000001'), name='ChatPromptTemplate', start_time=FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc), run_type='prompt', end_time=FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc), extra={}, error=None, serialized={'lc': 1, 'type': 'constructor', 'id': ['langchain', 'prompts', 'chat', 'ChatPromptTemplate'], 'kwargs': {'input_variables': ['question'], 'messages': [{'lc': 1, 'type': 'constructor', 'id': ['langchain', 'prompts', 'chat', 'SystemMessagePromptTemplate'], 'kwargs': {'prompt': {'lc': 1, 'type': 'constructor', 'id': ['langchain', 'prompts', 'prompt', 'PromptTemplate'], 'kwargs': {'input_variables': [], 'template': 'You are a nice assistant.', 'template_format': 'f-string'}, 'name': 'PromptTemplate'}}}, {'lc': 1, 'type': 'constructor', 'id': ['langchain', 'prompts', 'chat', 'HumanMessagePromptTemplate'], 'kwargs': {'prompt': {'lc': 1, 'type': 'constructor', 'id': ['langchain', 'prompts', 'prompt', 'PromptTemplate'], 'kwargs': {'input_variables': ['question'], 'template': '{question}', 'template_format': 'f-string'}, 'name': 'PromptTemplate'}}}]}, 'name': 'ChatPromptTemplate'}, events=[{'name': 'start', 'time': FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc)}, {'name': 'end', 'time': FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc)}], inputs={'question': 'What is your name?'}, outputs={'output': ChatPromptValue(messages=[SystemMessage(content='You are a nice assistant.', additional_kwargs={}, response_metadata={}), HumanMessage(content='What is your name?', additional_kwargs={}, response_metadata={})])}, reference_example_id=None, parent_run_id=UUID('00000000-0000-4000-8000-000000000000'), tags=['seq:step:1'], attachments={}, child_runs=[], session_name='default', session_id=None, dotted_order='20230101T000000000000Z00000000-0000-4000-8000-000000000000.20230101T000000000000Z00000000-0000-4000-8000-000000000001', trace_id=UUID('00000000-0000-4000-8000-000000000000')), RunTree(id=UUID('00000000-0000-4000-8000-000000000002'), name='FakeListLLM', start_time=FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc), run_type='llm', end_time=FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc), extra={'invocation_params': {'responses': ['foo', 'bar'], '_type': 'fake-list', 'stop': None}, 'options': {'stop': None}, 'batch_size': 1, 'metadata': {'ls_provider': 'fakelist', 'ls_model_type': 'llm'}}, error=None, serialized={'lc': 1, 'type': 'not_implemented', 'id': ['langchain_core', 'language_models', 'fake', 'FakeListLLM'], 'repr': "FakeListLLM(responses=['foo', 'bar'])", 'name': 'FakeListLLM'}, events=[{'name': 'start', 'time': FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc)}, {'name': 'end', 'time': FakeDatetime(2023, 1, 1, 0, 0, tzinfo=datetime.timezone.utc)}], inputs={'prompts': ['System: You are a nice assistant.\nHuman: What is your name?']}, outputs={'generations': [[{'text': 'foo', 'generation_info': None, 'type': 'Generation'}]], 'llm_output': None, 'run': None, 'type': 'LLMResult'}, reference_example_id=None, parent_run_id=UUID('00000000-0000-4000-8000-000000000000'), tags=['seq:step:2'], attachments={}, child_runs=[], session_name='default', session_id=None, dotted_order='20230101T000000000000Z00000000-0000-4000-8000-000000000000.20230101T000000000000Z00000000-0000-4000-8000-000000000002', trace_id=UUID('00000000-0000-4000-8000-000000000000'))], session_name='default', session_id=None, dotted_order='20230101T000000000000Z00000000-0000-4000-8000-000000000000', trace_id=UUID('00000000-0000-4000-8000-000000000000')),
  ])
# ---
# name: test_prompt_with_llm.2
  list([


"""Set of tests that complement the standard tests for vectorstore.

These tests verify that the base abstraction does appropriate delegation to
the relevant methods.
"""

from __future__ import annotations

import uuid
from collections.abc import Iterable, Sequence
from typing import Any, Optional

import pytest

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings, FakeEmbeddings
from langchain_core.vectorstores import VectorStore

class CustomAddTextsVectorstore(VectorStore):
    """A vectorstore that only implements add texts."""

    def __init__(self) -> None:
        self.store: dict[str, Document] = {}

    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[list[dict]] = None,
        ids: Optional[list[str]] = None,
        **kwargs: Any,
    ) -> list[str]:
        if not isinstance(texts, list):
            texts = list(texts)
        ids_iter = iter(ids or [])

        ids_ = []

        metadatas_ = metadatas or [{} for _ in texts]

        for text, metadata in zip(texts, metadatas_ or []):
            next_id = next(ids_iter, None)
            id_ = next_id or str(uuid.uuid4())
            self.store[id_] = Document(page_content=text, metadata=metadata, id=id_)
            ids_.append(id_)
        return ids_

    def get_by_ids(self, ids: Sequence[str], /) -> list[Document]:
        return [self.store[id] for id in ids if id in self.store]

    @classmethod
    def from_texts(  # type: ignore
        cls,
        texts: list[str],
        embedding: Embeddings,
        metadatas: Optional[list[dict]] = None,
        **kwargs: Any,
    ) -> CustomAddTextsVectorstore:
        vectorstore = CustomAddTextsVectorstore()
        vectorstore.add_texts(texts, metadatas=metadatas, **kwargs)
        return vectorstore

    def similarity_search(
        self, query: str, k: int = 4, **kwargs: Any
    ) -> list[Document]:
        raise NotImplementedError

class CustomAddDocumentsVectorstore(VectorStore):
    """A vectorstore that only implements add documents."""

    def __init__(self) -> None:
        self.store: dict[str, Document] = {}

    def add_documents(
        self,
        documents: list[Document],
        *,
        ids: Optional[list[str]] = None,
        **kwargs: Any,
    ) -> list[str]:
        ids_ = []
        ids_iter = iter(ids or [])
        for document in documents:
            id_ = next(ids_iter) if ids else document.id or str(uuid.uuid4())
            self.store[id_] = Document(
                id=id_, page_content=document.page_content, metadata=document.metadata
            )
            ids_.append(id_)
        return ids_

    def get_by_ids(self, ids: Sequence[str], /) -> list[Document]:
        return [self.store[id] for id in ids if id in self.store]

    @classmethod
    def from_texts(  # type: ignore
        cls,
        texts: list[str],
        embedding: Embeddings,
        metadatas: Optional[list[dict]] = None,
        **kwargs: Any,
    ) -> CustomAddDocumentsVectorstore:
        vectorstore = CustomAddDocumentsVectorstore()
        vectorstore.add_texts(texts, metadatas=metadatas, **kwargs)
        return vectorstore

    def similarity_search(
        self, query: str, k: int = 4, **kwargs: Any
    ) -> list[Document]:
        raise NotImplementedError

@pytest.mark.parametrize(
    "vs_class", [CustomAddTextsVectorstore, CustomAddDocumentsVectorstore]
)
def test_default_add_documents(vs_class: type[VectorStore]) -> None:
    """Test that we can implement the upsert method of the CustomVectorStore
    class without violating the Liskov Substitution Principle.
    """

    store = vs_class()

    # Check upsert with id
    assert store.add_documents([Document(id="1", page_content="hello")]) == ["1"]

    assert store.get_by_ids(["1"]) == [Document(id="1", page_content="hello")]

    # Check upsert without id
    ids = store.add_documents([Document(page_content="world")])
    assert len(ids) == 1
    assert store.get_by_ids(ids) == [Document(id=ids[0], page_content="world")]

    # Check that add_documents works
    assert store.add_documents([Document(id="5", page_content="baz")]) == ["5"]

    # Test add documents with id specified in both document and ids
    original_document = Document(id="7", page_content="baz")
    assert store.add_documents([original_document], ids=["6"]) == ["6"]
    assert original_document.id == "7"  # original document should not be modified
    assert store.get_by_ids(["6"]) == [Document(id="6", page_content="baz")]

@pytest.mark.parametrize(
    "vs_class", [CustomAddTextsVectorstore, CustomAddDocumentsVectorstore]
)
def test_default_add_texts(vs_class: type[VectorStore]) -> None:
    store = vs_class()
    # Check that default implementation of add_texts works
    assert store.add_texts(["hello", "world"], ids=["3", "4"]) == ["3", "4"]

    assert store.get_by_ids(["3", "4"]) == [
        Document(id="3", page_content="hello"),
        Document(id="4", page_content="world"),
    ]

    # Add texts without ids
    ids_ = store.add_texts(["foo", "bar"])
    assert len(ids_) == 2
    assert store.get_by_ids(ids_) == [
        Document(id=ids_[0], page_content="foo"),
        Document(id=ids_[1], page_content="bar"),
    ]

    # Add texts with metadatas
    ids_2 = store.add_texts(["foo", "bar"], metadatas=[{"foo": "bar"}] * 2)
    assert len(ids_2) == 2
    assert store.get_by_ids(ids_2) == [
        Document(id=ids_2[0], page_content="foo", metadata={"foo": "bar"}),
        Document(id=ids_2[1], page_content="bar", metadata={"foo": "bar"}),
    ]

@pytest.mark.parametrize(
    "vs_class", [CustomAddTextsVectorstore, CustomAddDocumentsVectorstore]
)
async def test_default_aadd_documents(vs_class: type[VectorStore]) -> None:
    """Test delegation to the synchronous method."""
    store = vs_class()

    # Check upsert with id
    assert await store.aadd_documents([Document(id="1", page_content="hello")]) == ["1"]

    assert await store.aget_by_ids(["1"]) == [Document(id="1", page_content="hello")]

    # Check upsert without id
    ids = await store.aadd_documents([Document(page_content="world")])
    assert len(ids) == 1
    assert await store.aget_by_ids(ids) == [Document(id=ids[0], page_content="world")]

    # Check that add_documents works
    assert await store.aadd_documents([Document(id="5", page_content="baz")]) == ["5"]

    # Test add documents with id specified in both document and ids
    original_document = Document(id="7", page_content="baz")
    assert await store.aadd_documents([original_document], ids=["6"]) == ["6"]
    assert original_document.id == "7"  # original document should not be modified
    assert await store.aget_by_ids(["6"]) == [Document(id="6", page_content="baz")]


from langchain_core.output_parsers import __all__

EXPECTED_ALL = [
    "BaseLLMOutputParser",
    "BaseGenerationOutputParser",
    "BaseOutputParser",
    "ListOutputParser",
    "CommaSeparatedListOutputParser",
    "NumberedListOutputParser",
    "MarkdownListOutputParser",
    "StrOutputParser",
    "BaseTransformOutputParser",
    "BaseCumulativeTransformOutputParser",
    "SimpleJsonOutputParser",
    "XMLOutputParser",
    "JsonOutputParser",
    "PydanticOutputParser",
    "JsonOutputToolsParser",
    "JsonOutputKeyToolsParser",
    "PydanticToolsParser",
]

def test_all_imports() -> None:
    assert set(__all__) == set(EXPECTED_ALL)


"""Test PydanticOutputParser"""

from enum import Enum
from typing import Literal, Optional

import pydantic
import pytest
from pydantic import BaseModel, Field
from pydantic.v1 import BaseModel as V1BaseModel

from langchain_core.exceptions import OutputParserException
from langchain_core.language_models import ParrotFakeChatModel
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.output_parsers.json import JsonOutputParser
from langchain_core.prompts.prompt import PromptTemplate
from langchain_core.utils.pydantic import TBaseModel

class ForecastV2(pydantic.BaseModel):
    temperature: int
    f_or_c: Literal["F", "C"]
    forecast: str

class ForecastV1(V1BaseModel):
    temperature: int
    f_or_c: Literal["F", "C"]
    forecast: str

@pytest.mark.parametrize("pydantic_object", [ForecastV2, ForecastV1])
def test_pydantic_parser_chaining(
    pydantic_object: TBaseModel,
) -> None:
    prompt = PromptTemplate(
        template="""{{
        "temperature": 20,
        "f_or_c": "C",
        "forecast": "Sunny"
    }}""",
        input_variables=[],
    )

    model = ParrotFakeChatModel()

    parser = PydanticOutputParser(pydantic_object=pydantic_object)  # type: ignore
    chain = prompt | model | parser

    res = chain.invoke({})
    assert type(res) is pydantic_object
    assert res.f_or_c == "C"
    assert res.temperature == 20
    assert res.forecast == "Sunny"

@pytest.mark.parametrize("pydantic_object", [ForecastV2, ForecastV1])
def test_pydantic_parser_validation(pydantic_object: TBaseModel) -> None:
    bad_prompt = PromptTemplate(
        template="""{{
        "temperature": "oof",
        "f_or_c": 1,
        "forecast": "Sunny"
    }}""",
        input_variables=[],
    )

    model = ParrotFakeChatModel()

    parser = PydanticOutputParser(pydantic_object=pydantic_object)  # type: ignore
    chain = bad_prompt | model | parser
    with pytest.raises(OutputParserException):
        chain.invoke({})

# JSON output parser tests
@pytest.mark.parametrize("pydantic_object", [ForecastV2, ForecastV1])
def test_json_parser_chaining(
    pydantic_object: TBaseModel,
) -> None:
    prompt = PromptTemplate(
        template="""{{
        "temperature": 20,
        "f_or_c": "C",
        "forecast": "Sunny"
    }}""",
        input_variables=[],
    )

    model = ParrotFakeChatModel()

    parser = JsonOutputParser(pydantic_object=pydantic_object)  # type: ignore
    chain = prompt | model | parser

    res = chain.invoke({})
    assert res["f_or_c"] == "C"
    assert res["temperature"] == 20
    assert res["forecast"] == "Sunny"

class Actions(Enum):
    SEARCH = "Search"
    CREATE = "Create"
    UPDATE = "Update"
    DELETE = "Delete"

class TestModel(BaseModel):
    action: Actions = Field(description="Action to be performed")
    action_input: str = Field(description="Input to be used in the action")
    additional_fields: Optional[str] = Field(
        description="Additional fields", default=None
    )
    for_new_lines: str = Field(description="To be used to test newlines")

# Prevent pytest from trying to run tests on TestModel
TestModel.__test__ = False  # type: ignore[attr-defined]

DEF_RESULT = """{
    "action": "Update",
    "action_input": "The PydanticOutputParser class is powerful",
    "additional_fields": null,
    "for_new_lines": "not_escape_newline:\n escape_newline: \\n"
}"""

# action 'update' with a lowercase 'u' to test schema validation failure.
DEF_RESULT_FAIL = """{
    "action": "update",
    "action_input": "The PydanticOutputParser class is powerful",
    "additional_fields": null
}"""

DEF_EXPECTED_RESULT = TestModel(
    action=Actions.UPDATE,
    action_input="The PydanticOutputParser class is powerful",
    additional_fields=None,
    for_new_lines="not_escape_newline:\n escape_newline: \n",
)

def test_pydantic_output_parser() -> None:
    """Test PydanticOutputParser."""

    pydantic_parser: PydanticOutputParser = PydanticOutputParser(
        pydantic_object=TestModel
    )

    result = pydantic_parser.parse(DEF_RESULT)
    print("parse_result:", result)  # noqa: T201
    assert result == DEF_EXPECTED_RESULT
    assert pydantic_parser.OutputType is TestModel

def test_pydantic_output_parser_fail() -> None:
    """Test PydanticOutputParser where completion result fails schema validation."""

    pydantic_parser: PydanticOutputParser = PydanticOutputParser(
        pydantic_object=TestModel
    )

    with pytest.raises(OutputParserException) as e:
        pydantic_parser.parse(DEF_RESULT_FAIL)
        assert "Failed to parse TestModel from completion" in str(e)

def test_pydantic_output_parser_type_inference() -> None:
    """Test pydantic output parser type inference."""

    class SampleModel(BaseModel):
        foo: int
        bar: str

    # Ignoring mypy error that appears in python 3.8, but not 3.11.
    # This seems to be functionally correct, so we'll ignore the error.
    pydantic_parser = PydanticOutputParser(pydantic_object=SampleModel)  # type: ignore
    schema = pydantic_parser.get_output_schema().model_json_schema()

    assert schema == {
        "properties": {
            "bar": {"title": "Bar", "type": "string"},
            "foo": {"title": "Foo", "type": "integer"},
        },
        "required": ["foo", "bar"],
        "title": "SampleModel",
        "type": "object",
    }

def test_format_instructions_preserves_language() -> None:
    """Test format instructions does not attempt to encode into ascii."""
    from pydantic import BaseModel, Field

    description = (
        "你好, こんにちは, नमस्ते, Bonjour, Hola, "
        "Olá, 안녕하세요, Jambo, Merhaba, Γειά σου"
    )

    class Foo(BaseModel):
        hello: str = Field(
            description=(
                "你好, こんにちは, नमस्ते, Bonjour, Hola, "
                "Olá, 안녕하세요, Jambo, Merhaba, Γειά σου"
            )
        )

    parser = PydanticOutputParser(pydantic_object=Foo)  # type: ignore
    assert description in parser.get_format_instructions()


import json
from typing import Any

import pytest
from pydantic import BaseModel

from langchain_core.exceptions import OutputParserException
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langchain_core.output_parsers.openai_functions import (
    JsonOutputFunctionsParser,
    PydanticOutputFunctionsParser,
)
from langchain_core.outputs import ChatGeneration

def test_json_output_function_parser() -> None:
    """Test the JSON output function parser is configured with robust defaults."""
    message = AIMessage(
        content="This is a test message",
        additional_kwargs={
            "function_call": {
                "name": "function_name",
                "arguments": '{"arg1": "code\ncode"}',
            }
        },
    )
    chat_generation = ChatGeneration(message=message)

    # Full output
    # Test that the parsers defaults are configured to parse in non-strict mode
    parser = JsonOutputFunctionsParser(args_only=False)
    result = parser.parse_result([chat_generation])
    assert result == {"arguments": {"arg1": "code\ncode"}, "name": "function_name"}

    # Args only
    parser = JsonOutputFunctionsParser(args_only=True)
    result = parser.parse_result([chat_generation])
    assert result == {"arg1": "code\ncode"}

    # Verify that the original message is not modified
    assert message.additional_kwargs == {
        "function_call": {
            "name": "function_name",
            "arguments": '{"arg1": "code\ncode"}',
        }
    }

@pytest.mark.parametrize(
    "config",
    [
        {
            "args_only": False,
            "strict": False,
            "args": '{"arg1": "value1"}',
            "result": {"arguments": {"arg1": "value1"}, "name": "function_name"},
            "exception": None,
        },
        {
            "args_only": True,
            "strict": False,
            "args": '{"arg1": "value1"}',
            "result": {"arg1": "value1"},
            "exception": None,
        },
        {
            "args_only": True,
            "strict": False,
            "args": '{"code": "print(2+\n2)"}',
            "result": {"code": "print(2+\n2)"},
            "exception": None,
        },
        {
            "args_only": True,
            "strict": False,
            "args": '{"code": "你好)"}',
            "result": {"code": "你好)"},
            "exception": None,
        },
        {
            "args_only": True,
            "strict": True,
            "args": '{"code": "print(2+\n2)"}',
            "exception": OutputParserException,
        },
    ],
)
def test_json_output_function_parser_strictness(config: dict[str, Any]) -> None:
    """Test parsing with JSON strictness on and off."""
    args = config["args"]

    message = AIMessage(
        content="This is a test message",
        additional_kwargs={
            "function_call": {"name": "function_name", "arguments": args}
        },
    )
    chat_generation = ChatGeneration(message=message)

    # Full output
    parser = JsonOutputFunctionsParser(
        strict=config["strict"], args_only=config["args_only"]
    )
    if config["exception"] is not None:
        with pytest.raises(config["exception"]):
            parser.parse_result([chat_generation])
    else:
        assert parser.parse_result([chat_generation]) == config["result"]

@pytest.mark.parametrize(
    "bad_message",
    [
        # Human message has no function call
        HumanMessage(content="This is a test message"),
        # AIMessage has no function call information.
        AIMessage(content="This is a test message", additional_kwargs={}),
        # Bad function call information (arguments should be a string)
        AIMessage(
            content="This is a test message",
            additional_kwargs={
                "function_call": {"name": "function_name", "arguments": {}}
            },
        ),
        # Bad function call information (arguments should be proper json)
        AIMessage(
            content="This is a test message",
            additional_kwargs={
                "function_call": {"name": "function_name", "arguments": "noqweqwe"}
            },
        ),
    ],
)
def test_exceptions_raised_while_parsing(bad_message: BaseMessage) -> None:
    """Test exceptions raised correctly while using JSON parser."""
    chat_generation = ChatGeneration(message=bad_message)

    with pytest.raises(OutputParserException):
        JsonOutputFunctionsParser().parse_result([chat_generation])

def test_pydantic_output_functions_parser() -> None:
    """Test pydantic output functions parser."""
    message = AIMessage(
        content="This is a test message",
        additional_kwargs={
            "function_call": {
                "name": "function_name",
                "arguments": json.dumps({"name": "value", "age": 10}),
            }
        },
    )
    chat_generation = ChatGeneration(message=message)

    class Model(BaseModel):
        """Test model."""

        name: str
        age: int

    # Full output
    parser = PydanticOutputFunctionsParser(pydantic_schema=Model)
    result = parser.parse_result([chat_generation])
    assert result == Model(name="value", age=10)

def test_pydantic_output_functions_parser_multiple_schemas() -> None:
    """Test that the parser works if providing multiple pydantic schemas."""

    message = AIMessage(
        content="This is a test message",
        additional_kwargs={
            "function_call": {
                "name": "cookie",
                "arguments": json.dumps({"name": "value", "age": 10}),
            }
        },
    )
    chat_generation = ChatGeneration(message=message)

    class Cookie(BaseModel):
        """Test model."""

        name: str
        age: int

    class Dog(BaseModel):
        """Test model."""

        species: str

    # Full output
    parser = PydanticOutputFunctionsParser(
        pydantic_schema={"cookie": Cookie, "dog": Dog}
    )
    result = parser.parse_result([chat_generation])
    assert result == Cookie(name="value", age=10)


"""Test few shot prompt template."""

from collections.abc import Sequence
from typing import Any

import pytest

from langchain_core.example_selectors import BaseExampleSelector
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.prompts import (
    AIMessagePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain_core.prompts.chat import SystemMessagePromptTemplate
from langchain_core.prompts.few_shot import (
    FewShotChatMessagePromptTemplate,
    FewShotPromptTemplate,
)
from langchain_core.prompts.prompt import PromptTemplate

EXAMPLE_PROMPT = PromptTemplate(
    input_variables=["question", "answer"], template="{question}: {answer}"
)

@pytest.fixture()
@pytest.mark.requires("jinja2")
def example_jinja2_prompt() -> tuple[PromptTemplate, list[dict[str, str]]]:
    example_template = "{{ word }}: {{ antonym }}"

    examples = [
        {"word": "happy", "antonym": "sad"},
        {"word": "tall", "antonym": "short"},
    ]

    return (
        PromptTemplate(
            input_variables=["word", "antonym"],
            template=example_template,
            template_format="jinja2",
        ),
        examples,
    )

def test_suffix_only() -> None:
    """Test prompt works with just a suffix."""
    suffix = "This is a {foo} test."
    input_variables = ["foo"]
    prompt = FewShotPromptTemplate(
        input_variables=input_variables,
        suffix=suffix,
        examples=[],
        example_prompt=EXAMPLE_PROMPT,
    )
    output = prompt.format(foo="bar")
    expected_output = "This is a bar test."
    assert output == expected_output

def test_auto_infer_input_variables() -> None:
    """Test prompt works with just a suffix."""
    suffix = "This is a {foo} test."
    prompt = FewShotPromptTemplate(
        suffix=suffix,
        examples=[],
        example_prompt=EXAMPLE_PROMPT,
    )
    assert prompt.input_variables == ["foo"]

def test_prompt_missing_input_variables() -> None:
    """Test error is raised when input variables are not provided."""
    # Test when missing in suffix
    template = "This is a {foo} test."
    with pytest.raises(ValueError):
        FewShotPromptTemplate(
            input_variables=[],
            suffix=template,
            examples=[],
            example_prompt=EXAMPLE_PROMPT,
            validate_template=True,
        )
    assert FewShotPromptTemplate(
        input_variables=[],
        suffix=template,
        examples=[],
        example_prompt=EXAMPLE_PROMPT,
    ).input_variables == ["foo"]

    # Test when missing in prefix
    template = "This is a {foo} test."
    with pytest.raises(ValueError):
        FewShotPromptTemplate(
            input_variables=[],
            suffix="foo",
            examples=[],
            prefix=template,
            example_prompt=EXAMPLE_PROMPT,
            validate_template=True,
        )
    assert FewShotPromptTemplate(
        input_variables=[],
        suffix="foo",
        examples=[],
        prefix=template,
        example_prompt=EXAMPLE_PROMPT,
    ).input_variables == ["foo"]

async def test_few_shot_functionality() -> None:
    """Test that few shot works with examples."""
    prefix = "This is a test about {content}."
    suffix = "Now you try to talk about {new_content}."
    examples = [
        {"question": "foo", "answer": "bar"},
        {"question": "baz", "answer": "foo"},
    ]
    prompt = FewShotPromptTemplate(
        suffix=suffix,
        prefix=prefix,
        input_variables=["content", "new_content"],
        examples=examples,
        example_prompt=EXAMPLE_PROMPT,
        example_separator="\n",
    )
    expected_output = (
        "This is a test about animals.\n"
        "foo: bar\n"
        "baz: foo\n"
        "Now you try to talk about party."
    )
    output = prompt.format(content="animals", new_content="party")
    assert output == expected_output
    output = await prompt.aformat(content="animals", new_content="party")
    assert output == expected_output

def test_partial_init_string() -> None:
    """Test prompt can be initialized with partial variables."""
    prefix = "This is a test about {content}."
    suffix = "Now you try to talk about {new_content}."
    examples = [
        {"question": "foo", "answer": "bar"},
        {"question": "baz", "answer": "foo"},
    ]
    prompt = FewShotPromptTemplate(
        suffix=suffix,
        prefix=prefix,
        input_variables=["new_content"],
        partial_variables={"content": "animals"},
        examples=examples,
        example_prompt=EXAMPLE_PROMPT,
        example_separator="\n",
    )
    output = prompt.format(new_content="party")
    expected_output = (
        "This is a test about animals.\n"
        "foo: bar\n"
        "baz: foo\n"
        "Now you try to talk about party."
    )
    assert output == expected_output

def test_partial_init_func() -> None:
    """Test prompt can be initialized with partial variables."""
    prefix = "This is a test about {content}."
    suffix = "Now you try to talk about {new_content}."
    examples = [
        {"question": "foo", "answer": "bar"},
        {"question": "baz", "answer": "foo"},
    ]
    prompt = FewShotPromptTemplate(
        suffix=suffix,
        prefix=prefix,
        input_variables=["new_content"],
        partial_variables={"content": lambda: "animals"},
        examples=examples,
        example_prompt=EXAMPLE_PROMPT,
        example_separator="\n",
    )
    output = prompt.format(new_content="party")
    expected_output = (
        "This is a test about animals.\n"
        "foo: bar\n"
        "baz: foo\n"
        "Now you try to talk about party."
    )
    assert output == expected_output

def test_partial() -> None:
    """Test prompt can be partialed."""
    prefix = "This is a test about {content}."
    suffix = "Now you try to talk about {new_content}."
    examples = [
        {"question": "foo", "answer": "bar"},
        {"question": "baz", "answer": "foo"},
    ]
    prompt = FewShotPromptTemplate(
        suffix=suffix,
        prefix=prefix,
        input_variables=["content", "new_content"],
        examples=examples,
        example_prompt=EXAMPLE_PROMPT,
        example_separator="\n",
    )
    new_prompt = prompt.partial(content="foo")
    new_output = new_prompt.format(new_content="party")
    expected_output = (
        "This is a test about foo.\n"
        "foo: bar\n"
        "baz: foo\n"
        "Now you try to talk about party."
    )
    assert new_output == expected_output
    output = prompt.format(new_content="party", content="bar")
    expected_output = (
        "This is a test about bar.\n"
        "foo: bar\n"
        "baz: foo\n"
        "Now you try to talk about party."
    )
    assert output == expected_output

@pytest.mark.requires("jinja2")
def test_prompt_jinja2_functionality(
    example_jinja2_prompt: tuple[PromptTemplate, list[dict[str, str]]],
) -> None:
    prefix = "Starting with {{ foo }}"
    suffix = "Ending with {{ bar }}"

    prompt = FewShotPromptTemplate(
        input_variables=["foo", "bar"],
        suffix=suffix,
        prefix=prefix,
        examples=example_jinja2_prompt[1],
        example_prompt=example_jinja2_prompt[0],
        template_format="jinja2",
    )
    output = prompt.format(foo="hello", bar="bye")
    expected_output = (
        "Starting with hello\n\n" "happy: sad\n\n" "tall: short\n\n" "Ending with bye"
    )

    assert output == expected_output


f test_mustache_prompt_from_template(snapshot: SnapshotAssertion) -> None:
    """Test prompts can be constructed from a template."""
    # Single input variable.
    template = "This is a {{foo}} test."
    prompt = PromptTemplate.from_template(template, template_format="mustache")
    assert prompt.format(foo="bar") == "This is a bar test."
    assert prompt.input_variables == ["foo"]
    assert prompt.get_input_jsonschema() == {
        "title": "PromptInput",
        "type": "object",
        "properties": {"foo": {"title": "Foo", "type": "string", "default": None}},
    }

    # Multiple input variables.
    template = "This {{bar}} is a {{foo}} test."
    prompt = PromptTemplate.from_template(template, template_format="mustache")
    assert prompt.format(bar="baz", foo="bar") == "This baz is a bar test."
    assert prompt.input_variables == ["bar", "foo"]
    assert prompt.get_input_jsonschema() == {
        "title": "PromptInput",
        "type": "object",
        "properties": {
            "bar": {"title": "Bar", "type": "string", "default": None},
            "foo": {"title": "Foo", "type": "string", "default": None},
        },
    }

    # Multiple input variables with repeats.
    template = "This {{bar}} is a {{foo}} test {{&foo}}."
    prompt = PromptTemplate.from_template(template, template_format="mustache")
    assert prompt.format(bar="baz", foo="bar") == "This baz is a bar test bar."
    assert prompt.input_variables == ["bar", "foo"]
    assert prompt.get_input_jsonschema() == {
        "title": "PromptInput",
        "type": "object",
        "properties": {
            "bar": {"title": "Bar", "type": "string", "default": None},
            "foo": {"title": "Foo", "type": "string", "default": None},
        },
    }

    # Nested variables.
    template = "This {{obj.bar}} is a {{obj.foo}} test {{{foo}}}."
    prompt = PromptTemplate.from_template(template, template_format="mustache")
    assert prompt.format(obj={"bar": "foo", "foo": "bar"}, foo="baz") == (
        "This foo is a bar test baz."
    )
    assert prompt.input_variables == ["foo", "obj"]
    if PYDANTIC_VERSION >= (2, 9):
        assert _normalize_schema(prompt.get_input_jsonschema()) == snapshot(
            name="schema_0"
        )

    # . variables
    template = "This {{.}} is a test."
    prompt = PromptTemplate.from_template(template, template_format="mustache")
    assert prompt.format(foo="baz") == ("This {'foo': 'baz'} is a test.")
    assert prompt.input_variables == []
    assert prompt.get_input_jsonschema() == {
        "title": "PromptInput",
        "type": "object",
        "properties": {},
    }

    # section/context variables
    template = """This{{#foo}}
        {{bar}}
    {{/foo}}is a test."""
    prompt = PromptTemplate.from_template(template, template_format="mustache")
    assert prompt.format(foo={"bar": "yo"}) == (
        """This
        yo
    is a test."""
    )
    assert prompt.input_variables == ["foo"]
    if PYDANTIC_VERSION >= (2, 9):
        assert _normalize_schema(prompt.get_input_jsonschema()) == snapshot(
            name="schema_2"
        )

    # more complex nested section/context variables
    template = """This{{#foo}}
        {{bar}}
        {{#baz}}
            {{qux}}
        {{/baz}}
        {{quux}}
    {{/foo}}is a test."""
    prompt = PromptTemplate.from_template(template, template_format="mustache")
    assert prompt.format(
        foo={"bar": "yo", "baz": [{"qux": "wassup"}], "quux": "hello"}
    ) == (
        """This
        yo
            wassup
        hello
    is a test."""
    )
    assert prompt.input_variables == ["foo"]
    if PYDANTIC_VERSION >= (2, 9):
        assert _normalize_schema(prompt.get_input_jsonschema()) == snapshot(
            name="schema_3"
        )

    # triply nested section/context variables
    template = """This{{#foo}}
        {{bar}}
        {{#baz.qux}}
            {{#barfoo}}
                {{foobar}}
            {{/barfoo}}
            {{foobar}}
        {{/baz.qux}}
        {{quux}}
    {{/foo}}is a test."""
    prompt = PromptTemplate.from_template(template, template_format="mustache")
    assert prompt.format(
        foo={
            "bar": "yo",
            "baz": {
                "qux": [
                    {"foobar": "wassup"},
                    {"foobar": "yoyo", "barfoo": {"foobar": "hello there"}},
                ]
            },
            "quux": "hello",
        }
    ) == (
        """This
        yo
            wassup
                hello there
            yoyo
        hello
    is a test."""
    )
    assert prompt.input_variables == ["foo"]
    if PYDANTIC_VERSION >= (2, 9):
        assert _normalize_schema(prompt.get_input_jsonschema()) == snapshot(
            name="schema_4"
        )

    # section/context variables with repeats
    template = """This{{#foo}}
        {{bar}}
    {{/foo}}is a test."""
    prompt = PromptTemplate.from_template(template, template_format="mustache")
    assert prompt.format(foo=[{"bar": "yo"}, {"bar": "hello"}]) == (
        """This
        yo

        hello
    is a test."""  # noqa: W293
    )
    assert prompt.input_variables == ["foo"]
    if PYDANTIC_VERSION >= (2, 9):
        assert _normalize_schema(prompt.get_input_jsonschema()) == snapshot(
            name="schema_5"
        )
    template = """This{{^foo}}
        no foos
    {{/foo}}is a test."""
    prompt = PromptTemplate.from_template(template, template_format="mustache")
    assert prompt.format() == (
        """This
        no foos
    is a test."""
    )
    assert prompt.input_variables == ["foo"]
    assert prompt.get_input_jsonschema() == {
        "properties": {"foo": {"default": None, "title": "Foo", "type": "object"}},
        "title": "PromptInput",
        "type": "object",
    }

def test_prompt_from_template_with_partial_variables() -> None:
    """Test prompts can be constructed from a template with partial variables."""
    # given
    template = "This is a {foo} test {bar}."
    partial_variables = {"bar": "baz"}
    # when
    prompt = PromptTemplate.from_template(template, partial_variables=partial_variables)
    # then
    expected_prompt = PromptTemplate(
        template=template,
        input_variables=["foo"],
        partial_variables=partial_variables,
    )
    assert prompt == expected_prompt

def test_prompt_missing_input_variables() -> None:
    """Test error is raised when input variables are not provided."""
    template = "This is a {foo} test."
    input_variables: list = []
    with pytest.raises(ValueError):
        PromptTemplate(
            input_variables=input_variables, template=template, validate_template=True
        )
    assert PromptTemplate(
        input_variables=input_variables, template=template
    ).input_variables == ["foo"]

def test_prompt_empty_input_variable() -> None:
    """Test error is raised when empty string input variable."""
    with pytest.raises(ValueError):
        PromptTemplate(input_variables=[""], template="{}", validate_template=True)

def test_prompt_wrong_input_variables() -> None:
    """Test error is raised when name of input variable is wrong."""
    template = "This is a {foo} test."
    input_variables = ["bar"]
    with pytest.raises(ValueError):
        PromptTemplate(
            input_variables=input_variables, template=template, validate_template=True
        )
    assert PromptTemplate(
        input_variables=input_variables, template=template
    ).input_variables == ["foo"]


"""Test few shot prompt template."""

import pytest

from langchain_core.prompts.few_shot_with_templates import FewShotPromptWithTemplates
from langchain_core.prompts.prompt import PromptTemplate

EXAMPLE_PROMPT = PromptTemplate(
    input_variables=["question", "answer"], template="{question}: {answer}"
)

async def test_prompttemplate_prefix_suffix() -> None:
    """Test that few shot works when prefix and suffix are PromptTemplates."""
    prefix = PromptTemplate(
        input_variables=["content"], template="This is a test about {content}."
    )
    suffix = PromptTemplate(
        input_variables=["new_content"],
        template="Now you try to talk about {new_content}.",
    )

    examples = [
        {"question": "foo", "answer": "bar"},
        {"question": "baz", "answer": "foo"},
    ]
    prompt = FewShotPromptWithTemplates(
        suffix=suffix,
        prefix=prefix,
        input_variables=["content", "new_content"],
        examples=examples,
        example_prompt=EXAMPLE_PROMPT,
        example_separator="\n",
    )
    expected_output = (
        "This is a test about animals.\n"
        "foo: bar\n"
        "baz: foo\n"
        "Now you try to talk about party."
    )
    output = prompt.format(content="animals", new_content="party")
    assert output == expected_output
    output = await prompt.aformat(content="animals", new_content="party")
    assert output == expected_output

def test_prompttemplate_validation() -> None:
    """Test that few shot works when prefix and suffix are PromptTemplates."""
    prefix = PromptTemplate(
        input_variables=["content"], template="This is a test about {content}."
    )
    suffix = PromptTemplate(
        input_variables=["new_content"],
        template="Now you try to talk about {new_content}.",
    )

    examples = [
        {"question": "foo", "answer": "bar"},
        {"question": "baz", "answer": "foo"},
    ]
    with pytest.raises(ValueError):
        FewShotPromptWithTemplates(
            suffix=suffix,
            prefix=prefix,
            input_variables=[],
            examples=examples,
            example_prompt=EXAMPLE_PROMPT,
            example_separator="\n",
            validate_template=True,
        )
    assert FewShotPromptWithTemplates(
        suffix=suffix,
        prefix=prefix,
        input_variables=[],
        examples=examples,
        example_prompt=EXAMPLE_PROMPT,
        example_separator="\n",
    ).input_variables == ["content", "new_content"]


import base64
import tempfile
from pathlib import Path
from typing import Any, Union, cast

import pytest
from pydantic import ValidationError
from syrupy import SnapshotAssertion

from langchain_core._api.deprecation import (
    LangChainPendingDeprecationWarning,
)
from langchain_core.load import dumpd, load
from langchain_core.messages import (
    AIMessage,
    BaseMessage,
    HumanMessage,
    SystemMessage,
    get_buffer_string,
)
from langchain_core.prompt_values import ChatPromptValue
from langchain_core.prompts import PromptTemplate
from langchain_core.prompts.chat import (
    AIMessagePromptTemplate,
    BaseMessagePromptTemplate,
    ChatMessage,
    ChatMessagePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
    _convert_to_message,
)
from tests.unit_tests.pydantic_utils import _normalize_schema

@pytest.fixture
def messages() -> list[BaseMessagePromptTemplate]:
    """Create messages."""
    system_message_prompt = SystemMessagePromptTemplate(
        prompt=PromptTemplate(
            template="Here's some context: {context}",
            input_variables=["context"],
        )
    )
    human_message_prompt = HumanMessagePromptTemplate(
        prompt=PromptTemplate(
            template="Hello {foo}, I'm {bar}. Thanks for the {context}",
            input_variables=["foo", "bar", "context"],
        )
    )
    ai_message_prompt = AIMessagePromptTemplate(
        prompt=PromptTemplate(
            template="I'm an AI. I'm {foo}. I'm {bar}.",
            input_variables=["foo", "bar"],
        )
    )
    chat_message_prompt = ChatMessagePromptTemplate(
        role="test",
        prompt=PromptTemplate(
            template="I'm a generic message. I'm {foo}. I'm {bar}.",
            input_variables=["foo", "bar"],
        ),
    )
    return [
        system_message_prompt,
        human_message_prompt,
        ai_message_prompt,
        chat_message_prompt,
    ]

@pytest.fixture
def chat_prompt_template(
    messages: list[BaseMessagePromptTemplate],
) -> ChatPromptTemplate:
    """Create a chat prompt template."""
    return ChatPromptTemplate(
        input_variables=["foo", "bar", "context"],
        messages=messages,  # type: ignore[arg-type]
    )

def test_create_chat_prompt_template_from_template() -> None:
    """Create a chat prompt template."""
    prompt = ChatPromptTemplate.from_template("hi {foo} {bar}")
    assert prompt.messages == [
        HumanMessagePromptTemplate.from_template("hi {foo} {bar}")
    ]

def test_create_chat_prompt_template_from_template_partial() -> None:
    """Create a chat prompt template with partials."""
    prompt = ChatPromptTemplate.from_template(
        "hi {foo} {bar}", partial_variables={"foo": "jim"}
    )
    expected_prompt = PromptTemplate(
        template="hi {foo} {bar}",
        input_variables=["bar"],
        partial_variables={"foo": "jim"},
    )
    assert len(prompt.messages) == 1
    output_prompt = prompt.messages[0]
    assert isinstance(output_prompt, HumanMessagePromptTemplate)
    assert output_prompt.prompt == expected_prompt

def test_create_system_message_prompt_template_from_template_partial() -> None:
    """Create a system message prompt template with partials."""

    graph_creator_content = """
    Your instructions are:
    {instructions}
    History:
    {history}
    """
    json_prompt_instructions: dict = {}
    graph_analyst_template = SystemMessagePromptTemplate.from_template(
        template=graph_creator_content,
        input_variables=["history"],
        partial_variables={"instructions": json_prompt_instructions},
    )
    assert graph_analyst_template.format(history="history") == SystemMessage(
        content="\n    Your instructions are:\n  "
        "  {}\n    History:\n    "
        "history\n    "
    )

def test_create_system_message_prompt_list_template() -> None:
    graph_creator_content1 = """
    This is the prompt for the first test:
    {variables}
    """
    graph_creator_content2 = """
    This is the prompt for the second test:
        {variables}
        """
    graph_analyst_template = SystemMessagePromptTemplate.from_template(
        template=[graph_creator_content1, graph_creator_content2],
        input_variables=["variables"],
    )
    assert graph_analyst_template.format(variables="foo") == SystemMessage(
        content=[
            {
                "type": "text",
                "text": "\n    This is the prompt for the first test:\n    foo\n    ",
            },
            {
                "type": "text",
                "text": "\n    This is the prompt for "
                "the second test:\n        foo\n        ",
            },
        ]
    )

def test_create_system_message_prompt_list_template_partial_variables_not_null() -> (
    None
):
    graph_creator_content1 = """
    This is the prompt for the first test:
    {variables}
    """
    graph_creator_content2 = """
    This is the prompt for the second test:
        {variables}
        """

    try:
        graph_analyst_template = SystemMessagePromptTemplate.from_template(
            template=[graph_creator_content1, graph_creator_content2],
            input_variables=["variables"],
            partial_variables={"variables": "foo"},
        )
        graph_analyst_template.format(variables="foo")
    except ValueError as e:
        assert str(e) == "Partial variables are not supported for list of templates."

def test_message_prompt_template_from_template_file() -> None:
    expected = ChatMessagePromptTemplate(
        prompt=PromptTemplate(
            template="Question: {question}\nAnswer:", input_variables=["question"]
        ),
        role="human",
    )
    actual = ChatMessagePromptTemplate.from_template_file(
        Path(__file__).parent.parent / "data" / "prompt_file.txt",
        ["question"],
        role="human",
    )
    assert expected == actual

async def test_chat_prompt_template(chat_prompt_template: ChatPromptTemplate) -> None:
    """Test chat prompt template."""
    prompt = chat_prompt_template.format_prompt(foo="foo", bar="bar", context="context")
    assert isinstance(prompt, ChatPromptValue)
    messages = prompt.to_messages()
    assert len(messages) == 4
    assert messages[0].content == "Here's some context: context"
    assert messages[1].content == "Hello foo, I'm bar. Thanks for the context"
    assert messages[2].content == "I'm an AI. I'm foo. I'm bar."
    assert messages[3].content == "I'm a generic message. I'm foo. I'm bar."

    async_prompt = await chat_prompt_template.aformat_prompt(
        foo="foo", bar="bar", context="context"
    )

    assert async_prompt.to_messages() == messages

    string = prompt.to_string()
    expected = (
        "System: Here's some context: context\n"
        "Human: Hello foo, I'm bar. Thanks for the context\n"
        "AI: I'm an AI. I'm foo. I'm bar.\n"
        "test: I'm a generic message. I'm foo. I'm bar."
    )
    assert string == expected

    string = chat_prompt_template.format(foo="foo", bar="bar", context="context")
    assert string == expected

    string = await chat_prompt_template.aformat(foo="foo", bar="bar", context="context")
    assert string == expected

def test_chat_prompt_template_from_messages(
    messages: list[BaseMessagePromptTemplate],
) -> None:
    """Test creating a chat prompt template from messages."""
    chat_prompt_template = ChatPromptTemplate.from_messages(messages)
    assert sorted(chat_prompt_template.input_variables) == sorted(
        ["context", "foo", "bar"]
    )
    assert len(chat_prompt_template.messages) == 4

async def test_chat_prompt_template_from_messages_using_role_strings() -> None:
    """Test creating a chat prompt template from role string messages."""
    template = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a helpful AI bot. Your name is {name}."),
            ("human", "Hello, how are you doing?"),
            ("ai", "I'm doing well, thanks!"),
            ("human", "{user_input}"),
        ]
    )

    expected = [
        SystemMessage(
            content="You are a helpful AI bot. Your name is Bob.", additional_kwargs={}
        ),
        HumanMessage(
            content="Hello, how are you doing?", additional_kwargs={}, example=False
        ),
        AIMessage(
            content="I'm doing well, thanks!", additional_kwargs={}, example=False
        ),
        HumanMessage(content="What is your name?", additional_kwargs={}, example=False),
    ]

    messages = template.format_messages(name="Bob", user_input="What is your name?")
    assert messages == expected

    messages = await template.aformat_messages(
        name="Bob", user_input="What is your name?"
    )
    assert messages == expected


from langchain_core.documents import Document

def test_str() -> None:
    assert str(Document(page_content="Hello, World!")) == "page_content='Hello, World!'"
    assert (
        str(Document(page_content="Hello, World!", metadata={"a": 3}))
        == "page_content='Hello, World!' metadata={'a': 3}"
    )

def test_repr() -> None:
    assert (
        repr(Document(page_content="Hello, World!"))
        == "Document(metadata={}, page_content='Hello, World!')"
    )
    assert (
        repr(Document(page_content="Hello, World!", metadata={"a": 3}))
        == "Document(metadata={'a': 3}, page_content='Hello, World!')"
    )


from langchain_core.documents import Document

def test_init() -> None:
    for doc in [
        Document(page_content="foo"),
        Document(page_content="foo", metadata={"a": 1}),
        Document(page_content="foo", id=None),
        Document(page_content="foo", id="1"),
        Document(page_content="foo", id=1),
    ]:
        assert isinstance(doc, Document)