langchain/libs/cli/langchain_cli/namespaces/integration.py

"""Develop integration packages for LangChain."""

import os
import re
import shutil
import subprocess
from pathlib import Path
from typing import Annotated, cast

import typer
from typing_extensions import TypedDict

from langchain_cli.utils.find_replace import replace_file, replace_glob

integration_cli = typer.Typer(no_args_is_help=True, add_completion=False)


class Replacements(TypedDict):
    """Replacements."""

    __package_name__: str
    __module_name__: str
    __ModuleName__: str
    __MODULE_NAME__: str
    __package_name_short__: str
    __package_name_short_snake__: str


def _process_name(name: str, *, community: bool = False) -> Replacements:
    preprocessed = name.replace("_", "-").lower()

    preprocessed = preprocessed.removeprefix("langchain-")

    if not re.match(r"^[a-z][a-z0-9-]*$", preprocessed):
        msg = (
            "Name should only contain lowercase letters (a-z), numbers, and hyphens"
            ", and start with a letter."
        )
        raise ValueError(msg)
    if preprocessed.endswith("-"):
        msg = "Name should not end with `-`."
        raise ValueError(msg)
    if preprocessed.find("--") != -1:
        msg = "Name should not contain consecutive hyphens."
        raise ValueError(msg)
    replacements: Replacements = {
        "__package_name__": f"langchain-{preprocessed}",
        "__module_name__": "langchain_" + preprocessed.replace("-", "_"),
        "__ModuleName__": preprocessed.title().replace("-", ""),
        "__MODULE_NAME__": preprocessed.upper().replace("-", ""),
        "__package_name_short__": preprocessed,
        "__package_name_short_snake__": preprocessed.replace("-", "_"),
    }
    if community:
        replacements["__module_name__"] = preprocessed.replace("-", "_")
    return replacements


@integration_cli.command()
def new(
    name: Annotated[
        str,
        typer.Option(
            help="The name of the integration to create (e.g. `my-integration`)",
            prompt="The name of the integration to create (e.g. `my-integration`)",
        ),
    ],
    name_class: Annotated[
        str | None,
        typer.Option(
            help="The name of the integration in PascalCase. e.g. `MyIntegration`."
            " This is used to name classes like `MyIntegrationVectorStore`",
        ),
    ] = None,
    src: Annotated[
        list[str] | None,
        typer.Option(
            help="The name of the single template file to copy."
            " e.g. `--src integration_template/chat_models.py "
            "--dst my_integration/chat_models.py`. Can be used multiple times.",
        ),
    ] = None,
    dst: Annotated[
        list[str] | None,
        typer.Option(
            help="The relative path to the integration package to place the new file in"
            ". e.g. `my-integration/my_integration.py`",
        ),
    ] = None,
) -> None:
    """Create a new integration package."""
    try:
        replacements = _process_name(name)
    except ValueError as e:
        typer.echo(e)
        raise typer.Exit(code=1) from None

    if name_class:
        if not re.match(r"^[A-Z][a-zA-Z0-9]*$", name_class):
            typer.echo(
                "Name should only contain letters (a-z, A-Z), numbers, and underscores"
                ", and start with a capital letter.",
            )
            raise typer.Exit(code=1)
        replacements["__ModuleName__"] = name_class
    else:
        replacements["__ModuleName__"] = typer.prompt(
            "Name of integration in PascalCase",
            default=replacements["__ModuleName__"],
        )

    project_template_dir = Path(__file__).parents[1] / "integration_template"
    destination_dir = Path.cwd() / replacements["__package_name__"]
    if not src and not dst:
        if destination_dir.exists():
            typer.echo(f"Folder {destination_dir} exists.")
            raise typer.Exit(code=1)

        # Copy over template from ../integration_template
        shutil.copytree(project_template_dir, destination_dir, dirs_exist_ok=False)

        # Folder movement
        package_dir = destination_dir / replacements["__module_name__"]
        shutil.move(destination_dir / "integration_template", package_dir)

        # Replacements in files
        replace_glob(destination_dir, "**/*", cast("dict[str, str]", replacements))

        # Dependency install
        try:
            # Use --no-progress to avoid tty issues in CI/test environments
            env = os.environ.copy()
            env.pop("UV_FROZEN", None)
            env.pop("VIRTUAL_ENV", None)
            subprocess.run(
                ["uv", "sync", "--dev", "--no-progress"],  # noqa: S607
                cwd=destination_dir,
                check=True,
                env=env,
            )
        except FileNotFoundError:
            typer.echo(
                "uv is not installed. Skipping dependency installation; run "
                "`uv sync --dev` manually if needed.",
            )
        except subprocess.CalledProcessError:
            typer.echo(
                "Failed to install dependencies. You may need to run "
                "`uv sync --dev` manually in the package directory.",
            )
    else:
        # Confirm src and dst are the same length
        if not src:
            typer.echo("Cannot provide --dst without --src.")
            raise typer.Exit(code=1)
        src_paths = [project_template_dir / p for p in src]
        if dst and len(src) != len(dst):
            typer.echo("Number of --src and --dst arguments must match.")
            raise typer.Exit(code=1)
        if not dst:
            # Assume we're in a package dir, copy to equivalent path
            dst_paths = [destination_dir / p for p in src]
        else:
            dst_paths = [Path.cwd() / p for p in dst]
            dst_paths = [
                p / f"{replacements['__package_name_short_snake__']}.ipynb"
                if not p.suffix
                else p
                for p in dst_paths
            ]

        # Confirm no duplicate dst_paths
        if len(dst_paths) != len(set(dst_paths)):
            typer.echo(
                "Duplicate destination paths provided or computed - please "
                "specify them explicitly with --dst.",
            )
            raise typer.Exit(code=1)

        # Confirm no files exist at dst_paths
        for dst_path in dst_paths:
            if dst_path.exists():
                typer.echo(f"File {dst_path} exists.")
                raise typer.Exit(code=1)

        for src_path, dst_path in zip(src_paths, dst_paths, strict=False):
            shutil.copy(src_path, dst_path)
            replace_file(dst_path, cast("dict[str, str]", replacements))


TEMPLATE_MAP: dict[str, str] = {
    "ChatModel": "chat.ipynb",
    "DocumentLoader": "document_loaders.ipynb",
    "Tool": "tools.ipynb",
    "VectorStore": "vectorstores.ipynb",
    "Embeddings": "text_embedding.ipynb",
    "ByteStore": "kv_store.ipynb",
    "LLM": "llms.ipynb",
    "Provider": "provider.ipynb",
    "Toolkit": "toolkits.ipynb",
    "Retriever": "retrievers.ipynb",
}

_component_types_str = ", ".join(f"`{k}`" for k in TEMPLATE_MAP)


@integration_cli.command()
def create_doc(
    name: Annotated[
        str,
        typer.Option(
            help=(
                "The kebab-case name of the integration (e.g. `openai`, "
                "`google-vertexai`). Do not include a 'langchain-' prefix."
            ),
            prompt=(
                "The kebab-case name of the integration (e.g. `openai`, "
                "`google-vertexai`). Do not include a 'langchain-' prefix."
            ),
        ),
    ],
    name_class: Annotated[
        str | None,
        typer.Option(
            help=(
                "The PascalCase name of the integration (e.g. `OpenAI`, "
                "`VertexAI`). Do not include a 'Chat', 'VectorStore', etc. "
                "prefix/suffix."
            ),
        ),
    ] = None,
    component_type: Annotated[
        str,
        typer.Option(
            help=(
                f"The type of component. Currently supported: {_component_types_str}."
            ),
        ),
    ] = "ChatModel",
    destination_dir: Annotated[
        str,
        typer.Option(
            help="The relative path to the docs directory to place the new file in.",
            prompt="The relative path to the docs directory to place the new file in.",
        ),
    ] = "docs/docs/integrations/chat/",
) -> None:
    """Create a new integration doc."""
    if component_type not in TEMPLATE_MAP:
        typer.echo(
            f"Unrecognized {component_type=}. Expected one of {_component_types_str}.",
        )
        raise typer.Exit(code=1)

    new(
        name=name,
        name_class=name_class,
        src=[f"docs/{TEMPLATE_MAP[component_type]}"],
        dst=[destination_dir],
    )