Skip to content

trulens.providers.cortex.endpoint

trulens.providers.cortex.endpoint

Classes

CortexCallback

Bases: EndpointCallback

Attributes
endpoint class-attribute instance-attribute
endpoint: Endpoint = Field(exclude=True)

The endpoint owning this callback.

cost class-attribute instance-attribute
cost: Cost = Field(default_factory=Cost)

Costs tracked by this callback.

Functions
__rich_repr__
__rich_repr__() -> Result

Requirement for pretty printing using the rich package.

handle
handle(response: Any) -> None

Called after each request.

handle_chunk
handle_chunk(response: Any) -> None

Called after receiving a chunk from a request.

handle_generation_chunk
handle_generation_chunk(response: Any) -> None

Called after receiving a chunk from a completion request.

handle_classification
handle_classification(response: Any) -> None

Called after each classification response.

handle_generation
handle_generation(response: dict) -> None

Get the usage information from Cortex LLM function response's usage field.

CortexEndpoint

Bases: Endpoint

Snowflake Cortex endpoint.

Attributes
tru_class_info instance-attribute
tru_class_info: Class

Class information of this pydantic object for use in deserialization.

Using this odd key to not pollute attribute names in whatever class we mix this into. Should be the same as CLASS_INFO.

instrumented_methods class-attribute
instrumented_methods: Dict[
    Any, List[Tuple[Callable, Callable, Type[Endpoint]]]
] = defaultdict(list)

Mapping of classes/module-methods that have been instrumented for cost tracking along with the wrapper methods and the class that instrumented them.

Key is the class or module owning the instrumented method. Tuple value has:

  • original function,

  • wrapped version,

  • endpoint that did the wrapping.

name instance-attribute
name: str

API/endpoint name.

rpm class-attribute instance-attribute

Requests per minute.

retries class-attribute instance-attribute
retries: int = 3

Retries (if performing requests using this class).

post_headers class-attribute instance-attribute
post_headers: Dict[str, str] = Field(
    default_factory=dict, exclude=True
)

Optional post headers for post requests if done by this class.

pace class-attribute instance-attribute
pace: Pace = Field(
    default_factory=lambda: Pace(
        marks_per_second=DEFAULT_RPM / 60.0,
        seconds_per_period=60.0,
    ),
    exclude=True,
)

Pacing instance to maintain a desired rpm.

global_callback class-attribute instance-attribute
global_callback: EndpointCallback = Field(exclude=True)

Track costs not run inside "track_cost" here.

Also note that Endpoints are singletons (one for each unique name argument) hence this global callback will track all requests for the named api even if you try to create multiple endpoints (with the same name).

callback_class class-attribute instance-attribute
callback_class: Type[EndpointCallback] = Field(exclude=True)

Callback class to use for usage tracking.

callback_name class-attribute instance-attribute
callback_name: str = Field(exclude=True)

Name of variable that stores the callback noted above.

Classes
EndpointSetup dataclass

Class for storing supported endpoint information.

See track_all_costs for usage.

Functions
warning
warning()

Issue warning that this singleton already exists.

delete_singleton_by_name staticmethod
delete_singleton_by_name(
    name: str, cls: Optional[Type[SingletonPerName]] = None
)

Delete the singleton instance with the given name.

This can be used for testing to create another singleton.

PARAMETER DESCRIPTION
name

The name of the singleton instance to delete.

TYPE: str

cls

The class of the singleton instance to delete. If not given, all instances with the given name are deleted.

TYPE: Optional[Type[SingletonPerName]] DEFAULT: None

delete_singleton
delete_singleton()

Delete the singleton instance. Can be used for testing to create another singleton.

__rich_repr__
__rich_repr__() -> Result

Requirement for pretty printing using the rich package.

load staticmethod
load(obj, *args, **kwargs)

Deserialize/load this object using the class information in tru_class_info to lookup the actual class that will do the deserialization.

model_validate classmethod
model_validate(*args, **kwargs) -> Any

Deserialized a jsonized version of the app into the instance of the class it was serialized from.

Note

This process uses extra information stored in the jsonized object and handled by WithClassInfo.

pace_me
pace_me() -> float

Block until we can make a request to this endpoint to keep pace with maximum rpm. Returns time in seconds since last call to this method returned.

run_in_pace
run_in_pace(
    func: Callable[[A], B], *args, **kwargs
) -> B

Run the given func on the given args and kwargs at pace with the endpoint-specified rpm. Failures will be retried self.retries times.

run_me
run_me(thunk: Thunk[T]) -> T

DEPRECATED: Run the given thunk, returning itse output, on pace with the api. Retries request multiple times if self.retries > 0.

DEPRECATED: Use run_in_pace instead.

print_instrumented classmethod
print_instrumented()

Print out all of the methods that have been instrumented for cost tracking. This is organized by the classes/modules containing them.

track_all_costs staticmethod
track_all_costs(
    __func: CallableMaybeAwaitable[A, T],
    *args,
    with_openai: bool = True,
    with_hugs: bool = True,
    with_litellm: bool = True,
    with_bedrock: bool = True,
    with_cortex: bool = True,
    **kwargs
) -> Tuple[T, Sequence[EndpointCallback]]

Track costs of all of the apis we can currently track, over the execution of thunk.

track_all_costs_tally staticmethod
track_all_costs_tally(
    __func: CallableMaybeAwaitable[A, T],
    *args,
    with_openai: bool = True,
    with_hugs: bool = True,
    with_litellm: bool = True,
    with_bedrock: bool = True,
    with_cortex: bool = True,
    **kwargs
) -> Tuple[T, Cost]

Track costs of all of the apis we can currently track, over the execution of thunk.

track_cost
track_cost(
    __func: CallableMaybeAwaitable[T], *args, **kwargs
) -> Tuple[T, EndpointCallback]

Tally only the usage performed within the execution of the given thunk. Returns the thunk's result alongside the EndpointCallback object that includes the usage information.

wrap_function
wrap_function(func)

Create a wrapper of the given function to perform cost tracking.