langchain-dev-utils 1.3.7__py3-none-any.whl

langchain_dev_utils/embeddings/base.py

@@ -0,0 +1,234 @@
+from typing import Any, Literal, NotRequired, Optional, TypedDict, Union
+
+from langchain.embeddings.base import _SUPPORTED_PROVIDERS, Embeddings, init_embeddings
+from langchain_core.utils import from_env
+
+from langchain_dev_utils._utils import (
+    _check_pkg_install,
+    _get_base_url_field_name,
+    _validate_provider_name,
+)
+
+_EMBEDDINGS_PROVIDERS_DICT = {}
+
+EmbeddingsType = Union[type[Embeddings], Literal["openai-compatible"]]
+
+
+class EmbeddingProvider(TypedDict):
+    provider_name: str
+    embeddings_model: EmbeddingsType
+    base_url: NotRequired[str]
+
+
+def _parse_model_string(model_name: str) -> tuple[str, str]:
+    """Parse model string into provider and model name.
+
+    Args:
+        model_name: Model name string in format 'provider:model-name'
+
+    Returns:
+        Tuple of (provider, model) parsed from the model_name
+
+    Raises:
+        ValueError: If model name format is invalid or model name is empty
+    """
+    if ":" not in model_name:
+        msg = (
+            f"Invalid model format '{model_name}'.\n"
+            f"Model name must be in format 'provider:model-name'\n"
+        )
+        raise ValueError(msg)
+
+    provider, model = model_name.split(":", 1)
+    provider = provider.lower().strip()
+    model = model.strip()
+    if not model:
+        msg = "Model name cannot be empty"
+        raise ValueError(msg)
+    return provider, model
+
+
+def register_embeddings_provider(
+    provider_name: str,
+    embeddings_model: EmbeddingsType,
+    base_url: Optional[str] = None,
+):
+    """Register an embeddings provider.
+
+    This function allows you to register custom embeddings providers that can be used
+    with the load_embeddings function. It supports both custom model classes and
+    string identifiers for supported providers.
+
+    Args:
+        provider_name: Name of the provider to register
+        embeddings_model: Either an Embeddings class or a string identifier
+            for a supported provider
+        base_url: The API address of the Embedding model provider (optional,
+            valid for both types of `embeddings_model`, but mainly used when
+            `embeddings_model` is a string and is "openai-compatible")
+
+    Raises:
+        ValueError: If base_url is not provided when embeddings_model is a string
+
+    Example:
+        # Register with custom model class:
+        >>> from langchain_dev_utils.embeddings import register_embeddings_provider, load_embeddings
+        >>> from langchain_core.embeddings.fake import FakeEmbeddings
+        >>>
+        >>> register_embeddings_provider("fakeembeddings", FakeEmbeddings)
+        >>> embeddings = load_embeddings("fakeembeddings:fake-embeddings",size=1024)
+        >>> embeddings.embed_query("hello world")
+
+        # Register with OpenAI-compatible API:
+        >>> register_embeddings_provider(
+        ...     "vllm",
+        ...     "openai-compatible",
+        ...     base_url="http://localhost:8000/v1"
+        ... )
+        >>> embeddings = load_embeddings("vllm:qwen3-embedding-4b")
+        >>> embeddings.embed_query("hello world")
+    """
+    _validate_provider_name(provider_name)
+    base_url = base_url or from_env(f"{provider_name.upper()}_API_BASE", default=None)()
+    if isinstance(embeddings_model, str):
+        if base_url is None:
+            raise ValueError(
+                f"base_url must be provided or set {provider_name.upper()}_API_BASE environment variable when embeddings_model is a string"
+            )
+
+        if embeddings_model != "openai-compatible":
+            raise ValueError(
+                "when embeddings_model is a string, the value must be 'openai-compatible'"
+            )
+
+        _check_pkg_install("langchain_openai")
+        from .adapters.openai_compatible import _create_openai_compatible_embedding
+
+        embeddings_model = _create_openai_compatible_embedding(
+            provider=provider_name,
+            base_url=base_url,
+        )
+        _EMBEDDINGS_PROVIDERS_DICT.update(
+            {
+                provider_name: {
+                    "embeddings_model": embeddings_model,
+                }
+            }
+        )
+    else:
+        if base_url is not None:
+            _EMBEDDINGS_PROVIDERS_DICT.update(
+                {
+                    provider_name: {
+                        "embeddings_model": embeddings_model,
+                        "base_url": base_url,
+                    }
+                }
+            )
+        else:
+            _EMBEDDINGS_PROVIDERS_DICT.update(
+                {provider_name: {"embeddings_model": embeddings_model}}
+            )
+
+
+def batch_register_embeddings_provider(
+    providers: list[EmbeddingProvider],
+):
+    """Batch register embeddings providers.
+
+    This function allows you to register multiple embeddings providers at once,
+    which is useful when setting up applications that need to work with multiple
+    embedding services.
+
+    Args:
+        providers: List of EmbeddingProvider dictionaries, each containing:
+            - provider_name: str - Provider name
+            - embeddings_model: Union[Type[Embeddings], str] - Model class or
+                provider string
+            - base_url: The API address of the Embedding model provider
+                (optional, valid for both types of `embeddings_model`, but
+                mainly used when `embeddings_model` is a string and is
+                "openai-compatible")
+
+    Raises:
+        ValueError: If any of the providers are invalid
+
+    Example:
+        # Register multiple providers at once:
+        >>> from langchain_dev_utils.embeddings import batch_register_embeddings_provider, load_embeddings
+        >>> from langchain_core.embeddings.fake import FakeEmbeddings
+        >>>
+        >>> batch_register_embeddings_provider(
+        ...     [
+        ...         {
+        ...             "provider_name": "fakeembeddings",
+        ...             "embeddings_model": FakeEmbeddings,
+        ...         },
+        ...         {
+        ...             "provider_name": "vllm",
+        ...             "embeddings_model":  "openai-compatible",
+        ...             "base_url": "http://localhost:8000/v1"
+        ...         },
+        ...     ]
+        ... )
+        >>> embeddings = load_embeddings("vllm:qwen3-embedding-4b")
+        >>> embeddings.embed_query("hello world")
+        >>> embeddings = load_embeddings("fakeembeddings:fake-embeddings", size=1024)
+        >>> embeddings.embed_query("hello world")
+    """
+    for provider in providers:
+        register_embeddings_provider(
+            provider["provider_name"],
+            provider["embeddings_model"],
+            provider.get("base_url"),
+        )
+
+
+def load_embeddings(
+    model: str,
+    *,
+    provider: Optional[str] = None,
+    **kwargs: Any,
+) -> Embeddings:
+    """Load embeddings model.
+
+    This function loads an embeddings model from the registered providers. The model parameter
+    must be specified in the format "provider:model-name" when provider is not specified separately.
+
+    Args:
+        model: Model name in format 'provider:model-name' if provider not specified separately
+        provider: Optional provider name (if not included in model parameter)
+        **kwargs: Additional arguments for model initialization (e.g., api_key)
+
+    Returns:
+        Embeddings: Initialized embeddings model instance
+
+    Raises:
+        ValueError: If provider is not registered or API key is not found
+
+    Example:
+        # Load model with provider prefix:
+        >>> from langchain_dev_utils.embeddings import load_embeddings
+        >>> embeddings = load_embeddings("vllm:qwen3-embedding-4b")
+        >>> embeddings.embed_query("hello world")
+
+        # Load model with separate provider parameter:
+        >>> embeddings = load_embeddings("qwen3-embedding-4b", provider="vllm")
+        >>> embeddings.embed_query("hello world")
+    """
+    if provider is None:
+        provider, model = _parse_model_string(model)
+    if provider not in list(_EMBEDDINGS_PROVIDERS_DICT.keys()) + list(
+        _SUPPORTED_PROVIDERS
+    ):
+        raise ValueError(f"Provider {provider} not registered")
+
+    if provider in _EMBEDDINGS_PROVIDERS_DICT:
+        embeddings = _EMBEDDINGS_PROVIDERS_DICT[provider]["embeddings_model"]
+        if base_url := _EMBEDDINGS_PROVIDERS_DICT[provider].get("base_url"):
+            url_key = _get_base_url_field_name(embeddings)
+            if url_key is not None:
+                kwargs.update({url_key: base_url})
+        return embeddings(model=model, **kwargs)
+    else:
+        return init_embeddings(model, provider=provider, **kwargs)

langchain_dev_utils/message_convert/__init__.py

@@ -0,0 +1,15 @@
+from .content import (
+    aconvert_reasoning_content_for_chunk_iterator,
+    convert_reasoning_content_for_ai_message,
+    convert_reasoning_content_for_chunk_iterator,
+    merge_ai_message_chunk,
+)
+from .format import format_sequence
+
+__all__ = [
+    "convert_reasoning_content_for_ai_message",
+    "convert_reasoning_content_for_chunk_iterator",
+    "aconvert_reasoning_content_for_chunk_iterator",
+    "merge_ai_message_chunk",
+    "format_sequence",
+]

langchain_dev_utils/message_convert/content.py

@@ -0,0 +1,201 @@
+from functools import reduce
+from typing import AsyncIterator, Iterator, Sequence, Tuple, cast
+
+from langchain_core.messages import AIMessage, AIMessageChunk
+
+
+def _get_reasoning_content(model_response: AIMessage | AIMessageChunk) -> str | None:
+    reasoning_content = None
+
+    reasoning_content_block = [
+        block for block in model_response.content_blocks if block["type"] == "reasoning"
+    ]
+    if reasoning_content_block:
+        reasoning_content = reasoning_content_block[0].get("reasoning")
+
+    if not reasoning_content:
+        reasoning_content = model_response.additional_kwargs.get("reasoning_content")
+
+    return reasoning_content
+
+
+def convert_reasoning_content_for_ai_message(
+    model_response: AIMessage,
+    think_tag: Tuple[str, str] = ("<think>", "</think>"),
+) -> AIMessage:
+    """Convert reasoning content in AI message to visible content.
+
+    This function extracts reasoning content from the additional_kwargs of an AI message
+    and merges it into the visible content, wrapping it with the specified tags.
+
+    Args:
+        model_response: AI message response from model
+        think_tag: Tuple of (opening_tag, closing_tag) to wrap reasoning content
+
+    Returns:
+        AIMessage: Modified AI message with reasoning content in visible content
+
+    Example:
+        # Basic usage with default tags:
+        >>> from langchain_dev_utils.message_convert import convert_reasoning_content_for_ai_message
+        >>> response = model.invoke("Explain quantum computing")
+        >>> response = convert_reasoning_content_for_ai_message(response)
+        >>> response.content
+
+        # Custom tags for reasoning content:
+        >>> response = convert_reasoning_content_for_ai_message(
+        ...     response, think_tag=('<reasoning>', '</reasoning>')
+        ... )
+        >>> response.content
+    """
+
+    reasoning_content = _get_reasoning_content(model_response)
+
+    if reasoning_content:
+        return model_response.model_copy(
+            update={
+                "content": f"{think_tag[0]}{reasoning_content}{think_tag[1]}{model_response.content}"
+            }
+        )
+    return model_response
+
+
+def convert_reasoning_content_for_chunk_iterator(
+    model_response: Iterator[AIMessageChunk | AIMessage],
+    think_tag: Tuple[str, str] = ("<think>", "</think>"),
+) -> Iterator[AIMessageChunk | AIMessage]:
+    """Convert reasoning content for streaming response chunks.
+
+    This function processes streaming response chunks and merges reasoning content
+    into the visible content, wrapping it with the specified tags. It handles
+    the first chunk, middle chunks, and last chunk differently to properly
+    format the reasoning content.
+
+    Args:
+        model_response: Iterator of message chunks from streaming response
+        think_tag: Tuple of (opening_tag, closing_tag) to wrap reasoning content
+
+    Yields:
+        BaseMessageChunk: Modified message chunks with reasoning content
+
+    Example:
+        # Process streaming response:
+        >>> from langchain_dev_utils.message_convert import convert_reasoning_content_for_chunk_iterator
+        >>> for chunk in convert_reasoning_content_for_chunk_iterator(
+        ...     model.stream("What is the capital of France?")
+        ... ):
+        ...     print(chunk.content, end="", flush=True)
+
+        # Custom tags for streaming:
+        >>> for chunk in convert_reasoning_content_for_chunk_iterator(
+        ...     model.stream("Explain quantum computing"),
+        ...     think_tag=('<reasoning>', '</reasoning>')
+        ... ):
+        ...     print(chunk.content, end="", flush=True)
+    """
+    isfirst = True
+    isend = True
+
+    for chunk in model_response:
+        if isinstance(chunk, AIMessageChunk):
+            reasoning_content = _get_reasoning_content(chunk)
+            if reasoning_content:
+                if isfirst:
+                    chunk = chunk.model_copy(
+                        update={"content": f"{think_tag[0]}{reasoning_content}"}
+                    )
+                    isfirst = False
+                else:
+                    chunk = chunk.model_copy(update={"content": reasoning_content})
+            elif chunk.content and isend and not isfirst:
+                chunk = chunk.model_copy(
+                    update={"content": f"{think_tag[1]}{chunk.content}"}
+                )
+                isend = False
+        yield chunk
+
+
+async def aconvert_reasoning_content_for_chunk_iterator(
+    model_response: AsyncIterator[AIMessageChunk | AIMessage],
+    think_tag: Tuple[str, str] = ("<think>", "</think>"),
+) -> AsyncIterator[AIMessageChunk | AIMessage]:
+    """Async convert reasoning content for streaming response chunks.
+
+    This is the async version of convert_reasoning_content_for_chunk_iterator.
+    It processes async streaming response chunks and merges reasoning content
+    into the visible content, wrapping it with the specified tags.
+
+    Args:
+        model_response: Async iterator of message chunks from streaming response
+        think_tag: Tuple of (opening_tag, closing_tag) to wrap reasoning content
+
+    Yields:
+        BaseMessageChunk: Modified message chunks with reasoning content
+
+    Example:
+        # Process async streaming response:
+        >>> from langchain_dev_utils.message_convert import aconvert_reasoning_content_for_chunk_iterator
+        >>> async for chunk in aconvert_reasoning_content_for_chunk_iterator(
+        ...     model.astream("What is the capital of France?")
+        ... ):
+        ...     print(chunk.content, end="", flush=True)
+
+        # Custom tags for async streaming:
+        >>> async for chunk in aconvert_reasoning_content_for_chunk_iterator(
+        ...     model.astream("Explain quantum computing"),
+        ...     think_tag=('<reasoning>', '</reasoning>')
+        ... ):
+        ...     print(chunk.content, end="", flush=True)
+    """
+    isfirst = True
+    isend = True
+
+    async for chunk in model_response:
+        if isinstance(chunk, AIMessageChunk):
+            reasoning_content = _get_reasoning_content(chunk)
+            if reasoning_content:
+                if isfirst:
+                    chunk = chunk.model_copy(
+                        update={"content": f"{think_tag[0]}{reasoning_content}"}
+                    )
+                    isfirst = False
+                else:
+                    chunk = chunk.model_copy(update={"content": reasoning_content})
+            elif chunk.content and isend and not isfirst:
+                chunk = chunk.model_copy(
+                    update={"content": f"{think_tag[1]}{chunk.content}"}
+                )
+                isend = False
+        yield chunk
+
+
+def merge_ai_message_chunk(chunks: Sequence[AIMessageChunk]) -> AIMessage:
+    """Merge a sequence of AIMessageChunk into a single AIMessage.
+
+    This function combines multiple message chunks into a single message,
+    preserving the content and metadata while handling tool calls appropriately.
+
+    Args:
+        chunks: Sequence of AIMessageChunk to merge
+
+    Returns:
+        AIMessage: Merged AIMessage
+
+    Example:
+        # Merge streaming chunks:
+        >>> from langchain_dev_utils.message_convert import merge_ai_message_chunk
+        >>> merged_message = merge_ai_message_chunk(list(model.stream("What is the capital of France?")))
+        >>> merged_message.content
+    """
+    ai_message_chunk = cast(AIMessageChunk, reduce(lambda x, y: x + y, chunks))
+    ai_message_chunk.additional_kwargs.pop("tool_calls", None)
+
+    data = {
+        "id": ai_message_chunk.id,
+        "content": ai_message_chunk.content,
+        "response_metadata": ai_message_chunk.response_metadata,
+        "additional_kwargs": ai_message_chunk.additional_kwargs,
+    }
+    if hasattr(ai_message_chunk, "tool_calls") and len(ai_message_chunk.tool_calls):
+        data["tool_calls"] = ai_message_chunk.tool_calls
+    return AIMessage.model_validate(data)

langchain_dev_utils/message_convert/format.py

@@ -0,0 +1,69 @@
+from typing import Sequence, Union
+
+from langchain_core.documents import Document
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+    ToolMessage,
+)
+
+
+def format_sequence(
+    inputs: Union[Sequence[Document], Sequence[BaseMessage], Sequence[str]],
+    separator: str = "-",
+    with_num: bool = False,
+) -> str:
+    """Convert a list of messages, documents, or strings into a formatted string.
+
+    This function extracts text content from various types (e.g., HumanMessage, Document)
+    and joins them into a single string. Optionally adds serial numbers and a custom
+    separator between items.
+
+    Args:
+        inputs: A list of inputs. Supported types:
+            - langchain_core.messages: HumanMessage, AIMessage, SystemMessage, ToolMessage
+            - langchain_core.documents.Document
+            - str
+        separator: The separator used to join the items. Defaults to "-".
+        with_num: If True, prefixes each item with a serial number (e.g., "1. Hello").
+                  Defaults to False.
+
+    Returns:
+        A formatted string composed of the input contents, joined by `separator`.
+
+    Example:
+        # Format messages with default separator:
+        >>> from langchain_dev_utils.message_convert import format_sequence
+        >>> from langchain_core.messages import HumanMessage, AIMessage
+        >>> messages = [
+        ...     HumanMessage(content="Hello, how are you?"),
+        ...     AIMessage(content="I'm doing well, thank you!")
+        ... ]
+        >>> formatted = format_sequence(messages)
+        >>> formatted
+
+        # Format with custom separator and numbering:
+        >>> formatted = format_sequence(messages, separator="---", with_num=True)
+        >>> formatted
+    """
+    if not inputs:
+        return ""
+
+    outputs = []
+
+    for input_item in inputs:
+        if isinstance(
+            input_item, (HumanMessage, AIMessage, SystemMessage, ToolMessage)
+        ):
+            outputs.append(input_item.content)
+        elif isinstance(input_item, Document):
+            outputs.append(input_item.page_content)
+        elif isinstance(input_item, str):
+            outputs.append(input_item)
+    if with_num:
+        outputs = [f"{i + 1}. {output}" for i, output in enumerate(outputs)]
+
+    str_ = "\n" + separator
+    return separator + str_.join(outputs)

langchain_dev_utils/pipeline/__init__.py

@@ -0,0 +1,7 @@
+from .parallel import create_parallel_pipeline
+from .sequential import create_sequential_pipeline
+
+__all__ = [
+    "create_parallel_pipeline",
+    "create_sequential_pipeline",
+]

langchain_dev_utils/pipeline/parallel.py

@@ -0,0 +1,135 @@
+from typing import Awaitable, Callable, Optional, Union
+
+from langgraph.cache.base import BaseCache
+from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.store.base import BaseStore
+from langgraph.types import Checkpointer, Send
+from langgraph.typing import ContextT, InputT, OutputT, StateT
+
+from .types import SubGraph
+
+
+def create_parallel_pipeline(
+    sub_graphs: list[SubGraph],
+    state_schema: type[StateT],
+    graph_name: Optional[str] = None,
+    branches_fn: Optional[
+        Union[
+            Callable[..., list[Send]],
+            Callable[..., Awaitable[list[Send]]],
+        ]
+    ] = None,
+    context_schema: type[ContextT] | None = None,
+    input_schema: type[InputT] | None = None,
+    output_schema: type[OutputT] | None = None,
+    checkpointer: Checkpointer | None = None,
+    store: BaseStore | None = None,
+    cache: BaseCache | None = None,
+) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]:
+    """
+    Create a parallel pipeline from a list of subgraphs.
+
+    This function allows you to compose multiple StateGraphs in a parallel fashion,
+    where subgraphs can execute concurrently. This is useful for creating complex
+    multi-agent workflows where agents can work independently or with dynamic branching.
+
+    Args:
+        sub_graphs: List of sub-graphs to execute in parallel
+        state_schema: state schema of the final constructed graph
+        graph_name: Name of the final constructed graph
+        branches_fn: Optional function to determine which sub-graphs to execute
+            in parallel
+        context_schema: context schema of the final constructed graph
+        input_schema: input schema of the final constructed graph
+        output_schema: output schema of the final constructed graph
+        checkpointer: Optional LangGraph checkpointer for the final constructed
+            graph
+        store: Optional LangGraph store for the final constructed graph
+        cache: Optional LangGraph cache for the final constructed graph
+
+    Returns:
+        CompiledStateGraph[StateT, ContextT, InputT, OutputT]: Compiled state
+            graph of the pipeline.
+
+    Example:
+        # Basic parallel pipeline: multiple specialized agents run concurrently
+        >>> from langchain_dev_utils.pipeline import create_parallel_pipeline
+        >>>
+        >>> graph = create_parallel_pipeline(
+        ...     sub_graphs=[
+        ...        time_agent, weather_agent, user_agent
+        ...     ],
+        ...     state_schema=AgentState,
+        ...     graph_name="parallel_agents_pipeline",
+        ... )
+        >>>
+        >>> response = graph.invoke({"messages": [HumanMessage("Hello")]})
+
+        # Dynamic parallel pipeline: decide which agents to run based on conditional branches
+        >>> graph = create_parallel_pipeline(
+        ...     sub_graphs=[
+        ...         time_agent, weather_agent, user_agent
+        ...     ],
+        ...     state_schema=AgentState,
+        ...     branches_fn=lambda state: [
+        ...         Send("weather_agent", arg={"messages": [HumanMessage("Get current weather in New York")]}),
+        ...         Send("time_agent", arg={"messages": [HumanMessage("Get current time")]}),
+        ...     ],
+        ...     graph_name="dynamic_parallel_pipeline",
+        ... )
+        >>>
+        >>> response = graph.invoke({"messages": [HumanMessage("Hello")]})
+    """
+    graph = StateGraph(
+        state_schema=state_schema,
+        context_schema=context_schema,
+        input_schema=input_schema,
+        output_schema=output_schema,
+    )
+
+    subgraphs_names = set()
+
+    compiled_subgraphs: list[CompiledStateGraph] = []
+    for subgraph in sub_graphs:
+        if isinstance(subgraph, StateGraph):
+            subgraph = subgraph.compile()
+
+        compiled_subgraphs.append(subgraph)
+        if subgraph.name is None or subgraph.name == "LangGraph":
+            raise ValueError(
+                "Please specify a name when you create your agent, either via `create_react_agent(..., name=agent_name)` "
+                "or via `graph.compile(name=name)`."
+            )
+
+        if subgraph.name in subgraphs_names:
+            raise ValueError(
+                f"Subgraph with name '{subgraph.name}' already exists. Subgraph names must be unique."
+            )
+
+        subgraphs_names.add(subgraph.name)
+
+    for sub_graph in compiled_subgraphs:
+        graph.add_node(sub_graph.name, sub_graph)
+
+    if branches_fn:
+        graph.add_conditional_edges(
+            "__start__",
+            branches_fn,
+            [subgraph.name for subgraph in compiled_subgraphs],
+        )
+        return graph.compile(
+            name=graph_name or "parallel graph",
+            checkpointer=checkpointer,
+            store=store,
+            cache=cache,
+        )
+    else:
+        for i in range(len(compiled_subgraphs)):
+            graph.add_edge("__start__", compiled_subgraphs[i].name)
+        return graph.compile(
+            name=graph_name or "parallel graph",
+            checkpointer=checkpointer,
+            store=store,
+            cache=cache,
+        )