865 lines
29 KiB
Python
865 lines
29 KiB
Python
|
import logging
|
||
|
|
import uuid
|
||
|
|
from typing import Any, Callable, Dict, List, Literal, Optional, TypeVar, Union, cast
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
from typing_extensions import deprecated
|
||
|
|
|
||
|
|
from literalai.api.base import BaseLiteralAPI, prepare_variables
|
||
|
|
from literalai.api.helpers.attachment_helpers import (
|
||
|
|
AttachmentUpload,
|
||
|
|
create_attachment_helper,
|
||
|
|
delete_attachment_helper,
|
||
|
|
get_attachment_helper,
|
||
|
|
update_attachment_helper,
|
||
|
|
)
|
||
|
|
from literalai.api.helpers.dataset_helpers import (
|
||
|
|
add_generation_to_dataset_helper,
|
||
|
|
add_step_to_dataset_helper,
|
||
|
|
create_dataset_helper,
|
||
|
|
create_dataset_item_helper,
|
||
|
|
create_experiment_helper,
|
||
|
|
create_experiment_item_helper,
|
||
|
|
delete_dataset_helper,
|
||
|
|
delete_dataset_item_helper,
|
||
|
|
get_dataset_helper,
|
||
|
|
get_dataset_item_helper,
|
||
|
|
update_dataset_helper,
|
||
|
|
)
|
||
|
|
from literalai.api.helpers.generation_helpers import (
|
||
|
|
create_generation_helper,
|
||
|
|
get_generations_helper,
|
||
|
|
)
|
||
|
|
from literalai.api.helpers.prompt_helpers import (
|
||
|
|
PromptRollout,
|
||
|
|
create_prompt_helper,
|
||
|
|
create_prompt_lineage_helper,
|
||
|
|
create_prompt_variant_helper,
|
||
|
|
get_prompt_ab_testing_helper,
|
||
|
|
get_prompt_helper,
|
||
|
|
get_prompt_lineage_helper,
|
||
|
|
update_prompt_ab_testing_helper,
|
||
|
|
)
|
||
|
|
from literalai.api.helpers.score_helpers import (
|
||
|
|
ScoreUpdate,
|
||
|
|
check_scores_finite,
|
||
|
|
create_score_helper,
|
||
|
|
create_scores_query_builder,
|
||
|
|
delete_score_helper,
|
||
|
|
get_scores_helper,
|
||
|
|
update_score_helper,
|
||
|
|
)
|
||
|
|
from literalai.api.helpers.step_helpers import (
|
||
|
|
create_step_helper,
|
||
|
|
delete_step_helper,
|
||
|
|
get_step_helper,
|
||
|
|
get_steps_helper,
|
||
|
|
send_steps_helper,
|
||
|
|
update_step_helper,
|
||
|
|
)
|
||
|
|
from literalai.api.helpers.thread_helpers import (
|
||
|
|
create_thread_helper,
|
||
|
|
delete_thread_helper,
|
||
|
|
get_thread_helper,
|
||
|
|
get_threads_helper,
|
||
|
|
list_threads_helper,
|
||
|
|
update_thread_helper,
|
||
|
|
upsert_thread_helper,
|
||
|
|
)
|
||
|
|
from literalai.api.helpers.user_helpers import (
|
||
|
|
create_user_helper,
|
||
|
|
delete_user_helper,
|
||
|
|
get_user_helper,
|
||
|
|
get_users_helper,
|
||
|
|
update_user_helper,
|
||
|
|
)
|
||
|
|
from literalai.context import active_steps_var, active_thread_var
|
||
|
|
from literalai.evaluation.dataset import Dataset, DatasetType
|
||
|
|
from literalai.evaluation.dataset_experiment import (
|
||
|
|
DatasetExperiment,
|
||
|
|
DatasetExperimentItem,
|
||
|
|
)
|
||
|
|
from literalai.evaluation.dataset_item import DatasetItem
|
||
|
|
from literalai.my_types import PaginatedResponse, User
|
||
|
|
from literalai.observability.filter import (
|
||
|
|
generations_filters,
|
||
|
|
generations_order_by,
|
||
|
|
scores_filters,
|
||
|
|
scores_order_by,
|
||
|
|
steps_filters,
|
||
|
|
steps_order_by,
|
||
|
|
threads_filters,
|
||
|
|
threads_order_by,
|
||
|
|
users_filters,
|
||
|
|
)
|
||
|
|
from literalai.observability.generation import (
|
||
|
|
BaseGeneration,
|
||
|
|
ChatGeneration,
|
||
|
|
CompletionGeneration,
|
||
|
|
GenerationMessage,
|
||
|
|
)
|
||
|
|
from literalai.observability.step import (
|
||
|
|
Attachment,
|
||
|
|
Score,
|
||
|
|
ScoreDict,
|
||
|
|
ScoreType,
|
||
|
|
Step,
|
||
|
|
StepDict,
|
||
|
|
StepType,
|
||
|
|
)
|
||
|
|
from literalai.observability.thread import Thread
|
||
|
|
from literalai.prompt_engineering.prompt import Prompt, ProviderSettings
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class LiteralAPI(BaseLiteralAPI):
|
||
|
|
"""
|
||
|
|
```python
|
||
|
|
from literalai import LiteralClient
|
||
|
|
# Initialize the client
|
||
|
|
literalai_client = LiteralClient(api_key="your_api_key_here")
|
||
|
|
# Access the API's methods
|
||
|
|
print(literalai_client.api)
|
||
|
|
```
|
||
|
|
"""
|
||
|
|
|
||
|
|
R = TypeVar("R")
|
||
|
|
|
||
|
|
def make_gql_call(
|
||
|
|
self,
|
||
|
|
description: str,
|
||
|
|
query: str,
|
||
|
|
variables: dict[str, Any],
|
||
|
|
timeout: Optional[int] = 10,
|
||
|
|
) -> dict:
|
||
|
|
def raise_error(error):
|
||
|
|
logger.error(f"Failed to {description}: {error}")
|
||
|
|
raise Exception(error)
|
||
|
|
|
||
|
|
variables = prepare_variables(variables)
|
||
|
|
with httpx.Client(follow_redirects=True) as client:
|
||
|
|
response = client.post(
|
||
|
|
self.graphql_endpoint,
|
||
|
|
json={"query": query, "variables": variables},
|
||
|
|
headers=self.headers,
|
||
|
|
timeout=timeout,
|
||
|
|
)
|
||
|
|
|
||
|
|
try:
|
||
|
|
response.raise_for_status()
|
||
|
|
except httpx.HTTPStatusError:
|
||
|
|
raise_error(f"Failed to {description}: {response.text}")
|
||
|
|
|
||
|
|
try:
|
||
|
|
json = response.json()
|
||
|
|
except ValueError as e:
|
||
|
|
raise_error(
|
||
|
|
f"Failed to parse JSON response: {e}, content: {response.content!r}"
|
||
|
|
)
|
||
|
|
|
||
|
|
if json.get("errors"):
|
||
|
|
raise_error(json["errors"])
|
||
|
|
|
||
|
|
if json.get("data"):
|
||
|
|
if isinstance(json["data"], dict):
|
||
|
|
for value in json["data"].values():
|
||
|
|
if value and value.get("ok") is False:
|
||
|
|
raise_error(
|
||
|
|
f"""Failed to {description}: {value.get("message")}"""
|
||
|
|
)
|
||
|
|
|
||
|
|
return json
|
||
|
|
|
||
|
|
def make_rest_call(self, subpath: str, body: Dict[str, Any]) -> Dict:
|
||
|
|
with httpx.Client(follow_redirects=True) as client:
|
||
|
|
response = client.post(
|
||
|
|
self.rest_endpoint + subpath,
|
||
|
|
json=body,
|
||
|
|
headers=self.headers,
|
||
|
|
timeout=20,
|
||
|
|
)
|
||
|
|
|
||
|
|
try:
|
||
|
|
response.raise_for_status()
|
||
|
|
except httpx.HTTPStatusError:
|
||
|
|
message = f"Failed to call {subpath}: {response.text}"
|
||
|
|
logger.error(message)
|
||
|
|
raise Exception(message)
|
||
|
|
|
||
|
|
try:
|
||
|
|
return response.json()
|
||
|
|
except ValueError as e:
|
||
|
|
raise ValueError(
|
||
|
|
f"Failed to parse JSON response: {e}, content: {response.content!r}"
|
||
|
|
)
|
||
|
|
|
||
|
|
def gql_helper(
|
||
|
|
self,
|
||
|
|
query: str,
|
||
|
|
description: str,
|
||
|
|
variables: Dict,
|
||
|
|
process_response: Callable[..., R],
|
||
|
|
timeout: Optional[int] = None,
|
||
|
|
) -> R:
|
||
|
|
response = self.make_gql_call(description, query, variables, timeout)
|
||
|
|
return process_response(response)
|
||
|
|
|
||
|
|
##################################################################################
|
||
|
|
# User APIs #
|
||
|
|
##################################################################################
|
||
|
|
|
||
|
|
def get_users(
|
||
|
|
self,
|
||
|
|
first: Optional[int] = None,
|
||
|
|
after: Optional[str] = None,
|
||
|
|
before: Optional[str] = None,
|
||
|
|
filters: Optional[users_filters] = None,
|
||
|
|
) -> PaginatedResponse["User"]:
|
||
|
|
return self.gql_helper(*get_users_helper(first, after, before, filters))
|
||
|
|
|
||
|
|
def get_user(
|
||
|
|
self, id: Optional[str] = None, identifier: Optional[str] = None
|
||
|
|
) -> "User":
|
||
|
|
return self.gql_helper(*get_user_helper(id, identifier))
|
||
|
|
|
||
|
|
def create_user(self, identifier: str, metadata: Optional[Dict] = None) -> "User":
|
||
|
|
return self.gql_helper(*create_user_helper(identifier, metadata))
|
||
|
|
|
||
|
|
def update_user(
|
||
|
|
self, id: str, identifier: Optional[str] = None, metadata: Optional[Dict] = None
|
||
|
|
) -> "User":
|
||
|
|
return self.gql_helper(*update_user_helper(id, identifier, metadata))
|
||
|
|
|
||
|
|
def delete_user(self, id: str) -> Dict:
|
||
|
|
return self.gql_helper(*delete_user_helper(id))
|
||
|
|
|
||
|
|
def get_or_create_user(
|
||
|
|
self, identifier: str, metadata: Optional[Dict] = None
|
||
|
|
) -> "User":
|
||
|
|
user = self.get_user(identifier=identifier)
|
||
|
|
if user:
|
||
|
|
return user
|
||
|
|
|
||
|
|
return self.create_user(identifier, metadata)
|
||
|
|
|
||
|
|
##################################################################################
|
||
|
|
# Thread APIs #
|
||
|
|
##################################################################################
|
||
|
|
|
||
|
|
def get_threads(
|
||
|
|
self,
|
||
|
|
first: Optional[int] = None,
|
||
|
|
after: Optional[str] = None,
|
||
|
|
before: Optional[str] = None,
|
||
|
|
filters: Optional[threads_filters] = None,
|
||
|
|
order_by: Optional[threads_order_by] = None,
|
||
|
|
step_types_to_keep: Optional[List[StepType]] = None,
|
||
|
|
) -> PaginatedResponse["Thread"]:
|
||
|
|
return self.gql_helper(
|
||
|
|
*get_threads_helper(
|
||
|
|
first, after, before, filters, order_by, step_types_to_keep
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
def list_threads(
|
||
|
|
self,
|
||
|
|
first: Optional[int] = None,
|
||
|
|
after: Optional[str] = None,
|
||
|
|
before: Optional[str] = None,
|
||
|
|
filters: Optional[threads_filters] = None,
|
||
|
|
order_by: Optional[threads_order_by] = None,
|
||
|
|
) -> PaginatedResponse["Thread"]:
|
||
|
|
return self.gql_helper(
|
||
|
|
*list_threads_helper(first, after, before, filters, order_by)
|
||
|
|
)
|
||
|
|
|
||
|
|
def get_thread(self, id: str) -> "Thread":
|
||
|
|
return self.gql_helper(*get_thread_helper(id))
|
||
|
|
|
||
|
|
def create_thread(
|
||
|
|
self,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
metadata: Optional[Dict] = None,
|
||
|
|
participant_id: Optional[str] = None,
|
||
|
|
tags: Optional[List[str]] = None,
|
||
|
|
) -> "Thread":
|
||
|
|
return self.gql_helper(
|
||
|
|
*create_thread_helper(name, metadata, participant_id, tags)
|
||
|
|
)
|
||
|
|
|
||
|
|
def upsert_thread(
|
||
|
|
self,
|
||
|
|
id: str,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
metadata: Optional[Dict] = None,
|
||
|
|
participant_id: Optional[str] = None,
|
||
|
|
tags: Optional[List[str]] = None,
|
||
|
|
) -> "Thread":
|
||
|
|
return self.gql_helper(
|
||
|
|
*upsert_thread_helper(id, name, metadata, participant_id, tags)
|
||
|
|
)
|
||
|
|
|
||
|
|
def update_thread(
|
||
|
|
self,
|
||
|
|
id: str,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
metadata: Optional[Dict] = None,
|
||
|
|
participant_id: Optional[str] = None,
|
||
|
|
tags: Optional[List[str]] = None,
|
||
|
|
) -> "Thread":
|
||
|
|
return self.gql_helper(
|
||
|
|
*update_thread_helper(id, name, metadata, participant_id, tags)
|
||
|
|
)
|
||
|
|
|
||
|
|
def delete_thread(self, id: str) -> bool:
|
||
|
|
return self.gql_helper(*delete_thread_helper(id))
|
||
|
|
|
||
|
|
##################################################################################
|
||
|
|
# Score APIs #
|
||
|
|
##################################################################################
|
||
|
|
|
||
|
|
def get_scores(
|
||
|
|
self,
|
||
|
|
first: Optional[int] = None,
|
||
|
|
after: Optional[str] = None,
|
||
|
|
before: Optional[str] = None,
|
||
|
|
filters: Optional[scores_filters] = None,
|
||
|
|
order_by: Optional[scores_order_by] = None,
|
||
|
|
) -> PaginatedResponse["Score"]:
|
||
|
|
return self.gql_helper(
|
||
|
|
*get_scores_helper(first, after, before, filters, order_by)
|
||
|
|
)
|
||
|
|
|
||
|
|
def create_scores(self, scores: List["ScoreDict"]):
|
||
|
|
check_scores_finite(scores)
|
||
|
|
|
||
|
|
query = create_scores_query_builder(scores)
|
||
|
|
variables = {}
|
||
|
|
for id, score in enumerate(scores):
|
||
|
|
for k, v in score.items():
|
||
|
|
variables[f"{k}_{id}"] = v
|
||
|
|
|
||
|
|
def process_response(response):
|
||
|
|
return [x for x in response["data"].values()]
|
||
|
|
|
||
|
|
return self.gql_helper(query, "create scores", variables, process_response)
|
||
|
|
|
||
|
|
def create_score(
|
||
|
|
self,
|
||
|
|
name: str,
|
||
|
|
value: float,
|
||
|
|
type: ScoreType,
|
||
|
|
step_id: Optional[str] = None,
|
||
|
|
generation_id: Optional[str] = None,
|
||
|
|
dataset_experiment_item_id: Optional[str] = None,
|
||
|
|
comment: Optional[str] = None,
|
||
|
|
tags: Optional[List[str]] = None,
|
||
|
|
) -> "Score":
|
||
|
|
if generation_id:
|
||
|
|
logger.warning(
|
||
|
|
"generation_id is deprecated and will be removed in a future version, please use step_id instead"
|
||
|
|
)
|
||
|
|
check_scores_finite([{"name": name, "value": value}])
|
||
|
|
|
||
|
|
return self.gql_helper(
|
||
|
|
*create_score_helper(
|
||
|
|
name,
|
||
|
|
value,
|
||
|
|
type,
|
||
|
|
step_id,
|
||
|
|
dataset_experiment_item_id,
|
||
|
|
comment,
|
||
|
|
tags,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
def update_score(
|
||
|
|
self,
|
||
|
|
id: str,
|
||
|
|
update_params: ScoreUpdate,
|
||
|
|
) -> "Score":
|
||
|
|
return self.gql_helper(*update_score_helper(id, update_params))
|
||
|
|
|
||
|
|
def delete_score(self, id: str) -> Dict:
|
||
|
|
return self.gql_helper(*delete_score_helper(id))
|
||
|
|
|
||
|
|
##################################################################################
|
||
|
|
# Attachment APIs #
|
||
|
|
##################################################################################
|
||
|
|
|
||
|
|
def upload_file(
|
||
|
|
self,
|
||
|
|
content: Union[bytes, str],
|
||
|
|
thread_id: Optional[str] = None,
|
||
|
|
mime: Optional[str] = "application/octet-stream",
|
||
|
|
) -> Dict:
|
||
|
|
id = str(uuid.uuid4())
|
||
|
|
body = {"fileName": id, "contentType": mime}
|
||
|
|
if thread_id:
|
||
|
|
body["threadId"] = thread_id
|
||
|
|
|
||
|
|
path = "/api/upload/file"
|
||
|
|
|
||
|
|
with httpx.Client(follow_redirects=True) as client:
|
||
|
|
response = client.post(
|
||
|
|
f"{self.url}{path}",
|
||
|
|
json=body,
|
||
|
|
headers=self.headers,
|
||
|
|
)
|
||
|
|
if response.status_code >= 400:
|
||
|
|
reason = response.text
|
||
|
|
logger.error(f"Failed to sign upload url: {reason}")
|
||
|
|
return {"object_key": None, "url": None}
|
||
|
|
json_res = response.json()
|
||
|
|
method = "put" if "put" in json_res else "post"
|
||
|
|
request_dict: Dict[str, Any] = json_res.get(method, {})
|
||
|
|
url: Optional[str] = request_dict.get("url")
|
||
|
|
|
||
|
|
if not url:
|
||
|
|
raise Exception("Invalid server response")
|
||
|
|
headers: Optional[Dict] = request_dict.get("headers")
|
||
|
|
fields: Dict = request_dict.get("fields", {})
|
||
|
|
object_key: Optional[str] = fields.get("key")
|
||
|
|
upload_type: Literal["raw", "multipart"] = cast(
|
||
|
|
Literal["raw", "multipart"], request_dict.get("uploadType", "multipart")
|
||
|
|
)
|
||
|
|
signed_url: Optional[str] = json_res.get("signedUrl")
|
||
|
|
|
||
|
|
# Prepare form data
|
||
|
|
form_data = (
|
||
|
|
{}
|
||
|
|
) # type: Dict[str, Union[tuple[Union[str, None], Any], tuple[Union[str, None], Any, Any]]]
|
||
|
|
for field_name, field_value in fields.items():
|
||
|
|
form_data[field_name] = (None, field_value)
|
||
|
|
|
||
|
|
# Add file to the form_data
|
||
|
|
# Note: The content_type parameter is not needed here, as the correct MIME type should be set
|
||
|
|
# in the 'Content-Type' field from upload_details
|
||
|
|
form_data["file"] = (id, content, mime)
|
||
|
|
|
||
|
|
with httpx.Client(follow_redirects=True) as client:
|
||
|
|
if upload_type == "raw":
|
||
|
|
upload_response = client.request(
|
||
|
|
url=url,
|
||
|
|
headers=headers,
|
||
|
|
method=method,
|
||
|
|
data=content, # type: ignore
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
upload_response = client.request(
|
||
|
|
url=url,
|
||
|
|
headers=headers,
|
||
|
|
method=method,
|
||
|
|
files=form_data,
|
||
|
|
) # type: ignore
|
||
|
|
try:
|
||
|
|
upload_response.raise_for_status()
|
||
|
|
return {"object_key": object_key, "url": signed_url}
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Failed to upload file: {str(e)}")
|
||
|
|
return {"object_key": None, "url": None}
|
||
|
|
|
||
|
|
def create_attachment(
|
||
|
|
self,
|
||
|
|
thread_id: Optional[str] = None,
|
||
|
|
step_id: Optional[str] = None,
|
||
|
|
id: Optional[str] = None,
|
||
|
|
metadata: Optional[Dict] = None,
|
||
|
|
mime: Optional[str] = None,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
object_key: Optional[str] = None,
|
||
|
|
url: Optional[str] = None,
|
||
|
|
content: Optional[Union[bytes, str]] = None,
|
||
|
|
path: Optional[str] = None,
|
||
|
|
) -> "Attachment":
|
||
|
|
if not thread_id:
|
||
|
|
if active_thread := active_thread_var.get(None):
|
||
|
|
thread_id = active_thread.id
|
||
|
|
|
||
|
|
if not step_id:
|
||
|
|
if active_steps := active_steps_var.get():
|
||
|
|
step_id = active_steps[-1].id
|
||
|
|
else:
|
||
|
|
raise Exception("No step_id provided and no active step found.")
|
||
|
|
|
||
|
|
(
|
||
|
|
query,
|
||
|
|
description,
|
||
|
|
variables,
|
||
|
|
content,
|
||
|
|
process_response,
|
||
|
|
) = create_attachment_helper(
|
||
|
|
thread_id=thread_id,
|
||
|
|
step_id=step_id,
|
||
|
|
id=id,
|
||
|
|
metadata=metadata,
|
||
|
|
mime=mime,
|
||
|
|
name=name,
|
||
|
|
object_key=object_key,
|
||
|
|
url=url,
|
||
|
|
content=content,
|
||
|
|
path=path,
|
||
|
|
)
|
||
|
|
|
||
|
|
if content:
|
||
|
|
uploaded = self.upload_file(content=content, thread_id=thread_id, mime=mime)
|
||
|
|
|
||
|
|
if uploaded["object_key"] is None or uploaded["url"] is None:
|
||
|
|
raise Exception("Failed to upload file")
|
||
|
|
|
||
|
|
object_key = uploaded["object_key"]
|
||
|
|
if object_key:
|
||
|
|
variables["objectKey"] = object_key
|
||
|
|
else:
|
||
|
|
variables["url"] = uploaded["url"]
|
||
|
|
|
||
|
|
response = self.make_gql_call(description, query, variables)
|
||
|
|
return process_response(response)
|
||
|
|
|
||
|
|
def update_attachment(
|
||
|
|
self, id: str, update_params: AttachmentUpload
|
||
|
|
) -> "Attachment":
|
||
|
|
return self.gql_helper(*update_attachment_helper(id, update_params))
|
||
|
|
|
||
|
|
def get_attachment(self, id: str) -> Optional["Attachment"]:
|
||
|
|
return self.gql_helper(*get_attachment_helper(id))
|
||
|
|
|
||
|
|
def delete_attachment(self, id: str) -> Dict:
|
||
|
|
return self.gql_helper(*delete_attachment_helper(id))
|
||
|
|
|
||
|
|
##################################################################################
|
||
|
|
# Step APIs #
|
||
|
|
##################################################################################
|
||
|
|
|
||
|
|
def create_step(
|
||
|
|
self,
|
||
|
|
thread_id: Optional[str] = None,
|
||
|
|
type: Optional[StepType] = "undefined",
|
||
|
|
start_time: Optional[str] = None,
|
||
|
|
end_time: Optional[str] = None,
|
||
|
|
input: Optional[Dict] = None,
|
||
|
|
output: Optional[Dict] = None,
|
||
|
|
metadata: Optional[Dict] = None,
|
||
|
|
parent_id: Optional[str] = None,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
tags: Optional[List[str]] = None,
|
||
|
|
root_run_id: Optional[str] = None,
|
||
|
|
) -> Step:
|
||
|
|
return self.gql_helper(
|
||
|
|
*create_step_helper(
|
||
|
|
thread_id=thread_id,
|
||
|
|
type=type,
|
||
|
|
start_time=start_time,
|
||
|
|
end_time=end_time,
|
||
|
|
input=input,
|
||
|
|
output=output,
|
||
|
|
metadata=metadata,
|
||
|
|
parent_id=parent_id,
|
||
|
|
name=name,
|
||
|
|
tags=tags,
|
||
|
|
root_run_id=root_run_id,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
def update_step(
|
||
|
|
self,
|
||
|
|
id: str,
|
||
|
|
type: Optional[StepType] = None,
|
||
|
|
input: Optional[str] = None,
|
||
|
|
output: Optional[str] = None,
|
||
|
|
metadata: Optional[Dict] = None,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
tags: Optional[List[str]] = None,
|
||
|
|
start_time: Optional[str] = None,
|
||
|
|
end_time: Optional[str] = None,
|
||
|
|
parent_id: Optional[str] = None,
|
||
|
|
) -> "Step":
|
||
|
|
return self.gql_helper(
|
||
|
|
*update_step_helper(
|
||
|
|
id=id,
|
||
|
|
type=type,
|
||
|
|
input=input,
|
||
|
|
output=output,
|
||
|
|
metadata=metadata,
|
||
|
|
name=name,
|
||
|
|
tags=tags,
|
||
|
|
start_time=start_time,
|
||
|
|
end_time=end_time,
|
||
|
|
parent_id=parent_id,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
def get_steps(
|
||
|
|
self,
|
||
|
|
first: Optional[int] = None,
|
||
|
|
after: Optional[str] = None,
|
||
|
|
before: Optional[str] = None,
|
||
|
|
filters: Optional[steps_filters] = None,
|
||
|
|
order_by: Optional[steps_order_by] = None,
|
||
|
|
) -> PaginatedResponse["Step"]:
|
||
|
|
return self.gql_helper(
|
||
|
|
*get_steps_helper(first, after, before, filters, order_by)
|
||
|
|
)
|
||
|
|
|
||
|
|
def get_step(
|
||
|
|
self,
|
||
|
|
id: str,
|
||
|
|
) -> Optional["Step"]:
|
||
|
|
return self.gql_helper(*get_step_helper(id=id))
|
||
|
|
|
||
|
|
def delete_step(
|
||
|
|
self,
|
||
|
|
id: str,
|
||
|
|
) -> bool:
|
||
|
|
return self.gql_helper(*delete_step_helper(id=id))
|
||
|
|
|
||
|
|
def send_steps(self, steps: List[Union["StepDict", "Step"]]):
|
||
|
|
return self.gql_helper(*send_steps_helper(steps=steps))
|
||
|
|
|
||
|
|
##################################################################################
|
||
|
|
# Generation APIs #
|
||
|
|
##################################################################################
|
||
|
|
|
||
|
|
def get_generations(
|
||
|
|
self,
|
||
|
|
first: Optional[int] = None,
|
||
|
|
after: Optional[str] = None,
|
||
|
|
before: Optional[str] = None,
|
||
|
|
filters: Optional[generations_filters] = None,
|
||
|
|
order_by: Optional[generations_order_by] = None,
|
||
|
|
) -> PaginatedResponse["BaseGeneration"]:
|
||
|
|
return self.gql_helper(
|
||
|
|
*get_generations_helper(first, after, before, filters, order_by)
|
||
|
|
)
|
||
|
|
|
||
|
|
def create_generation(
|
||
|
|
self, generation: Union["ChatGeneration", "CompletionGeneration"]
|
||
|
|
) -> Union["ChatGeneration", "CompletionGeneration"]:
|
||
|
|
return self.gql_helper(*create_generation_helper(generation))
|
||
|
|
|
||
|
|
##################################################################################
|
||
|
|
# Dataset APIs #
|
||
|
|
##################################################################################
|
||
|
|
|
||
|
|
def create_dataset(
|
||
|
|
self,
|
||
|
|
name: str,
|
||
|
|
description: Optional[str] = None,
|
||
|
|
metadata: Optional[Dict] = None,
|
||
|
|
type: DatasetType = "key_value",
|
||
|
|
) -> "Dataset":
|
||
|
|
return self.gql_helper(
|
||
|
|
*create_dataset_helper(self, name, description, metadata, type)
|
||
|
|
)
|
||
|
|
|
||
|
|
def get_dataset(
|
||
|
|
self, id: Optional[str] = None, name: Optional[str] = None
|
||
|
|
) -> Optional["Dataset"]:
|
||
|
|
subpath, _, variables, process_response = get_dataset_helper(
|
||
|
|
self, id=id, name=name
|
||
|
|
)
|
||
|
|
response = self.make_rest_call(subpath, variables)
|
||
|
|
return process_response(response)
|
||
|
|
|
||
|
|
def update_dataset(
|
||
|
|
self,
|
||
|
|
id: str,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
description: Optional[str] = None,
|
||
|
|
metadata: Optional[Dict] = None,
|
||
|
|
) -> "Dataset":
|
||
|
|
return self.gql_helper(
|
||
|
|
*update_dataset_helper(self, id, name, description, metadata)
|
||
|
|
)
|
||
|
|
|
||
|
|
def delete_dataset(self, id: str) -> "Dataset":
|
||
|
|
return self.gql_helper(*delete_dataset_helper(self, id))
|
||
|
|
|
||
|
|
##################################################################################
|
||
|
|
# Experiment APIs #
|
||
|
|
##################################################################################
|
||
|
|
|
||
|
|
def create_experiment(
|
||
|
|
self,
|
||
|
|
name: str,
|
||
|
|
dataset_id: Optional[str] = None,
|
||
|
|
prompt_variant_id: Optional[str] = None,
|
||
|
|
params: Optional[Dict] = None,
|
||
|
|
) -> "DatasetExperiment":
|
||
|
|
return self.gql_helper(
|
||
|
|
*create_experiment_helper(
|
||
|
|
api=self,
|
||
|
|
name=name,
|
||
|
|
dataset_id=dataset_id,
|
||
|
|
prompt_variant_id=prompt_variant_id,
|
||
|
|
params=params,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
def create_experiment_item(
|
||
|
|
self, experiment_item: DatasetExperimentItem
|
||
|
|
) -> "DatasetExperimentItem":
|
||
|
|
# Create the dataset experiment item
|
||
|
|
result = self.gql_helper(
|
||
|
|
*create_experiment_item_helper(
|
||
|
|
dataset_experiment_id=experiment_item.dataset_experiment_id,
|
||
|
|
dataset_item_id=experiment_item.dataset_item_id,
|
||
|
|
experiment_run_id=experiment_item.experiment_run_id,
|
||
|
|
input=experiment_item.input,
|
||
|
|
output=experiment_item.output,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
for score in experiment_item.scores:
|
||
|
|
score["datasetExperimentItemId"] = result.id
|
||
|
|
|
||
|
|
# Create the scores and add to experiment item.
|
||
|
|
result.scores = self.create_scores(experiment_item.scores)
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
##################################################################################
|
||
|
|
# Dataset Item APIs #
|
||
|
|
##################################################################################
|
||
|
|
|
||
|
|
def create_dataset_item(
|
||
|
|
self,
|
||
|
|
dataset_id: str,
|
||
|
|
input: Dict,
|
||
|
|
expected_output: Optional[Dict] = None,
|
||
|
|
metadata: Optional[Dict] = None,
|
||
|
|
) -> "DatasetItem":
|
||
|
|
return self.gql_helper(
|
||
|
|
*create_dataset_item_helper(dataset_id, input, expected_output, metadata)
|
||
|
|
)
|
||
|
|
|
||
|
|
def get_dataset_item(self, id: str) -> Optional["DatasetItem"]:
|
||
|
|
return self.gql_helper(*get_dataset_item_helper(id))
|
||
|
|
|
||
|
|
def delete_dataset_item(self, id: str) -> "DatasetItem":
|
||
|
|
return self.gql_helper(*delete_dataset_item_helper(id))
|
||
|
|
|
||
|
|
def add_step_to_dataset(
|
||
|
|
self, dataset_id: str, step_id: str, metadata: Optional[Dict] = None
|
||
|
|
) -> "DatasetItem":
|
||
|
|
return self.gql_helper(
|
||
|
|
*add_step_to_dataset_helper(dataset_id, step_id, metadata)
|
||
|
|
)
|
||
|
|
|
||
|
|
def add_generation_to_dataset(
|
||
|
|
self, dataset_id: str, generation_id: str, metadata: Optional[Dict] = None
|
||
|
|
) -> "DatasetItem":
|
||
|
|
return self.gql_helper(
|
||
|
|
*add_generation_to_dataset_helper(dataset_id, generation_id, metadata)
|
||
|
|
)
|
||
|
|
|
||
|
|
##################################################################################
|
||
|
|
# Prompt APIs #
|
||
|
|
##################################################################################
|
||
|
|
|
||
|
|
def get_or_create_prompt_lineage(
|
||
|
|
self, name: str, description: Optional[str] = None
|
||
|
|
) -> Dict:
|
||
|
|
return self.gql_helper(*create_prompt_lineage_helper(name, description))
|
||
|
|
|
||
|
|
@deprecated("Use get_or_create_prompt_lineage instead")
|
||
|
|
def create_prompt_lineage(
|
||
|
|
self, name: str, description: Optional[str] = None
|
||
|
|
) -> Dict:
|
||
|
|
return self.get_or_create_prompt_lineage(name, description)
|
||
|
|
|
||
|
|
def get_or_create_prompt(
|
||
|
|
self,
|
||
|
|
name: str,
|
||
|
|
template_messages: List[GenerationMessage],
|
||
|
|
settings: Optional[ProviderSettings] = None,
|
||
|
|
tools: Optional[List[Dict]] = None,
|
||
|
|
) -> "Prompt":
|
||
|
|
lineage = self.get_or_create_prompt_lineage(name)
|
||
|
|
lineage_id = lineage["id"]
|
||
|
|
return self.gql_helper(
|
||
|
|
*create_prompt_helper(self, lineage_id, template_messages, settings, tools)
|
||
|
|
)
|
||
|
|
|
||
|
|
@deprecated("Please use `get_or_create_prompt` instead.")
|
||
|
|
def create_prompt(
|
||
|
|
self,
|
||
|
|
name: str,
|
||
|
|
template_messages: List[GenerationMessage],
|
||
|
|
settings: Optional[ProviderSettings] = None,
|
||
|
|
) -> "Prompt":
|
||
|
|
return self.get_or_create_prompt(name, template_messages, settings)
|
||
|
|
|
||
|
|
def get_prompt(
|
||
|
|
self,
|
||
|
|
id: Optional[str] = None,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
version: Optional[int] = None,
|
||
|
|
) -> "Prompt":
|
||
|
|
if not (id or name):
|
||
|
|
raise ValueError("At least the `id` or the `name` must be provided.")
|
||
|
|
|
||
|
|
(
|
||
|
|
get_prompt_query,
|
||
|
|
description,
|
||
|
|
variables,
|
||
|
|
process_response,
|
||
|
|
timeout,
|
||
|
|
cached_prompt,
|
||
|
|
) = get_prompt_helper(
|
||
|
|
api=self, id=id, name=name, version=version, cache=self.cache
|
||
|
|
)
|
||
|
|
|
||
|
|
try:
|
||
|
|
if id:
|
||
|
|
prompt = self.gql_helper(
|
||
|
|
get_prompt_query, description, variables, process_response, timeout
|
||
|
|
)
|
||
|
|
elif name:
|
||
|
|
prompt = self.gql_helper(
|
||
|
|
get_prompt_query, description, variables, process_response, timeout
|
||
|
|
)
|
||
|
|
|
||
|
|
return prompt
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
if cached_prompt:
|
||
|
|
logger.warning("Failed to get prompt from API, returning cached prompt")
|
||
|
|
return cached_prompt
|
||
|
|
|
||
|
|
raise e
|
||
|
|
|
||
|
|
def create_prompt_variant(
|
||
|
|
self,
|
||
|
|
name: str,
|
||
|
|
template_messages: List[GenerationMessage],
|
||
|
|
settings: Optional[ProviderSettings] = None,
|
||
|
|
tools: Optional[List[Dict]] = None,
|
||
|
|
) -> Optional[str]:
|
||
|
|
lineage = self.gql_helper(*get_prompt_lineage_helper(name))
|
||
|
|
lineage_id = lineage["id"] if lineage else None
|
||
|
|
return self.gql_helper(
|
||
|
|
*create_prompt_variant_helper(
|
||
|
|
lineage_id, template_messages, settings, tools
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
def get_prompt_ab_testing(self, name: str) -> List["PromptRollout"]:
|
||
|
|
return self.gql_helper(*get_prompt_ab_testing_helper(name=name))
|
||
|
|
|
||
|
|
def update_prompt_ab_testing(
|
||
|
|
self, name: str, rollouts: List["PromptRollout"]
|
||
|
|
) -> Dict:
|
||
|
|
return self.gql_helper(
|
||
|
|
*update_prompt_ab_testing_helper(name=name, rollouts=rollouts)
|
||
|
|
)
|
||
|
|
|
||
|
|
##################################################################################
|
||
|
|
# Misc APIs #
|
||
|
|
##################################################################################
|
||
|
|
|
||
|
|
def get_my_project_id(self) -> str:
|
||
|
|
response = self.make_rest_call("/my-project", {})
|
||
|
|
return response["projectId"]
|