test_loading.py — langchain Source File

Source Code

"""Test loading functionality."""

import os
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path

import pytest

from langchain_core.prompts.few_shot import FewShotPromptTemplate
from langchain_core.prompts.loading import load_prompt
from langchain_core.prompts.prompt import PromptTemplate

EXAMPLE_DIR = (Path(__file__).parent.parent / "examples").absolute()


@contextmanager
def change_directory(dir_path: Path) -> Iterator[None]:
    """Change the working directory to the right folder."""
    origin = Path().absolute()
    try:
        os.chdir(dir_path)
        yield
    finally:
        os.chdir(origin)


def test_loading_from_yaml() -> None:
    """Test loading from yaml file."""
    prompt = load_prompt(EXAMPLE_DIR / "simple_prompt.yaml")
    expected_prompt = PromptTemplate(
        input_variables=["adjective"],
        partial_variables={"content": "dogs"},
        template="Tell me a {adjective} joke about {content}.",
    )
    assert prompt == expected_prompt


def test_loading_from_json() -> None:
    """Test loading from json file."""
    prompt = load_prompt(EXAMPLE_DIR / "simple_prompt.json")
    expected_prompt = PromptTemplate(
        input_variables=["adjective", "content"],
        template="Tell me a {adjective} joke about {content}.",
    )
    assert prompt == expected_prompt


def test_loading_jinja_from_json() -> None:
    """Test that loading jinja2 format prompts from JSON raises ValueError."""
    prompt_path = EXAMPLE_DIR / "jinja_injection_prompt.json"
    with pytest.raises(ValueError, match=r".*can lead to arbitrary code execution.*"):
        load_prompt(prompt_path)


def test_loading_jinja_from_yaml() -> None:
    """Test that loading jinja2 format prompts from YAML raises ValueError."""
    prompt_path = EXAMPLE_DIR / "jinja_injection_prompt.yaml"
    with pytest.raises(ValueError, match=r".*can lead to arbitrary code execution.*"):
        load_prompt(prompt_path)
// ... (120 more lines)