Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Python] Limit #552

Merged
merged 1 commit into from Mar 26, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
143 changes: 98 additions & 45 deletions python/langsmith/client.py
Expand Up @@ -1438,6 +1438,7 @@ def list_runs(
start_time: Optional[datetime.datetime] = None,
error: Optional[bool] = None,
run_ids: Optional[List[ID_TYPE]] = None,
limit: Optional[int] = None,
**kwargs: Any,
) -> Iterator[ls_schemas.Run]:
"""List runs from the LangSmith API.
Expand Down Expand Up @@ -1479,6 +1480,8 @@ def list_runs(
Whether to filter by error status.
run_ids : List[str or UUID] or None, default=None
The IDs of the runs to filter by.
limit : int or None, default=None
The maximum number of runs to return.
**kwargs : Any
Additional keyword arguments.

Expand Down Expand Up @@ -1571,12 +1574,14 @@ def list_runs(
**kwargs,
}
body_query = {k: v for k, v in body_query.items() if v is not None}
yield from (
ls_schemas.Run(**run, _host_url=self._host_url)
for run in self._get_cursor_paginated_list(
for i, run in enumerate(
self._get_cursor_paginated_list(
"/runs/query", body=body_query, request_method="post"
)
)
):
yield ls_schemas.Run(**run, _host_url=self._host_url)
if limit is not None and i + 1 >= limit:
break

def get_run_url(
self,
Expand Down Expand Up @@ -1794,6 +1799,7 @@ def list_shared_projects(
project_ids: Optional[List[ID_TYPE]] = None,
name: Optional[str] = None,
name_contains: Optional[str] = None,
limit: Optional[int] = None,
) -> Iterator[ls_schemas.TracerSessionResult]:
"""List shared projects.

Expand All @@ -1806,19 +1812,22 @@ def list_shared_projects(
Name of the project to filter the results, by default None.
name_contains : str, optional
Substring to search for in project names, by default None.
limit : int, optional

Yields:
TracerSessionResult: The shared projects.
"""
params = {"id": project_ids, "name": name, "name_contains": name_contains}
share_token = _as_uuid(dataset_share_token, "dataset_share_token")
yield from [
ls_schemas.TracerSessionResult(**project, _host_url=self._host_url)
for project in self._get_paginated_list(
for i, project in enumerate(
self._get_paginated_list(
f"/public/{share_token}/datasets/sessions",
params=params,
)
]
):
yield ls_schemas.TracerSessionResult(**project, _host_url=self._host_url)
if limit is not None and i + 1 >= limit:
break

def create_project(
self,
Expand Down Expand Up @@ -2093,6 +2102,7 @@ def list_projects(
reference_dataset_id: Optional[ID_TYPE] = None,
reference_dataset_name: Optional[str] = None,
reference_free: Optional[bool] = None,
limit: Optional[int] = None,
) -> Iterator[ls_schemas.TracerSession]:
"""List projects from the LangSmith API.

Expand All @@ -2110,13 +2120,17 @@ def list_projects(
The name of the reference dataset to filter by, by default None
reference_free : Optional[bool], optional
Whether to filter for only projects not associated with a dataset.
limit : Optional[int], optional
The maximum number of projects to return, by default None

Yields:
------
TracerSession
The projects.
"""
params: Dict[str, Any] = {}
params: Dict[str, Any] = {
"limit": min(limit, 100) if limit is not None else 100
}
if project_ids is not None:
params["id"] = project_ids
if name is not None:
Expand All @@ -2137,10 +2151,12 @@ def list_projects(
params["reference_dataset"] = reference_dataset_id
if reference_free is not None:
params["reference_free"] = reference_free
yield from (
ls_schemas.TracerSessionResult(**project, _host_url=self._host_url)
for project in self._get_paginated_list("/sessions", params=params)
)
for i, project in enumerate(
self._get_paginated_list("/sessions", params=params)
):
yield ls_schemas.TracerSession(**project, _host_url=self._host_url)
if limit is not None and i + 1 >= limit:
break

@ls_utils.xor_args(("project_name", "project_id"))
def delete_project(
Expand Down Expand Up @@ -2389,6 +2405,7 @@ def list_datasets(
data_type: Optional[str] = None,
dataset_name: Optional[str] = None,
dataset_name_contains: Optional[str] = None,
limit: Optional[int] = None,
) -> Iterator[ls_schemas.Dataset]:
"""List the datasets on the LangSmith API.

Expand All @@ -2397,7 +2414,9 @@ def list_datasets(
Dataset
The datasets.
"""
params: Dict[str, Any] = {}
params: Dict[str, Any] = {
"limit": min(limit, 100) if limit is not None else 100
}
if dataset_ids is not None:
params["id"] = dataset_ids
if data_type is not None:
Expand All @@ -2406,15 +2425,16 @@ def list_datasets(
params["name"] = dataset_name
if dataset_name_contains is not None:
params["name_contains"] = dataset_name_contains

yield from (
ls_schemas.Dataset(
for i, dataset in enumerate(
self._get_paginated_list("/datasets", params=params)
):
yield ls_schemas.Dataset(
**dataset,
_host_url=self._host_url,
_tenant_id=self._get_optional_tenant_id(),
)
for dataset in self._get_paginated_list("/datasets", params=params)
)
if limit is not None and i + 1 >= limit:
break

@ls_utils.xor_args(("dataset_id", "dataset_name"))
def delete_dataset(
Expand Down Expand Up @@ -2503,27 +2523,34 @@ def list_dataset_versions(
dataset_id: Optional[ID_TYPE] = None,
dataset_name: Optional[str] = None,
search: Optional[str] = None,
limit: Optional[int] = None,
) -> Iterator[ls_schemas.DatasetVersion]:
"""List dataset versions.

Args:
dataset_id (Optional[ID_TYPE]): The ID of the dataset.
dataset_name (Optional[str]): The name of the dataset.
search (Optional[str]): The search query.
limit (Optional[int]): The maximum number of versions to return.

Returns:
Iterator[ls_schemas.DatasetVersion]: An iterator of dataset versions.
"""
if dataset_id is None:
dataset_id = self.read_dataset(dataset_name=dataset_name).id
params = {"search": search}
yield from (
ls_schemas.DatasetVersion(**version)
for version in self._get_paginated_list(
params = {
"search": search,
"limit": min(limit, 100) if limit is not None else 100,
}
for i, version in enumerate(
self._get_paginated_list(
f"/datasets/{_as_uuid(dataset_id, 'dataset_id')}/versions",
params=params,
)
)
):
yield ls_schemas.DatasetVersion(**version)
if limit is not None and i + 1 >= limit:
break

def read_dataset_version(
self,
Expand Down Expand Up @@ -2944,6 +2971,7 @@ def list_examples(
example_ids: Optional[Sequence[ID_TYPE]] = None,
as_of: Optional[Union[datetime.datetime, str]] = None,
inline_s3_urls: bool = True,
limit: Optional[int] = None,
**kwargs: Any,
) -> Iterator[ls_schemas.Example]:
"""Retrieve the example rows of the specified dataset.
Expand All @@ -2961,6 +2989,7 @@ def list_examples(
of the tagged (or timestamped) version.
inline_s3_urls (bool, optional): Whether to inline S3 URLs.
Defaults to True.
limit (int, optional): The maximum number of examples to return.

Yields:
Example: The examples.
Expand All @@ -2972,6 +3001,7 @@ def list_examples(
as_of.isoformat() if isinstance(as_of, datetime.datetime) else as_of
),
"inline_s3_urls": inline_s3_urls,
"limit": min(limit, 100) if limit is not None else 100,
}
if dataset_id is not None:
params["dataset"] = dataset_id
Expand All @@ -2980,14 +3010,16 @@ def list_examples(
params["dataset"] = dataset_id
else:
pass
yield from (
ls_schemas.Example(
for i, example in enumerate(
self._get_paginated_list("/examples", params=params)
):
yield ls_schemas.Example(
**example,
_host_url=self._host_url,
_tenant_id=self._get_optional_tenant_id(),
)
for example in self._get_paginated_list("/examples", params=params)
)
if limit is not None and i + 1 >= limit:
break

def update_example(
self,
Expand Down Expand Up @@ -3436,6 +3468,7 @@ def list_feedback(
run_ids: Optional[Sequence[ID_TYPE]] = None,
feedback_key: Optional[Sequence[str]] = None,
feedback_source_type: Optional[Sequence[ls_schemas.FeedbackSourceType]] = None,
limit: Optional[int] = None,
**kwargs: Any,
) -> Iterator[ls_schemas.Feedback]:
"""List the feedback objects on the LangSmith API.
Expand All @@ -3450,6 +3483,7 @@ def list_feedback(
feedback_source_type: List[FeedbackSourceType] or None, default=None
The type of feedback source, such as model
(for model-generated feedback) or API.
limit : int or None, default=None
**kwargs : Any
Additional keyword arguments.

Expand All @@ -3460,16 +3494,19 @@ def list_feedback(
"""
params: dict = {
"run": run_ids,
"limit": min(limit, 100) if limit is not None else 100,
**kwargs,
}
if feedback_key is not None:
params["key"] = feedback_key
if feedback_source_type is not None:
params["source"] = feedback_source_type
yield from (
ls_schemas.Feedback(**feedback)
for feedback in self._get_paginated_list("/feedback", params=params)
)
for i, feedback in enumerate(
self._get_paginated_list("/feedback", params=params)
):
yield ls_schemas.Feedback(**feedback)
if limit is not None and i + 1 >= limit:
break

def delete_feedback(self, feedback_id: ID_TYPE) -> None:
"""Delete a feedback by ID.
Expand Down Expand Up @@ -3602,23 +3639,29 @@ def create_presigned_feedback_token(
def list_presigned_feedback_tokens(
self,
run_id: ID_TYPE,
*,
limit: Optional[int] = None,
) -> Iterator[ls_schemas.FeedbackIngestToken]:
"""List the feedback ingest tokens for a run.

Args:
run_id: The ID of the run to filter by.
limit: The maximum number of tokens to return.

Yields:
FeedbackIngestToken
The feedback ingest tokens.
"""
params = {
"run_id": _as_uuid(run_id, "run_id"),
"limit": min(limit, 100) if limit is not None else 100,
}
yield from (
ls_schemas.FeedbackIngestToken(**token)
for token in self._get_paginated_list("/feedback/tokens", params=params)
)
for i, token in enumerate(
self._get_paginated_list("/feedback/tokens", params=params)
):
yield ls_schemas.FeedbackIngestToken(**token)
if limit is not None and i + 1 >= limit:
break

# Annotation Queue API

Expand All @@ -3628,6 +3671,7 @@ def list_annotation_queues(
queue_ids: Optional[List[ID_TYPE]] = None,
name: Optional[str] = None,
name_contains: Optional[str] = None,
limit: Optional[int] = None,
) -> Iterator[ls_schemas.AnnotationQueue]:
"""List the annotation queues on the LangSmith API.

Expand All @@ -3638,6 +3682,7 @@ def list_annotation_queues(
The name of the queue to filter by.
name_contains : str or None, default=None
The substring that the queue name should contain.
limit : int or None, default=None

Yields:
AnnotationQueue
Expand All @@ -3651,15 +3696,18 @@ def list_annotation_queues(
),
"name": name,
"name_contains": name_contains,
"limit": min(limit, 100) if limit is not None else 100,
}
yield from (
ls_schemas.AnnotationQueue(
for i, queue in enumerate(
self._get_paginated_list("/annotation-queues", params=params)
):
yield ls_schemas.AnnotationQueue(
**queue,
_host_url=self._host_url,
_tenant_id=self._get_optional_tenant_id(),
)
for queue in self._get_paginated_list("/annotation-queues", params=params)
)
if limit is not None and i + 1 >= limit:
break

def create_annotation_queue(
self,
Expand Down Expand Up @@ -3773,7 +3821,7 @@ def add_runs_to_annotation_queue(
ls_utils.raise_for_status_with_text(response)

def list_runs_from_annotation_queue(
self, queue_id: ID_TYPE
self, queue_id: ID_TYPE, *, limit: Optional[int] = None
) -> Iterator[ls_schemas.RunWithAnnotationQueueInfo]:
"""List runs from an annotation queue with the specified queue ID.

Expand All @@ -3785,10 +3833,15 @@ def list_runs_from_annotation_queue(
annotation queue.
"""
path = f"/annotation-queues/{_as_uuid(queue_id, 'queue_id')}/runs"
yield from (
ls_schemas.RunWithAnnotationQueueInfo(**run)
for run in self._get_paginated_list(path, params={"headers": self._headers})
)
limit_ = min(limit, 100) if limit is not None else 100
for i, run in enumerate(
self._get_paginated_list(
path, params={"headers": self._headers, "limit": limit_}
)
):
yield ls_schemas.RunWithAnnotationQueueInfo(**run)
if limit is not None and i + 1 >= limit:
break

async def arun_on_dataset(
self,
Expand Down