Source code for langchain_community.chat_loaders.langsmith

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Dict, Iterable, Iterator, List, Optional, Union, cast

from langchain_core.chat_loaders import BaseChatLoader
from langchain_core.chat_sessions import ChatSession
from langchain_core.load.load import load

if TYPE_CHECKING:
    from langsmith.client import Client
    from langsmith.schemas import Run

logger = logging.getLogger(__name__)


class LangSmithRunChatLoader(BaseChatLoader):
    """从LangSmith的“llm”运行列表中加载聊天会话。

属性：
    runs（Iterable[Union[str, Run]]）：LLM运行ID或运行对象的列表。
    client（Client）：用于获取数据的LangSmith客户端实例。"""

    def __init__(
        self, runs: Iterable[Union[str, Run]], client: Optional["Client"] = None
    ):
        """初始化一个新的LangSmithRunChatLoader实例。

:param runs: LLM运行ID或运行对象的列表。
:param client: LangSmith客户端的一个实例，如果未提供，
    将创建一个新的客户端实例。
"""
        from langsmith.client import Client

        self.runs = runs
        self.client = client or Client()

    def _load_single_chat_session(self, llm_run: "Run") -> ChatSession:
        """将单个LangSmith LLM运行转换为ChatSession。

:param llm_run: LLM运行对象。
:return: 表示运行数据的聊天会话。
"""
        chat_session = LangSmithRunChatLoader._get_messages_from_llm_run(llm_run)
        functions = LangSmithRunChatLoader._get_functions_from_llm_run(llm_run)
        if functions:
            chat_session["functions"] = functions
        return chat_session

    @staticmethod
    def _get_messages_from_llm_run(llm_run: "Run") -> ChatSession:
        """从LangSmith LLM运行中提取消息。

:param llm_run: LLM运行对象。
:return: 包含提取消息的ChatSession。
"""
        if llm_run.run_type != "llm":
            raise ValueError(f"Expected run of type llm. Got: {llm_run.run_type}")
        if "messages" not in llm_run.inputs:
            raise ValueError(f"Run has no 'messages' inputs. Got {llm_run.inputs}")
        if not llm_run.outputs:
            raise ValueError("Cannot convert pending run")
        messages = load(llm_run.inputs)["messages"]
        message_chunk = load(llm_run.outputs)["generations"][0]["message"]
        return ChatSession(messages=messages + [message_chunk])

    @staticmethod
    def _get_functions_from_llm_run(llm_run: "Run") -> Optional[List[Dict]]:
        """从LangSmith LLM运行中提取函数（如果存在）。

:param llm_run: LLM运行对象。
:return: 运行中的函数或None。
"""
        if llm_run.run_type != "llm":
            raise ValueError(f"Expected run of type llm. Got: {llm_run.run_type}")
        return (llm_run.extra or {}).get("invocation_params", {}).get("functions")

    def lazy_load(self) -> Iterator[ChatSession]:
        """从运行ID的可迭代对象中延迟加载聊天会话。

该方法获取运行并即时将其转换为聊天会话，每次生成一个会话。

:return: 包含消息的聊天会话的迭代器。
"""
        from langsmith.schemas import Run

        for run_obj in self.runs:
            try:
                if hasattr(run_obj, "id"):
                    run = run_obj
                else:
                    run = self.client.read_run(run_obj)
                session = self._load_single_chat_session(cast(Run, run))
                yield session
            except ValueError as e:
                logger.warning(f"Could not load run {run_obj}: {repr(e)}")
                continue


class LangSmithDatasetChatLoader(BaseChatLoader):
    """从LangSmith数据集中加载具有“chat”数据类型的聊天会话。

属性：
    dataset_name（str）：LangSmith数据集的名称。
    client（Client）：用于获取数据的LangSmith客户端实例。"""

    def __init__(self, *, dataset_name: str, client: Optional["Client"] = None):
        """初始化一个新的LangSmithChatDatasetLoader实例。

:param dataset_name: LangSmith数据集的名称。
:param client: LangSmith客户端的一个实例；如果未提供，
    将创建一个新的客户端实例。
"""
        try:
            from langsmith.client import Client
        except ImportError as e:
            raise ImportError(
                "The LangSmith client is required to load LangSmith datasets.\n"
                "Please install it with `pip install langsmith`"
            ) from e

        self.dataset_name = dataset_name
        self.client = client or Client()

    def lazy_load(self) -> Iterator[ChatSession]:
        """从指定的LangSmith数据集中延迟加载聊天会话。

该方法从数据集中获取聊天数据，并在运行时将每个数据点转换为聊天会话，一次生成一个会话。

:return: 包含消息的聊天会话的迭代器。
"""
        from langchain_community.adapters import openai as oai_adapter  # noqa: E402

        data = self.client.read_dataset_openai_finetuning(
            dataset_name=self.dataset_name
        )
        for data_point in data:
            yield ChatSession(
                messages=[
                    oai_adapter.convert_dict_to_message(m)
                    for m in data_point.get("messages", [])
                ],
                functions=data_point.get("functions"),
            )