tensorzero 0.0.1__tar.gz

tensorzero-0.0.1/.gitignore

@@ -0,0 +1,27 @@
+# Generated by Cargo
+# will have compiled files and executables
+debug/
+target/
+
+# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
+# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
+Cargo.lock
+
+# These are backup files generated by rustfmt
+**/*.rs.bk
+
+# MSVC Windows builds of rustc generate these, which store debugging information
+*.pdb
+.envrc
+.python-version
+
+# VScode settings
+.vscode/
+
+.DS_Store
+.credentials/
+
+# Python cache files
+__pycache__/
+*.py[cod]
+*$py.class

tensorzero-0.0.1/PKG-INFO

@@ -0,0 +1,56 @@
+Metadata-Version: 2.3
+Name: tensorzero
+Version: 0.0.1
+Summary: The Python client for TensorZero
+Author-email: Viraj Mehta <viraj@tensorzero.com>, Gabriel Bianconi <gabriel@tensorzero.com>
+Requires-Python: >=3.10
+Requires-Dist: httpx>=0.27.0
+Description-Content-Type: text/markdown
+
+# TensorZero Python Client
+
+This is an async Python client for the TensorZero gateway. Check out the [docs](https://tensorzero.com/docs/) for more information. This client allows you to easily make inference requests and assign feedback to them via the TensorZero gateway.
+
+## Installation
+
+```bash
+pip install tensorzero
+```
+
+## Basic Usage
+
+### Non-Streaming Inference
+
+```python
+from tensorzero import AsyncTensorZero
+
+with AsyncTensorZero("http://localhost:3000") as client:
+    result = await client.inference(
+        function_name="basic_test",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [{"role": "user", "content": "Hello"}],
+        },
+    )
+episode_id = result.episode_id
+output = result.output
+print(output[0].text)  # Prints the text of the first content block returned by TensorZero
+```
+
+### Streaming Inference
+
+```python
+from tensorzero import AsyncTensorZero
+
+async with AsyncTensorZero() as client:
+    stream = await client.chat.completions.create(
+        function_name="basic_test",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [{"role": "user", "content": "Hello"}],
+        },
+        stream=True,
+    )
+    async for chunk in stream:
+        print(chunk.content[0].text)  # Prints the text in each chunk returned by TensorZero
+```

tensorzero-0.0.1/README.md

@@ -0,0 +1,47 @@
+# TensorZero Python Client
+
+This is an async Python client for the TensorZero gateway. Check out the [docs](https://tensorzero.com/docs/) for more information. This client allows you to easily make inference requests and assign feedback to them via the TensorZero gateway.
+
+## Installation
+
+```bash
+pip install tensorzero
+```
+
+## Basic Usage
+
+### Non-Streaming Inference
+
+```python
+from tensorzero import AsyncTensorZero
+
+with AsyncTensorZero("http://localhost:3000") as client:
+    result = await client.inference(
+        function_name="basic_test",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [{"role": "user", "content": "Hello"}],
+        },
+    )
+episode_id = result.episode_id
+output = result.output
+print(output[0].text)  # Prints the text of the first content block returned by TensorZero
+```
+
+### Streaming Inference
+
+```python
+from tensorzero import AsyncTensorZero
+
+async with AsyncTensorZero() as client:
+    stream = await client.chat.completions.create(
+        function_name="basic_test",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [{"role": "user", "content": "Hello"}],
+        },
+        stream=True,
+    )
+    async for chunk in stream:
+        print(chunk.content[0].text)  # Prints the text in each chunk returned by TensorZero
+```

tensorzero-0.0.1/pyproject.toml

@@ -0,0 +1,18 @@
+[project]
+name = "tensorzero"
+version = "0.0.1"
+description = "The Python client for TensorZero"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = ["httpx>=0.27.0"]
+authors = [
+    { name = "Viraj Mehta", email = "viraj@tensorzero.com" },
+    { name = "Gabriel Bianconi", email = "gabriel@tensorzero.com" },
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.uv]
+dev-dependencies = ["uuid7>=0.1.0", "pytest>=8.3.2", "pytest-asyncio>=0.24.0"]

tensorzero-0.0.1/src/tensorzero/__init__.py

@@ -0,0 +1,29 @@
+from .client import AsyncTensorZero
+from .types import (
+    ChatInferenceResponse,
+    ContentBlock,
+    FeedbackResponse,
+    InferenceChunk,
+    InferenceResponse,
+    JsonInferenceOutput,
+    Text,
+    TextChunk,
+    ToolCall,
+    ToolCallChunk,
+    Usage,
+)
+
+__all__ = [
+    "ChatInferenceResponse",
+    "ContentBlock",
+    "FeedbackResponse",
+    "InferenceChunk",
+    "InferenceResponse",
+    "JsonInferenceOutput",
+    "AsyncTensorZero",
+    "Text",
+    "TextChunk",
+    "ToolCall",
+    "ToolCallChunk",
+    "Usage",
+]

tensorzero-0.0.1/src/tensorzero/client.py

@@ -0,0 +1,194 @@
+"""
+TensorZero Client
+
+This module provides an asynchronous client for interacting with the TensorZero gateway.
+It includes functionality for making inference requests and sending feedback.
+
+The main class, AsyncTensorZero, offers methods for:
+- Initializing the client with a base URL
+- Making inference requests (with optional streaming)
+- Sending feedback on episodes or inferences
+- Managing the client session using async context managers
+
+Usage:
+    async with TensorZero(base_url) as client:
+        response = await client.inference(...)
+        feedback = await client.feedback(...)
+"""
+
+import json
+import logging
+from typing import Any, AsyncGenerator, Dict, List, Literal, Optional, Union
+from urllib.parse import urljoin
+from uuid import UUID
+
+import httpx
+
+from .types import (
+    FeedbackResponse,
+    InferenceChunk,
+    InferenceResponse,
+    parse_inference_chunk,
+    parse_inference_response,
+)
+
+
+class AsyncTensorZero:
+    def __init__(self, base_url: str):
+        """
+        Initialize the TensorZero client.
+
+        :param base_url: The base URL of the TensorZero gateway. Example:"http://localhost:3000"
+        """
+        self.base_url = base_url
+        self.client = httpx.AsyncClient()
+        self.logger = logging.getLogger(__name__)
+
+    async def inference(
+        self,
+        function_name: str,
+        input: Dict[str, Any],
+        episode_id: Optional[UUID] = None,
+        stream: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        variant_name: Optional[str] = None,
+        dryrun: Optional[bool] = None,
+        allowed_tools: Optional[List[str]] = None,
+        additional_tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[
+            Union[Literal["auto", "required", "off"], Dict[Literal["specific"], str]]
+        ] = None,
+        parallel_tool_calls: Optional[bool] = None,
+    ) -> Union[InferenceResponse, AsyncGenerator[InferenceChunk, None]]:
+        """
+        Make a POST request to the /inference endpoint.
+
+        :param function_name: The name of the function to call
+        :param input: The input to the function
+                      Structure: {"system": Optional[str], "messages": List[{"role": "user" | "assistant", "content": Any}]}
+                      The input will be validated server side against the input schema of the function being called.
+        :param episode_id: The episode ID to use for the inference.
+                           If this is the first inference in an episode, leave this field blank. The TensorZero gateway will generate and return a new episode ID.
+                           Note: Only use episode IDs generated by the TensorZero gateway. Don't generate them yourself.
+        :param stream: If set, the TensorZero gateway will stream partial message deltas (e.g. generated tokens) as it receives them from model providers.
+        :param params: Override inference-time parameters for a particular variant type. Currently, we support:
+                        {"chat_completion": {"temperature": float, "max_tokens": int, "seed": int}}
+        :param variant_name: If set, pins the inference request to a particular variant.
+                             Note: You should generally not do this, and instead let the TensorZero gateway assign a
+                             particular variant. This field is primarily used for testing or debugging purposes.
+        :param dryrun: If true, the request will be executed but won't be stored to the database.
+        :param allowed_tools: If set, restricts the tools available during this inference request.
+                              The list of names should be a subset of the tools configured for the function.
+                              Tools provided at inference time in `additional_tools` (if any) are always available.
+        :param additional_tools: A list of additional tools to use for the request. Each element should look like {"name": str, "parameters": valid JSON Schema, "description": str}
+        :param tool_choice: If set, overrides the tool choice strategy for the request.
+                            It should be one of: "auto", "required", "off", or {"specific": str}. The last option pins the request to a specific tool name.
+        :param parallel_tool_calls: If true, the request will allow for multiple tool calls in a single inference request.
+        :return: If stream is false, returns an InferenceResponse.
+                 If stream is true, returns an async generator that yields InferenceChunks as they come in.
+        """
+        url = urljoin(self.base_url, "inference")
+        data = {
+            "function_name": function_name,
+            "input": input,
+        }
+        if episode_id is not None:
+            data["episode_id"] = str(episode_id)
+        if stream is not None:
+            data["stream"] = stream
+        if params is not None:
+            data["params"] = params
+        if variant_name is not None:
+            data["variant_name"] = variant_name
+        if dryrun is not None:
+            data["dryrun"] = dryrun
+        if allowed_tools is not None:
+            data["allowed_tools"] = allowed_tools
+        if additional_tools is not None:
+            data["additional_tools"] = additional_tools
+        if tool_choice is not None:
+            data["tool_choice"] = tool_choice
+        if parallel_tool_calls is not None:
+            data["parallel_tool_calls"] = parallel_tool_calls
+        response = await self.client.post(url, json=data)
+        response.raise_for_status()
+        if not stream:
+            return parse_inference_response(response.json())
+        else:
+            return self._stream_sse(response)
+
+    async def feedback(
+        self,
+        metric_name: str,
+        value: Any,
+        inference_id: Optional[UUID] = None,
+        episode_id: Optional[UUID] = None,
+        dryrun: Optional[bool] = None,
+    ) -> Dict[str, Any]:
+        """
+        Make a POST request to the /feedback endpoint.
+
+        :param metric_name: The name of the metric to provide feedback for
+        :param value: The value of the feedback. It should correspond to the metric type.
+        :param inference_id: The inference ID to assign the feedback to.
+                             Only use inference IDs that were returned by the TensorZero gateway.
+                             Note: You can assign feedback to either an episode or an inference, but not both.
+        :param episode_id: The episode ID to use for the request
+                           Only use episode IDs that were returned by the TensorZero gateway.
+                           Note: You can assign feedback to either an episode or an inference, but not both.
+        :param dryrun: If true, the feedback request will be executed but won't be stored to the database (i.e. no-op).
+        :return: {"feedback_id": str}
+        """
+        if episode_id is None and inference_id is None:
+            raise ValueError("Either episode_id or inference_id must be provided")
+        if episode_id is not None and inference_id is not None:
+            raise ValueError(
+                "Only one of episode_id or inference_id can be provided, not both"
+            )
+        data = {
+            "metric_name": metric_name,
+            "value": value,
+        }
+        if dryrun is not None:
+            data["dryrun"] = dryrun
+        if episode_id is not None:
+            data["episode_id"] = str(episode_id)
+        if inference_id is not None:
+            data["inference_id"] = str(inference_id)
+        url = urljoin(self.base_url, "feedback")
+        response = await self.client.post(url, json=data)
+        response.raise_for_status()
+        feedback_result = FeedbackResponse(**response.json())
+        return feedback_result
+
+    async def close(self):
+        """
+        Close the connection to the TensorZero gateway.
+        """
+        await self.client.aclose()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+    async def _stream_sse(
+        self, response: httpx.Response
+    ) -> AsyncGenerator[InferenceChunk, None]:
+        """
+        Parse the SSE stream from the response.
+
+        :param response: The httpx.Response object
+        :yield: Parsed SSE events as dictionaries
+        """
+        async for line in response.aiter_lines():
+            if line.startswith("data: "):
+                data = line[6:].strip()
+                if data == "[DONE]":
+                    break
+                try:
+                    data = json.loads(data)
+                    yield parse_inference_chunk(data)
+                except json.JSONDecodeError:
+                    self.logger.error(f"Failed to parse SSE data: {data}")

tensorzero-0.0.1/src/tensorzero/types.py

@@ -0,0 +1,175 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+from uuid import UUID
+
+# Types for non-streaming inference responses
+
+
+@dataclass
+class Usage:
+    input_tokens: int
+    output_tokens: int
+
+
+@dataclass
+class Text:
+    text: str
+
+
+@dataclass
+class ToolCall:
+    name: str
+    arguments: Dict[str, Any]
+    id: str
+    parsed_name: Optional[str]
+    parsed_arguments: Optional[Dict[str, Any]]
+
+
+ContentBlock = Union[Text, ToolCall]
+
+
+@dataclass
+class JsonInferenceOutput:
+    raw: str
+    parsed: Optional[Dict[str, Any]]
+
+
+@dataclass
+class ChatInferenceResponse:
+    inference_id: UUID
+    episode_id: UUID
+    variant_name: str
+    output: List[ContentBlock]
+    usage: Usage
+
+
+@dataclass
+class JsonInferenceResponse:
+    inference_id: UUID
+    episode_id: UUID
+    variant_name: str
+    output: JsonInferenceOutput
+    usage: Usage
+
+
+InferenceResponse = Union[ChatInferenceResponse, JsonInferenceResponse]
+
+
+def parse_inference_response(data: Dict[str, Any]) -> InferenceResponse:
+    if "output" in data and isinstance(data["output"], list):
+        return ChatInferenceResponse(
+            inference_id=UUID(data["inference_id"]),
+            episode_id=UUID(data["episode_id"]),
+            variant_name=data["variant_name"],
+            output=[parse_content_block(block) for block in data["output"]],
+            usage=Usage(**data["usage"]),
+        )
+    elif "output" in data and isinstance(data["output"], dict):
+        return JsonInferenceResponse(
+            inference_id=UUID(data["inference_id"]),
+            episode_id=UUID(data["episode_id"]),
+            variant_name=data["variant_name"],
+            output=JsonInferenceOutput(**data["output"]),
+            usage=Usage(**data["usage"]),
+        )
+    else:
+        raise ValueError("Unable to determine response type")
+
+
+def parse_content_block(block: Dict[str, Any]) -> ContentBlock:
+    block_type = block["type"]
+    if block_type == "text":
+        return Text(text=block["text"])
+    elif block_type == "tool_call":
+        return ToolCall(
+            name=block["name"],
+            arguments=block["arguments"],
+            id=block["id"],
+            parsed_name=block.get("parsed_name"),
+            parsed_arguments=block.get("parsed_arguments"),
+        )
+    else:
+        raise ValueError(f"Unknown content block type: {block}")
+
+
+# Types for streaming inference responses
+
+
+@dataclass
+class TextChunk:
+    # In the possibility that multiple text messages are sent in a single streaming response,
+    # this `id` will be used to disambiguate them
+    id: str
+    text: str
+
+
+@dataclass
+class ToolCallChunk:
+    name: str
+    # This is the tool call ID that many LLM APIs use to associate tool calls with tool responses
+    id: str
+    # `arguments` will come as partial JSON
+    arguments: str
+
+
+ContentBlockChunk = Union[TextChunk, ToolCallChunk]
+
+
+@dataclass
+class ChatChunk:
+    inference_id: UUID
+    episode_id: UUID
+    variant_name: str
+    content: List[ContentBlockChunk]
+    usage: Optional[Usage]
+
+
+@dataclass
+class JsonChunk:
+    inference_id: UUID
+    episode_id: UUID
+    variant_name: str
+    raw: str
+    usage: Optional[Usage]
+
+
+InferenceChunk = Union[ChatChunk, JsonChunk]
+
+
+def parse_inference_chunk(chunk: Dict[str, Any]) -> InferenceChunk:
+    if "content" in chunk:
+        return ChatChunk(
+            inference_id=UUID(chunk["inference_id"]),
+            episode_id=UUID(chunk["episode_id"]),
+            variant_name=chunk["variant_name"],
+            content=[parse_content_block_chunk(block) for block in chunk["content"]],
+            usage=Usage(**chunk["usage"]) if "usage" in chunk else None,
+        )
+    elif "raw" in chunk:
+        return JsonChunk(
+            inference_id=UUID(chunk["inference_id"]),
+            episode_id=UUID(chunk["episode_id"]),
+            variant_name=chunk["variant_name"],
+            raw=chunk["raw"],
+            usage=Usage(**chunk["usage"]) if "usage" in chunk else None,
+        )
+    else:
+        raise ValueError(f"Unable to determine response type: {chunk}")
+
+
+def parse_content_block_chunk(block: Dict[str, Any]) -> ContentBlockChunk:
+    block_type = block["type"]
+    if block_type == "text":
+        return TextChunk(id=block["id"], text=block["text"])
+    elif block_type == "tool_call":
+        return ToolCallChunk(
+            name=block["name"], id=block["id"], arguments=block["arguments"]
+        )
+    else:
+        raise ValueError(f"Unknown content block type: {block}")
+
+
+# Types for feedback
+@dataclass
+class FeedbackResponse:
+    feedback_id: UUID

tensorzero-0.0.1/tests/test_client.py

@@ -0,0 +1,334 @@
+"""
+Tests for the TensorZero client
+
+We use pytest and pytest-asyncio to run the tests.
+
+These tests cover the major functionality of the client but do not
+attempt to comprehensively cover all of TensorZero's functionality.
+See the tests across the Rust codebase for more comprehensive tests.
+
+To run:
+```
+pytest
+```
+or
+```
+uv run pytest
+```
+"""
+
+from uuid import UUID
+
+import pytest
+import pytest_asyncio
+from tensorzero import (
+    AsyncTensorZero,
+    FeedbackResponse,
+    Text,
+    TextChunk,
+    ToolCall,
+    ToolCallChunk,
+)
+from uuid_extensions import uuid7
+
+
+@pytest_asyncio.fixture
+async def client():
+    async with AsyncTensorZero("http://localhost:3000") as client:
+        yield client
+
+
+@pytest.mark.asyncio
+async def test_basic_inference(client):
+    result = await client.inference(
+        function_name="basic_test",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [{"role": "user", "content": "Hello"}],
+        },
+    )
+    assert result.variant_name == "test"
+    output = result.output
+    assert len(output) == 1
+    assert isinstance(output[0], Text)
+    assert (
+        output[0].text
+        == "Megumin gleefully chanted her spell, unleashing a thunderous explosion that lit up the sky and left a massive crater in its wake."
+    )
+    usage = result.usage
+    assert usage.input_tokens == 10
+    assert usage.output_tokens == 10
+
+
+@pytest.mark.asyncio
+async def test_inference_streaming(client):
+    stream = await client.inference(
+        function_name="basic_test",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [{"role": "user", "content": "Hello"}],
+        },
+        stream=True,
+    )
+    chunks = [chunk async for chunk in stream]
+    expected_text = [
+        "Wally,",
+        " the",
+        " golden",
+        " retriever,",
+        " wagged",
+        " his",
+        " tail",
+        " excitedly",
+        " as",
+        " he",
+        " devoured",
+        " a",
+        " slice",
+        " of",
+        " cheese",
+        " pizza.",
+    ]
+    previous_inference_id = None
+    previous_episode_id = None
+    for i, chunk in enumerate(chunks):
+        if previous_inference_id is not None:
+            assert chunk.inference_id == previous_inference_id
+        if previous_episode_id is not None:
+            assert chunk.episode_id == previous_episode_id
+        previous_inference_id = chunk.inference_id
+        previous_episode_id = chunk.episode_id
+        variant_name = chunk.variant_name
+        assert variant_name == "test"
+        if i + 1 < len(chunks):
+            assert len(chunk.content) == 1
+            assert isinstance(chunk.content[0], TextChunk)
+            assert chunk.content[0].text == expected_text[i]
+        else:
+            assert len(chunk.content) == 0
+            assert chunk.usage.input_tokens == 10
+            assert chunk.usage.output_tokens == 16
+
+
+@pytest.mark.asyncio
+async def test_tool_call_inference(client):
+    result = await client.inference(
+        function_name="weather_helper",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [
+                {
+                    "role": "user",
+                    "content": "Hi I'm visiting Brooklyn from Brazil. What's the weather?",
+                }
+            ],
+        },
+    )
+    assert result.variant_name == "variant"
+    output = result.output
+    assert len(output) == 1
+    assert isinstance(output[0], ToolCall)
+    assert output[0].name == "get_temperature"
+    assert output[0].id == "0"
+    assert output[0].arguments == '{"location":"Brooklyn","units":"celsius"}'
+    assert output[0].parsed_name == "get_temperature"
+    assert output[0].parsed_arguments == {"location": "Brooklyn", "units": "celsius"}
+    usage = result.usage
+    assert usage.input_tokens == 10
+    assert usage.output_tokens == 10
+
+
+@pytest.mark.asyncio
+async def test_malformed_tool_call_inference(client):
+    result = await client.inference(
+        function_name="weather_helper",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [
+                {
+                    "role": "user",
+                    "content": "Hi I'm visiting Brooklyn from Brazil. What's the weather?",
+                }
+            ],
+        },
+        variant_name="bad_tool",
+    )
+    assert result.variant_name == "bad_tool"
+    output = result.output
+    assert len(output) == 1
+    assert isinstance(output[0], ToolCall)
+    assert output[0].name == "get_temperature"
+    assert output[0].id == "0"
+    assert output[0].arguments == '{"location":"Brooklyn","units":"Celsius"}'
+    assert output[0].parsed_name == "get_temperature"
+    assert output[0].parsed_arguments is None
+    usage = result.usage
+    assert usage.input_tokens == 10
+    assert usage.output_tokens == 10
+
+
+@pytest.mark.asyncio
+async def test_tool_call_streaming(client):
+    stream = await client.inference(
+        function_name="weather_helper",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [
+                {
+                    "role": "user",
+                    "content": "Hi I'm visiting Brooklyn from Brazil. What's the weather?",
+                }
+            ],
+        },
+        stream=True,
+    )
+    chunks = [chunk async for chunk in stream]
+    expected_text = [
+        '{"location"',
+        ':"Brooklyn"',
+        ',"units"',
+        ':"celsius',
+        '"}',
+    ]
+    previous_inference_id = None
+    previous_episode_id = None
+    for i, chunk in enumerate(chunks):
+        if previous_inference_id is not None:
+            assert chunk.inference_id == previous_inference_id
+        if previous_episode_id is not None:
+            assert chunk.episode_id == previous_episode_id
+        previous_inference_id = chunk.inference_id
+        previous_episode_id = chunk.episode_id
+        variant_name = chunk.variant_name
+        assert variant_name == "variant"
+        if i + 1 < len(chunks):
+            assert len(chunk.content) == 1
+            assert isinstance(chunk.content[0], ToolCallChunk)
+            assert chunk.content[0].name == "get_temperature"
+            assert chunk.content[0].id == "0"
+            assert chunk.content[0].arguments == expected_text[i]
+        else:
+            assert len(chunk.content) == 0
+            assert chunk.usage.input_tokens == 10
+            assert chunk.usage.output_tokens == 5
+
+
+@pytest.mark.asyncio
+async def test_json_streaming(client):
+    # We don't actually have a streaming JSON function implemented in `dummy.rs` but it doesn't matter for this test since
+    # TensorZero doesn't parse the JSON output of the function for streaming calls.
+    stream = await client.inference(
+        function_name="json_success",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [{"role": "user", "content": "Hello, world!"}],
+        },
+        stream=True,
+    )
+    chunks = [chunk async for chunk in stream]
+    expected_text = [
+        "Wally,",
+        " the",
+        " golden",
+        " retriever,",
+        " wagged",
+        " his",
+        " tail",
+        " excitedly",
+        " as",
+        " he",
+        " devoured",
+        " a",
+        " slice",
+        " of",
+        " cheese",
+        " pizza.",
+    ]
+    previous_inference_id = None
+    previous_episode_id = None
+    for i, chunk in enumerate(chunks):
+        if previous_inference_id is not None:
+            assert chunk.inference_id == previous_inference_id
+        if previous_episode_id is not None:
+            assert chunk.episode_id == previous_episode_id
+        previous_inference_id = chunk.inference_id
+        previous_episode_id = chunk.episode_id
+        variant_name = chunk.variant_name
+        assert variant_name == "test"
+        if i + 1 < len(chunks):
+            assert chunk.raw == expected_text[i]
+        else:
+            assert chunk.usage.input_tokens == 10
+            assert chunk.usage.output_tokens == 16
+
+
+@pytest.mark.asyncio
+async def test_json_sucess(client):
+    result = await client.inference(
+        function_name="json_success",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [{"role": "user", "content": "Hello, world!"}],
+        },
+        stream=False,
+    )
+    assert result.variant_name == "test"
+    assert result.output.raw == '{"answer":"Hello"}'
+    assert result.output.parsed == {"answer": "Hello"}
+    assert result.usage.input_tokens == 10
+    assert result.usage.output_tokens == 10
+
+
+@pytest.mark.asyncio
+async def test_json_failure(client):
+    result = await client.inference(
+        function_name="json_fail",
+        input={
+            "system": {"assistant_name": "Alfred Pennyworth"},
+            "messages": [{"role": "user", "content": "Hello, world!"}],
+        },
+        stream=False,
+    )
+    assert result.variant_name == "test"
+    assert (
+        result.output.raw
+        == "Megumin gleefully chanted her spell, unleashing a thunderous explosion that lit up the sky and left a massive crater in its wake."
+    )
+    assert result.output.parsed is None
+    assert result.usage.input_tokens == 10
+    assert result.usage.output_tokens == 10
+
+
+@pytest.mark.asyncio
+async def test_feedback(client):
+    result = await client.feedback(
+        metric_name="user_rating", value=5, episode_id=uuid7()
+    )
+    assert isinstance(result, FeedbackResponse)
+
+    result = await client.feedback(
+        metric_name="task_success", value=True, inference_id=uuid7()
+    )
+    assert isinstance(result, FeedbackResponse)
+
+    result = await client.feedback(
+        metric_name="demonstration", value="hi how are you", inference_id=uuid7()
+    )
+    assert isinstance(result, FeedbackResponse)
+
+
+@pytest.mark.asyncio
+async def test_feedback_invalid_input(client):
+    with pytest.raises(ValueError):
+        await client.feedback(metric_name="test_metric", value=5)
+
+    with pytest.raises(ValueError):
+        await client.feedback(
+            metric_name="test_metric",
+            value=5,
+            episode_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+            inference_id=UUID("123e4567-e89b-12d3-a456-426614174001"),
+        )
+
+
+# Add more tests as needed for edge cases, error handling, etc.

tensorzero-0.0.1/uv.lock

@@ -0,0 +1,208 @@
+version = 1
+requires-python = ">=3.10"
+
+[[package]]
+name = "anyio"
+version = "4.4.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "exceptiongroup", marker = "python_full_version < '3.11'" },
+    { name = "idna" },
+    { name = "sniffio" },
+    { name = "typing-extensions", marker = "python_full_version < '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/e6/e3/c4c8d473d6780ef1853d630d581f70d655b4f8d7553c6997958c283039a2/anyio-4.4.0.tar.gz", hash = "sha256:5aadc6a1bbb7cdb0bede386cac5e2940f5e2ff3aa20277e991cf028e0585ce94", size = 163930 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7b/a2/10639a79341f6c019dedc95bd48a4928eed9f1d1197f4c04f546fc7ae0ff/anyio-4.4.0-py3-none-any.whl", hash = "sha256:c1b2d8f46a8a812513012e1107cb0e68c17159a7a594208005a57dc776e1bdc7", size = 86780 },
+]
+
+[[package]]
+name = "certifi"
+version = "2024.7.4"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/c2/02/a95f2b11e207f68bc64d7aae9666fed2e2b3f307748d5123dffb72a1bbea/certifi-2024.7.4.tar.gz", hash = "sha256:5a1e7645bc0ec61a09e26c36f6106dd4cf40c6db3a1fb6352b0244e7fb057c7b", size = 164065 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/1c/d5/c84e1a17bf61d4df64ca866a1c9a913874b4e9bdc131ec689a0ad013fb36/certifi-2024.7.4-py3-none-any.whl", hash = "sha256:c198e21b1289c2ab85ee4e67bb4b4ef3ead0892059901a8d5b622f24a1101e90", size = 162960 },
+]
+
+[[package]]
+name = "colorama"
+version = "0.4.6"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335 },
+]
+
+[[package]]
+name = "exceptiongroup"
+version = "1.2.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/09/35/2495c4ac46b980e4ca1f6ad6db102322ef3ad2410b79fdde159a4b0f3b92/exceptiongroup-1.2.2.tar.gz", hash = "sha256:47c2edf7c6738fafb49fd34290706d1a1a2f4d1c6df275526b62cbb4aa5393cc", size = 28883 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/02/cc/b7e31358aac6ed1ef2bb790a9746ac2c69bcb3c8588b41616914eb106eaf/exceptiongroup-1.2.2-py3-none-any.whl", hash = "sha256:3111b9d131c238bec2f8f516e123e14ba243563fb135d3fe885990585aa7795b", size = 16453 },
+]
+
+[[package]]
+name = "h11"
+version = "0.14.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/f5/38/3af3d3633a34a3316095b39c8e8fb4853a28a536e55d347bd8d8e9a14b03/h11-0.14.0.tar.gz", hash = "sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d", size = 100418 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/95/04/ff642e65ad6b90db43e668d70ffb6736436c7ce41fcc549f4e9472234127/h11-0.14.0-py3-none-any.whl", hash = "sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761", size = 58259 },
+]
+
+[[package]]
+name = "httpcore"
+version = "1.0.5"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "certifi" },
+    { name = "h11" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/17/b0/5e8b8674f8d203335a62fdfcfa0d11ebe09e23613c3391033cbba35f7926/httpcore-1.0.5.tar.gz", hash = "sha256:34a38e2f9291467ee3b44e89dd52615370e152954ba21721378a87b2960f7a61", size = 83234 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/78/d4/e5d7e4f2174f8a4d63c8897d79eb8fe2503f7ecc03282fee1fa2719c2704/httpcore-1.0.5-py3-none-any.whl", hash = "sha256:421f18bac248b25d310f3cacd198d55b8e6125c107797b609ff9b7a6ba7991b5", size = 77926 },
+]
+
+[[package]]
+name = "httpx"
+version = "0.27.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "anyio" },
+    { name = "certifi" },
+    { name = "httpcore" },
+    { name = "idna" },
+    { name = "sniffio" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/5c/2d/3da5bdf4408b8b2800061c339f240c1802f2e82d55e50bd39c5a881f47f0/httpx-0.27.0.tar.gz", hash = "sha256:a0cb88a46f32dc874e04ee956e4c2764aba2aa228f650b06788ba6bda2962ab5", size = 126413 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/41/7b/ddacf6dcebb42466abd03f368782142baa82e08fc0c1f8eaa05b4bae87d5/httpx-0.27.0-py3-none-any.whl", hash = "sha256:71d5465162c13681bff01ad59b2cc68dd838ea1f10e51574bac27103f00c91a5", size = 75590 },
+]
+
+[[package]]
+name = "idna"
+version = "3.8"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/e8/ac/e349c5e6d4543326c6883ee9491e3921e0d07b55fdf3cce184b40d63e72a/idna-3.8.tar.gz", hash = "sha256:d838c2c0ed6fced7693d5e8ab8e734d5f8fda53a039c0164afb0b82e771e3603", size = 189467 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/22/7e/d71db821f177828df9dea8c42ac46473366f191be53080e552e628aad991/idna-3.8-py3-none-any.whl", hash = "sha256:050b4e5baadcd44d760cedbd2b8e639f2ff89bbc7a5730fcc662954303377aac", size = 66894 },
+]
+
+[[package]]
+name = "iniconfig"
+version = "2.0.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d7/4b/cbd8e699e64a6f16ca3a8220661b5f83792b3017d0f79807cb8708d33913/iniconfig-2.0.0.tar.gz", hash = "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3", size = 4646 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ef/a6/62565a6e1cf69e10f5727360368e451d4b7f58beeac6173dc9db836a5b46/iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374", size = 5892 },
+]
+
+[[package]]
+name = "packaging"
+version = "24.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/51/65/50db4dda066951078f0a96cf12f4b9ada6e4b811516bf0262c0f4f7064d4/packaging-24.1.tar.gz", hash = "sha256:026ed72c8ed3fcce5bf8950572258698927fd1dbda10a5e981cdf0ac37f4f002", size = 148788 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/08/aa/cc0199a5f0ad350994d660967a8efb233fe0416e4639146c089643407ce6/packaging-24.1-py3-none-any.whl", hash = "sha256:5b8f2217dbdbd2f7f384c41c628544e6d52f2d0f53c6d0c3ea61aa5d1d7ff124", size = 53985 },
+]
+
+[[package]]
+name = "pluggy"
+version = "1.5.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/96/2d/02d4312c973c6050a18b314a5ad0b3210edb65a906f868e31c111dede4a6/pluggy-1.5.0.tar.gz", hash = "sha256:2cffa88e94fdc978c4c574f15f9e59b7f4201d439195c3715ca9e2486f1d0cf1", size = 67955 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/88/5f/e351af9a41f866ac3f1fac4ca0613908d9a41741cfcf2228f4ad853b697d/pluggy-1.5.0-py3-none-any.whl", hash = "sha256:44e1ad92c8ca002de6377e165f3e0f1be63266ab4d554740532335b9d75ea669", size = 20556 },
+]
+
+[[package]]
+name = "pytest"
+version = "8.3.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "exceptiongroup", marker = "python_full_version < '3.11'" },
+    { name = "iniconfig" },
+    { name = "packaging" },
+    { name = "pluggy" },
+    { name = "tomli", marker = "python_full_version < '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/b4/8c/9862305bdcd6020bc7b45b1b5e7397a6caf1a33d3025b9a003b39075ffb2/pytest-8.3.2.tar.gz", hash = "sha256:c132345d12ce551242c87269de812483f5bcc87cdbb4722e48487ba194f9fdce", size = 1439314 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/0f/f9/cf155cf32ca7d6fa3601bc4c5dd19086af4b320b706919d48a4c79081cf9/pytest-8.3.2-py3-none-any.whl", hash = "sha256:4ba08f9ae7dcf84ded419494d229b48d0903ea6407b030eaec46df5e6a73bba5", size = 341802 },
+]
+
+[[package]]
+name = "pytest-asyncio"
+version = "0.24.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pytest" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/52/6d/c6cf50ce320cf8611df7a1254d86233b3df7cc07f9b5f5cbcb82e08aa534/pytest_asyncio-0.24.0.tar.gz", hash = "sha256:d081d828e576d85f875399194281e92bf8a68d60d72d1a2faf2feddb6c46b276", size = 49855 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/96/31/6607dab48616902f76885dfcf62c08d929796fc3b2d2318faf9fd54dbed9/pytest_asyncio-0.24.0-py3-none-any.whl", hash = "sha256:a811296ed596b69bf0b6f3dc40f83bcaf341b155a269052d82efa2b25ac7037b", size = 18024 },
+]
+
+[[package]]
+name = "sniffio"
+version = "1.3.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/a2/87/a6771e1546d97e7e041b6ae58d80074f81b7d5121207425c964ddf5cfdbd/sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc", size = 20372 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235 },
+]
+
+[[package]]
+name = "tensorzero"
+version = "0.1.0"
+source = { editable = "." }
+dependencies = [
+    { name = "httpx" },
+]
+
+[package.dev-dependencies]
+dev = [
+    { name = "pytest" },
+    { name = "pytest-asyncio" },
+    { name = "uuid7" },
+]
+
+[package.metadata]
+requires-dist = [{ name = "httpx", specifier = ">=0.27.0" }]
+
+[package.metadata.requires-dev]
+dev = [
+    { name = "pytest", specifier = ">=8.3.2" },
+    { name = "pytest-asyncio", specifier = ">=0.24.0" },
+    { name = "uuid7", specifier = ">=0.1.0" },
+]
+
+[[package]]
+name = "tomli"
+version = "2.0.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/c0/3f/d7af728f075fb08564c5949a9c95e44352e23dee646869fa104a3b2060a3/tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f", size = 15164 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/97/75/10a9ebee3fd790d20926a90a2547f0bf78f371b2f13aa822c759680ca7b9/tomli-2.0.1-py3-none-any.whl", hash = "sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc", size = 12757 },
+]
+
+[[package]]
+name = "typing-extensions"
+version = "4.12.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/df/db/f35a00659bc03fec321ba8bce9420de607a1d37f8342eee1863174c69557/typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8", size = 85321 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/26/9f/ad63fc0248c5379346306f8668cda6e2e2e9c95e01216d2b8ffd9ff037d0/typing_extensions-4.12.2-py3-none-any.whl", hash = "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d", size = 37438 },
+]
+
+[[package]]
+name = "uuid7"
+version = "0.1.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/5c/19/7472bd526591e2192926247109dbf78692e709d3e56775792fec877a7720/uuid7-0.1.0.tar.gz", hash = "sha256:8c57aa32ee7456d3cc68c95c4530bc571646defac01895cfc73545449894a63c", size = 14052 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b5/77/8852f89a91453956582a85024d80ad96f30a41fed4c2b3dce0c9f12ecc7e/uuid7-0.1.0-py2.py3-none-any.whl", hash = "sha256:5e259bb63c8cb4aded5927ff41b444a80d0c7124e8a0ced7cf44efa1f5cccf61", size = 7477 },
+]