codekingpro/portable-devtools

1"""LangSmith document loader."""2 3import datetime4import json5import uuid6from collections.abc import Callable, Iterator, Sequence7from typing import Any8 9from langsmith import Client as LangSmithClient10from typing_extensions import override11 12from langchain_core.document_loaders.base import BaseLoader13from langchain_core.documents import Document14from langchain_core.tracers._compat import pydantic_to_dict15 16 17class LangSmithLoader(BaseLoader):18    """Load LangSmith Dataset examples as `Document` objects.19 20    Loads the example inputs as the `Document` page content and places the entire21    example into the `Document` metadata. This allows you to easily create few-shot22    example retrievers from the loaded documents.23 24    ??? example "Lazy loading"25 26        ```python27        from langchain_core.document_loaders import LangSmithLoader28 29        loader = LangSmithLoader(dataset_id="...", limit=100)30        docs = []31        for doc in loader.lazy_load():32            docs.append(doc)33        ```34 35        ```python36        # -> [Document("...", metadata={"inputs": {...}, "outputs": {...}, ...}), ...]37        ```38    """39 40    def __init__(41        self,42        *,43        dataset_id: uuid.UUID | str | None = None,44        dataset_name: str | None = None,45        example_ids: Sequence[uuid.UUID | str] | None = None,46        as_of: datetime.datetime | str | None = None,47        splits: Sequence[str] | None = None,48        inline_s3_urls: bool = True,49        offset: int = 0,50        limit: int | None = None,51        metadata: dict | None = None,52        filter: str | None = None,  # noqa: A00253        content_key: str = "",54        format_content: Callable[..., str] | None = None,55        client: LangSmithClient | None = None,56        **client_kwargs: Any,57    ) -> None:58        """Create a LangSmith loader.59 60        Args:61            dataset_id: The ID of the dataset to filter by.62            dataset_name: The name of the dataset to filter by.63            content_key: The inputs key to set as `Document` page content.64 65                `'.'` characters are interpreted as nested keys, e.g.66                `content_key="first.second"` will result in67                `Document(page_content=format_content(example.inputs["first"]["second"]))`68            format_content: Function for converting the content extracted from the example69                inputs into a string.70 71                Defaults to JSON-encoding the contents.72            example_ids: The IDs of the examples to filter by.73            as_of: The dataset version tag or timestamp to retrieve the examples as of.74 75                Response examples will only be those that were present at the time of76                the tagged (or timestamped) version.77            splits: A list of dataset splits, which are divisions of your dataset such78                as `train`, `test`, or `validation`.79 80                Returns examples only from the specified splits.81            inline_s3_urls: Whether to inline S3 URLs.82            offset: The offset to start from.83            limit: The maximum number of examples to return.84            metadata: Metadata to filter by.85            filter: A structured filter string to apply to the examples.86            client: LangSmith Client.87 88                If not provided will be initialized from below args.89            client_kwargs: Keyword args to pass to LangSmith client init.90 91                Should only be specified if `client` isn't.92 93        Raises:94            ValueError: If both `client` and `client_kwargs` are provided.95        """  # noqa: E50196        if client and client_kwargs:97            raise ValueError98        self._client = client or LangSmithClient(**client_kwargs)99        self.content_key = list(content_key.split(".")) if content_key else []100        self.format_content = format_content or _stringify101        self.dataset_id = dataset_id102        self.dataset_name = dataset_name103        self.example_ids = example_ids104        self.as_of = as_of105        self.splits = splits106        self.inline_s3_urls = inline_s3_urls107        self.offset = offset108        self.limit = limit109        self.metadata = metadata110        self.filter = filter111 112    @override113    def lazy_load(self) -> Iterator[Document]:114        for example in self._client.list_examples(115            dataset_id=self.dataset_id,116            dataset_name=self.dataset_name,117            example_ids=self.example_ids,118            as_of=self.as_of,119            splits=self.splits,120            inline_s3_urls=self.inline_s3_urls,121            offset=self.offset,122            limit=self.limit,123            metadata=self.metadata,124            filter=self.filter,125        ):126            content: Any = example.inputs127            for key in self.content_key:128                content = content[key]129            content_str = self.format_content(content)130            metadata = pydantic_to_dict(example)131            # Stringify datetime and UUID types.132            for k in ("dataset_id", "created_at", "modified_at", "source_run_id", "id"):133                metadata[k] = str(metadata[k]) if metadata[k] else metadata[k]134            yield Document(content_str, metadata=metadata)135 136 137def _stringify(x: str | dict[str, Any]) -> str:138    if isinstance(x, str):139        return x140    try:141        return json.dumps(x, indent=2)142    except Exception:143        return str(x)144