diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 988f3ef6..51257e02 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,7 +18,7 @@ jobs: lint: timeout-minutes: 10 name: lint - runs-on: ${{ github.repository == 'stainless-sdks/turbopuffer-python' && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} + runs-on: ${{ startsWith(github.repository, 'stainless-sdks/') && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} if: (github.event_name == 'push' || github.event.pull_request.head.repo.fork) && (github.event_name != 'push' || github.event.head_commit.message != 'codegen metadata') steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -41,7 +41,7 @@ jobs: permissions: contents: read id-token: write - runs-on: ${{ github.repository == 'stainless-sdks/turbopuffer-python' && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} + runs-on: ${{ startsWith(github.repository, 'stainless-sdks/') && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -78,7 +78,7 @@ jobs: test: timeout-minutes: 10 name: test - runs-on: ${{ github.repository == 'stainless-sdks/turbopuffer-python' && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} + runs-on: ${{ startsWith(github.repository, 'stainless-sdks/') && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} if: github.event_name == 'push' || github.event.pull_request.head.repo.fork steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 diff --git a/.stats.yml b/.stats.yml index dd607793..3b87118b 100644 --- a/.stats.yml +++ b/.stats.yml @@ -1,4 +1,4 @@ configured_endpoints: 14 -openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/turbopuffer-benesch/turbopuffer-45f4429a37dcc469269006f05ea38ee3973087f22456b1cdb60b682902997372.yml -openapi_spec_hash: a606206d2ac6886fe31c3392a1f590ea -config_hash: d3fc3b93270f7ee8ae68224dafcfeb6d +openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/turbopuffer-benesch/turbopuffer-c0051833b669a15d2f206974c559f74189cad6df3d636ae7eb620415d7cf6eb9.yml +openapi_spec_hash: 8459bc996ccbfcd259819e2629b7ab84 +config_hash: 7324a6da70bb03cd6ca0e8bd3ee13230 diff --git a/README.md b/README.md index 9f185409..dbced4bc 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,15 @@ and offers both synchronous and asynchronous clients powered by [httpx](https:// It is generated with [Stainless](https://www.stainless.com/). +## MCP Server + +Use the Turbopuffer MCP Server to enable AI assistants to interact with this API, allowing them to explore endpoints, make test requests, and use documentation to help integrate this SDK into your application. + +[![Add to Cursor](https://cursor.com/deeplink/mcp-install-dark.svg)](https://cursor.com/en-US/install-mcp?name=%40turbopuffer%2Fturbopuffer-mcp&config=eyJuYW1lIjoiQHR1cmJvcHVmZmVyL3R1cmJvcHVmZmVyLW1jcCIsInRyYW5zcG9ydCI6Imh0dHAiLCJ1cmwiOiJodHRwczovL3R1cmJvcHVmZmVyLnN0bG1jcC5jb20iLCJoZWFkZXJzIjp7IngtdHVyYm9wdWZmZXItYXBpLWtleSI6InRwdWZfQTEuLi4ifX0) +[![Install in VS Code](https://img.shields.io/badge/_-Add_to_VS_Code-blue?style=for-the-badge&logo=data:image/svg%2bxml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIGZpbGw9Im5vbmUiIHZpZXdCb3g9IjAgMCA0MCA0MCI+PHBhdGggZmlsbD0iI0VFRSIgZmlsbC1ydWxlPSJldmVub2RkIiBkPSJNMzAuMjM1IDM5Ljg4NGEyLjQ5MSAyLjQ5MSAwIDAgMS0xLjc4MS0uNzNMMTIuNyAyNC43OGwtMy40NiAyLjYyNC0zLjQwNiAyLjU4MmExLjY2NSAxLjY2NSAwIDAgMS0xLjA4Mi4zMzggMS42NjQgMS42NjQgMCAwIDEtMS4wNDYtLjQzMWwtMi4yLTJhMS42NjYgMS42NjYgMCAwIDEgMC0yLjQ2M0w3LjQ1OCAyMCA0LjY3IDE3LjQ1MyAxLjUwNyAxNC41N2ExLjY2NSAxLjY2NSAwIDAgMSAwLTIuNDYzbDIuMi0yYTEuNjY1IDEuNjY1IDAgMCAxIDIuMTMtLjA5N2w2Ljg2MyA1LjIwOUwyOC40NTIuODQ0YTIuNDg4IDIuNDg4IDAgMCAxIDEuODQxLS43MjljLjM1MS4wMDkuNjk5LjA5MSAxLjAxOS4yNDVsOC4yMzYgMy45NjFhMi41IDIuNSAwIDAgMSAxLjQxNSAyLjI1M3YuMDk5LS4wNDVWMzMuMzd2LS4wNDUuMDk1YTIuNTAxIDIuNTAxIDAgMCAxLTEuNDE2IDIuMjU3bC04LjIzNSAzLjk2MWEyLjQ5MiAyLjQ5MiAwIDAgMS0xLjA3Ny4yNDZabS43MTYtMjguOTQ3LTExLjk0OCA5LjA2MiAxMS45NTIgOS4wNjUtLjAwNC0xOC4xMjdaIi8+PC9zdmc+)](https://vscode.stainless.com/mcp/%7B%22name%22%3A%22%40turbopuffer%2Fturbopuffer-mcp%22%2C%22type%22%3A%22http%22%2C%22url%22%3A%22https%3A%2F%2Fturbopuffer.stlmcp.com%22%2C%22headers%22%3A%7B%22x-turbopuffer-api-key%22%3A%22tpuf_A1...%22%7D%7D) + +> Note: You may need to set environment variables in your MCP client. + ## Documentation The HTTP API documentation can be found at [turbopuffer.com/docs/api-overview](https://turbopuffer.com/docs/api-overview). diff --git a/api.md b/api.md index 54d42f3d..32ea7708 100644 --- a/api.md +++ b/api.md @@ -37,6 +37,9 @@ from turbopuffer.types import ( FullTextSearchConfig, FuzzyMaxEditDistance, FuzzyParams, + HighlightConfigParams, + HighlightFragmentBy, + HighlightOffsetUnits, ID, IncludeAttributes, Language, diff --git a/scripts/gen b/scripts/gen index a0e66a12..77776859 100755 --- a/scripts/gen +++ b/scripts/gen @@ -4,7 +4,7 @@ set -e cd "$(dirname "$0")/.." -apigen_image=ghcr.io/turbopuffer/turbopuffer-apigen:1b58f2aa9172bf7a668bf862c271b852e95a846b +apigen_image=ghcr.io/turbopuffer/turbopuffer-apigen:50f6977f4020491663b87ecfc2650cc737f03437 apigen() { if [[ "$TURBOPUFFER_DEV_APIGEN" ]]; then diff --git a/scripts/lint b/scripts/lint index eab812fd..cd7bfa8a 100755 --- a/scripts/lint +++ b/scripts/lint @@ -21,7 +21,7 @@ else fi echo "==> Running pyright" -uv run pyright +uv run pyright -p . echo "==> Running mypy" uv run mypy . diff --git a/src/turbopuffer/resources/namespaces.py b/src/turbopuffer/resources/namespaces.py index 3f1293fe..0fbb5e82 100644 --- a/src/turbopuffer/resources/namespaces.py +++ b/src/turbopuffer/resources/namespaces.py @@ -222,6 +222,7 @@ def explain_query( *, namespace: str | None = None, aggregate_by: Dict[str, AggregateBy] | Omit = omit, + compute_attributes: Dict[str, object] | Omit = omit, consistency: namespace_explain_query_params.Consistency | Omit = omit, distance_metric: DistanceMetric | Omit = omit, exclude_attributes: SequenceNotStr[str] | Omit = omit, @@ -246,6 +247,10 @@ def explain_query( aggregate_by: Aggregations to compute over all documents in the namespace that match the filters. + compute_attributes: Computes additional values on documents returned by a query. Each key is the + name of the computed attribute; each value is an expression describing how to + compute it. + consistency: The consistency level for a query. distance_metric: A function used to calculate vector similarity. @@ -286,6 +291,7 @@ def explain_query( body=maybe_transform( { "aggregate_by": aggregate_by, + "compute_attributes": compute_attributes, "consistency": consistency, "distance_metric": distance_metric, "exclude_attributes": exclude_attributes, @@ -434,6 +440,7 @@ def query( *, namespace: str | None = None, aggregate_by: Dict[str, AggregateBy] | Omit = omit, + compute_attributes: Dict[str, object] | Omit = omit, consistency: namespace_query_params.Consistency | Omit = omit, distance_metric: DistanceMetric | Omit = omit, exclude_attributes: SequenceNotStr[str] | Omit = omit, @@ -458,6 +465,10 @@ def query( aggregate_by: Aggregations to compute over all documents in the namespace that match the filters. + compute_attributes: Computes additional values on documents returned by a query. Each key is the + name of the computed attribute; each value is an expression describing how to + compute it. + consistency: The consistency level for a query. distance_metric: A function used to calculate vector similarity. @@ -498,6 +509,7 @@ def query( body=maybe_transform( { "aggregate_by": aggregate_by, + "compute_attributes": compute_attributes, "consistency": consistency, "distance_metric": distance_metric, "exclude_attributes": exclude_attributes, @@ -992,6 +1004,7 @@ async def explain_query( *, namespace: str | None = None, aggregate_by: Dict[str, AggregateBy] | Omit = omit, + compute_attributes: Dict[str, object] | Omit = omit, consistency: namespace_explain_query_params.Consistency | Omit = omit, distance_metric: DistanceMetric | Omit = omit, exclude_attributes: SequenceNotStr[str] | Omit = omit, @@ -1016,6 +1029,10 @@ async def explain_query( aggregate_by: Aggregations to compute over all documents in the namespace that match the filters. + compute_attributes: Computes additional values on documents returned by a query. Each key is the + name of the computed attribute; each value is an expression describing how to + compute it. + consistency: The consistency level for a query. distance_metric: A function used to calculate vector similarity. @@ -1056,6 +1073,7 @@ async def explain_query( body=await async_maybe_transform( { "aggregate_by": aggregate_by, + "compute_attributes": compute_attributes, "consistency": consistency, "distance_metric": distance_metric, "exclude_attributes": exclude_attributes, @@ -1204,6 +1222,7 @@ async def query( *, namespace: str | None = None, aggregate_by: Dict[str, AggregateBy] | Omit = omit, + compute_attributes: Dict[str, object] | Omit = omit, consistency: namespace_query_params.Consistency | Omit = omit, distance_metric: DistanceMetric | Omit = omit, exclude_attributes: SequenceNotStr[str] | Omit = omit, @@ -1228,6 +1247,10 @@ async def query( aggregate_by: Aggregations to compute over all documents in the namespace that match the filters. + compute_attributes: Computes additional values on documents returned by a query. Each key is the + name of the computed attribute; each value is an expression describing how to + compute it. + consistency: The consistency level for a query. distance_metric: A function used to calculate vector similarity. @@ -1268,6 +1291,7 @@ async def query( body=await async_maybe_transform( { "aggregate_by": aggregate_by, + "compute_attributes": compute_attributes, "consistency": consistency, "distance_metric": distance_metric, "exclude_attributes": exclude_attributes, diff --git a/src/turbopuffer/types/__init__.py b/src/turbopuffer/types/__init__.py index 69704124..b05237d6 100644 --- a/src/turbopuffer/types/__init__.py +++ b/src/turbopuffer/types/__init__.py @@ -37,15 +37,18 @@ from .namespace_metadata import NamespaceMetadata as NamespaceMetadata from .pinning_config_param import PinningConfigParam as PinningConfigParam from .attribute_embed_param import AttributeEmbedParam as AttributeEmbedParam +from .highlight_fragment_by import HighlightFragmentBy as HighlightFragmentBy from .sharding_config_param import ShardingConfigParam as ShardingConfigParam from .attribute_embed_config import AttributeEmbedConfig as AttributeEmbedConfig from .attribute_schema_param import AttributeSchemaParam as AttributeSchemaParam from .full_text_search_param import FullTextSearchParam as FullTextSearchParam +from .highlight_offset_units import HighlightOffsetUnits as HighlightOffsetUnits from .namespace_query_params import NamespaceQueryParams as NamespaceQueryParams from .namespace_write_params import NamespaceWriteParams as NamespaceWriteParams from .sparse_distance_metric import SparseDistanceMetric as SparseDistanceMetric from .attribute_schema_config import AttributeSchemaConfig as AttributeSchemaConfig from .full_text_search_config import FullTextSearchConfig as FullTextSearchConfig +from .highlight_config_params import HighlightConfigParams as HighlightConfigParams from .namespace_recall_params import NamespaceRecallParams as NamespaceRecallParams from .client_namespaces_params import ClientNamespacesParams as ClientNamespacesParams from .include_attributes_param import IncludeAttributesParam as IncludeAttributesParam diff --git a/src/turbopuffer/types/custom.py b/src/turbopuffer/types/custom.py index 417347ca..f3fdb6ee 100644 --- a/src/turbopuffer/types/custom.py +++ b/src/turbopuffer/types/custom.py @@ -8,12 +8,23 @@ from .fuzzy_params import FuzzyParams from .saturate_params import SaturateParams from .bm25_clause_params import Bm25ClauseParams +from .highlight_config_params import HighlightConfigParams from .contains_any_token_filter_params import ContainsAnyTokenFilterParams from .contains_all_tokens_filter_params import ContainsAllTokensFilterParams AggregateBy = Union[Tuple[Literal["Count"]], Tuple[Literal["Sum"], str], Tuple[Literal["Count"], str]] +ComputeAttributesVectorDist = Tuple[str, Literal["VectorDist"], Sequence[float]] +ComputeAttributesHighlight = Tuple[Literal["Highlight"], str] +ComputeAttributesHighlightWithConfig = Tuple[Literal["Highlight"], str, HighlightConfigParams] +RankByAnn = Tuple[str, Literal["ANN"], Sequence[float]] +RankByAnnMulti = Tuple[str, Literal["ANN"], Sequence[Sequence[float]]] ExprRefNew = TypedDict("ExprRefNew", {"$ref_new": str}) Expr = Union[ExprRefNew, Tuple[Literal["Embed"], str], Tuple[Literal["Embed"], str, EmbedParams]] +RankByAnnExpr = Tuple[str, Literal["ANN"], Expr] +RankByKnn = Tuple[str, Literal["kNN"], Sequence[float]] +RankByKnnMulti = Tuple[str, Literal["kNN"], Sequence[Sequence[float]]] +RankByKnnExpr = Tuple[str, Literal["kNN"], Expr] +RankBySparseKnn = Tuple[str, Literal["SparseKNN"], Mapping[str, float]] Filter = Union[ Tuple[str, Literal["Eq"], Any], Tuple[str, Literal["NotEq"], Any], @@ -51,15 +62,6 @@ Tuple[Literal["And"], Sequence["Filter"]], Tuple[Literal["Or"], Sequence["Filter"]], ] -GroupByFunction = Tuple[Literal["ForEachUnique"], str] -GroupBy = Union[str, Mapping[str, GroupByFunction]] -RankByAnn = Tuple[str, Literal["ANN"], Sequence[float]] -RankByAnnMulti = Tuple[str, Literal["ANN"], Sequence[Sequence[float]]] -RankByAnnExpr = Tuple[str, Literal["ANN"], Expr] -RankByKnn = Tuple[str, Literal["kNN"], Sequence[float]] -RankByKnnMulti = Tuple[str, Literal["kNN"], Sequence[Sequence[float]]] -RankByKnnExpr = Tuple[str, Literal["kNN"], Expr] -RankBySparseKnn = Tuple[str, Literal["SparseKNN"], Mapping[str, float]] RankByText = Union[ Tuple[str, Literal["BM25"], str], Tuple[str, Literal["BM25"], Sequence[str]], @@ -90,4 +92,12 @@ RankByAttribute, RankByAttributes, ] +ComputeAttributes = Union[ + ComputeAttributesVectorDist, + ComputeAttributesHighlight, + ComputeAttributesHighlightWithConfig, + RankBy, +] +GroupByFunction = Tuple[Literal["ForEachUnique"], str] +GroupBy = Union[str, Mapping[str, GroupByFunction]] RerankBy = Union[Tuple[Literal["RRF"]], Tuple[Literal["RRF"], RrfParams]] diff --git a/src/turbopuffer/types/highlight_config_params.py b/src/turbopuffer/types/highlight_config_params.py new file mode 100644 index 00000000..c76b970d --- /dev/null +++ b/src/turbopuffer/types/highlight_config_params.py @@ -0,0 +1,29 @@ +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from __future__ import annotations + +from typing_extensions import TypedDict + +from .highlight_fragment_by import HighlightFragmentBy +from .highlight_offset_units import HighlightOffsetUnits + +__all__ = ["HighlightConfigParams"] + + +class HighlightConfigParams(TypedDict, total=False): + """Additional (optional) parameters for the Highlight compute expression.""" + + fragment_by: HighlightFragmentBy + """How to split a text attribute into fragments for highlighting.""" + + fragment_limit: int + """The maximum number of fragments to return. Defaults to `3`.""" + + include_offsets: HighlightOffsetUnits + """The units to report highlighted fragment offsets in.""" + + rank_fragments_by: object + """ + How to rank candidate fragments within the attribute before selecting the top + `fragment_limit`. Defaults to the query's `rank_by`. + """ diff --git a/src/turbopuffer/types/highlight_fragment_by.py b/src/turbopuffer/types/highlight_fragment_by.py new file mode 100644 index 00000000..ee15dc29 --- /dev/null +++ b/src/turbopuffer/types/highlight_fragment_by.py @@ -0,0 +1,7 @@ +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from typing_extensions import Literal, TypeAlias + +__all__ = ["HighlightFragmentBy"] + +HighlightFragmentBy: TypeAlias = Literal["none", "sentence", "paragraph", "word"] diff --git a/src/turbopuffer/types/highlight_offset_units.py b/src/turbopuffer/types/highlight_offset_units.py new file mode 100644 index 00000000..9947526a --- /dev/null +++ b/src/turbopuffer/types/highlight_offset_units.py @@ -0,0 +1,7 @@ +# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. + +from typing_extensions import Literal, TypeAlias + +__all__ = ["HighlightOffsetUnits"] + +HighlightOffsetUnits: TypeAlias = Literal["utf-8", "utf-16", "codepoints"] diff --git a/src/turbopuffer/types/namespace_explain_query_params.py b/src/turbopuffer/types/namespace_explain_query_params.py index 9b65b484..325d9ac3 100644 --- a/src/turbopuffer/types/namespace_explain_query_params.py +++ b/src/turbopuffer/types/namespace_explain_query_params.py @@ -5,7 +5,7 @@ from typing import Dict, Union, Iterable from typing_extensions import Literal, TypeAlias, TypedDict -from .custom import GroupBy +from .custom import GroupBy, ComputeAttributes from .._types import SequenceNotStr from .limit_param import LimitParam from .distance_metric import DistanceMetric @@ -24,6 +24,13 @@ class NamespaceExplainQueryParams(TypedDict, total=False): filters. """ + compute_attributes: Dict[str, ComputeAttributes] + """Computes additional values on documents returned by a query. + + Each key is the name of the computed attribute; each value is an expression + describing how to compute it. + """ + consistency: Consistency """The consistency level for a query.""" diff --git a/src/turbopuffer/types/namespace_multi_query_params.py b/src/turbopuffer/types/namespace_multi_query_params.py index 140fd46b..2418d970 100644 --- a/src/turbopuffer/types/namespace_multi_query_params.py +++ b/src/turbopuffer/types/namespace_multi_query_params.py @@ -5,7 +5,7 @@ from typing import Dict, Union, Iterable from typing_extensions import Literal, Required, TypeAlias, TypedDict -from .custom import GroupBy +from .custom import GroupBy, ComputeAttributes from .._types import SequenceNotStr from .limit_param import LimitParam from .distance_metric import DistanceMetric @@ -42,6 +42,13 @@ class Query(TypedDict, total=False): filters. """ + compute_attributes: Dict[str, ComputeAttributes] + """Computes additional values on documents returned by a query. + + Each key is the name of the computed attribute; each value is an expression + describing how to compute it. + """ + distance_metric: DistanceMetric """A function used to calculate vector similarity.""" diff --git a/src/turbopuffer/types/namespace_query_params.py b/src/turbopuffer/types/namespace_query_params.py index 015f4d1d..8371f627 100644 --- a/src/turbopuffer/types/namespace_query_params.py +++ b/src/turbopuffer/types/namespace_query_params.py @@ -5,7 +5,7 @@ from typing import Dict, Union, Iterable from typing_extensions import Literal, TypeAlias, TypedDict -from .custom import Filter, GroupBy, AggregateBy +from .custom import Filter, GroupBy, AggregateBy, ComputeAttributes from .._types import SequenceNotStr from .limit_param import LimitParam from .distance_metric import DistanceMetric @@ -24,6 +24,13 @@ class NamespaceQueryParams(TypedDict, total=False): filters. """ + compute_attributes: Dict[str, ComputeAttributes] + """Computes additional values on documents returned by a query. + + Each key is the name of the computed attribute; each value is an expression + describing how to compute it. + """ + consistency: Consistency """The consistency level for a query.""" diff --git a/tests/api_resources/test_namespaces.py b/tests/api_resources/test_namespaces.py index fea7b631..c6988863 100644 --- a/tests/api_resources/test_namespaces.py +++ b/tests/api_resources/test_namespaces.py @@ -172,6 +172,8 @@ def test_method_explain_query(self, client: Turbopuffer) -> None: @parametrize def test_method_explain_query_with_all_params(self, client: Turbopuffer) -> None: namespace = client.namespace("namespace").explain_query( + aggregate_by={"foo": ("Sum", "bar")}, + compute_attributes={"foo": ("Highlight", "bar")}, consistency={"level": "strong"}, distance_metric="cosine_distance", exclude_attributes=["string"], @@ -294,6 +296,8 @@ def test_method_multi_query_with_all_params(self, client: Turbopuffer) -> None: namespace="namespace", queries=[ { + "aggregate_by": {"foo": "bar"}, + "compute_attributes": {"foo": ("Highlight", "bar")}, "distance_metric": "cosine_distance", "exclude_attributes": ["string"], "include_attributes": True, @@ -355,6 +359,8 @@ def test_method_query(self, client: Turbopuffer) -> None: @parametrize def test_method_query_with_all_params(self, client: Turbopuffer) -> None: namespace = client.namespace("namespace").query( + aggregate_by={"foo": ("Sum", "bar")}, + compute_attributes={"foo": ("Highlight", "bar")}, consistency={"level": "strong"}, distance_metric="cosine_distance", exclude_attributes=["string"], @@ -793,6 +799,8 @@ async def test_method_explain_query(self, async_client: AsyncTurbopuffer) -> Non @parametrize async def test_method_explain_query_with_all_params(self, async_client: AsyncTurbopuffer) -> None: namespace = await async_client.namespace("namespace").explain_query( + aggregate_by={"foo": ("Sum", "bar")}, + compute_attributes={"foo": ("Highlight", "bar")}, consistency={"level": "strong"}, distance_metric="cosine_distance", exclude_attributes=["string"], @@ -918,6 +926,8 @@ async def test_method_multi_query_with_all_params(self, async_client: AsyncTurbo namespace="namespace", queries=[ { + "aggregate_by": {"foo": "bar"}, + "compute_attributes": {"foo": ("Highlight", "bar")}, "distance_metric": "cosine_distance", "exclude_attributes": ["string"], "include_attributes": True, @@ -979,6 +989,8 @@ async def test_method_query(self, async_client: AsyncTurbopuffer) -> None: @parametrize async def test_method_query_with_all_params(self, async_client: AsyncTurbopuffer) -> None: namespace = await async_client.namespace("namespace").query( + aggregate_by={"foo": ("Sum", "bar")}, + compute_attributes={"foo": ("Highlight", "bar")}, consistency={"level": "strong"}, distance_metric="cosine_distance", exclude_attributes=["string"], diff --git a/tests/custom/test_bm25.py b/tests/custom/test_bm25.py index 8964d120..c30f014c 100644 --- a/tests/custom/test_bm25.py +++ b/tests/custom/test_bm25.py @@ -257,7 +257,7 @@ def test_bm25_pre_tokenized_array(tpuf: Turbopuffer): assert len(result.rows) == 2 with pytest.raises( - turbopuffer.APIError, match=r"""invalid input \\'jumped\\' for rank_by field "content", expecting \[\]string""" + turbopuffer.APIError, match=r"""invalid input .* for rank_by field "content", expecting \[\]string""" ): # Query must be an array. ns.query( diff --git a/tests/custom/test_vectors.py b/tests/custom/test_vectors.py index a5970390..9d8618cb 100644 --- a/tests/custom/test_vectors.py +++ b/tests/custom/test_vectors.py @@ -12,7 +12,6 @@ Vector, RowParam, ColumnsParam, - QueryBilling, VectorEncoding, NamespaceQueryResponse, namespace_query_params, @@ -204,10 +203,12 @@ def check_results(vector_set: NamespaceQueryResponse, expected: List[Row]): include_attributes=["hello", "vector"], ) check_results(vector_set, expected) - assert vector_set.billing == QueryBilling( - billable_logical_bytes_queried=256000000, - billable_logical_bytes_returned=105, - ) + assert vector_set.billing is not None + # Production returns 256 MB until turbopuffer#10448 is deployed, then 1.28 GB. + # Accept exactly those two values so this client transition can land first. + # TODO: Remove the 256 MB case after the server rollout is complete. + assert vector_set.billing.billable_logical_bytes_queried in (256_000_000, 1_280_000_000) + assert vector_set.billing.billable_logical_bytes_returned == 105 # Test query with dict vector_set = ns.query( diff --git a/tests/test_compute_attributes.py b/tests/test_compute_attributes.py new file mode 100644 index 00000000..cd47dd86 --- /dev/null +++ b/tests/test_compute_attributes.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +import json +from typing import Any, Dict, List, Tuple, cast + +import pytest + +from turbopuffer.types import ComputeAttributes, NamespaceQueryParams +from turbopuffer._utils import maybe_transform +from turbopuffer.types.highlight_config_params import HighlightConfigParams + +# An empty HighlightConfigParams (all fields optional) is valid for the +# HighlightWithConfig variant. Annotate it so it unifies with the union. +_empty_highlight_config: HighlightConfigParams = {} + +# Each case: (variant name, ComputeAttributes value, expected JSON wire form). +# +# ComputeAttributes = Union[ +# ComputeAttributesVectorDist, # ("vec", "VectorDist", [0.5]) +# ComputeAttributesHighlight, # ("Highlight", "body") +# ComputeAttributesHighlightWithConfig, # ("Highlight", "body", {}) (3rd elem is HighlightConfigParams) +# RankBy, # e.g. ("vec", "ANN", [0.5]) +# ] +# +# The variants are positional tuples that serialize to JSON arrays on the wire. +CASES: List[Tuple[str, ComputeAttributes, List[Any]]] = [ + ("VectorDist", ("vec", "VectorDist", [0.5]), ["vec", "VectorDist", [0.5]]), + ("Highlight", ("Highlight", "body"), ["Highlight", "body"]), + ("HighlightWithConfig", ("Highlight", "body", _empty_highlight_config), ["Highlight", "body", {}]), + ("RankBy", ("vec", "ANN", [0.5]), ["vec", "ANN", [0.5]]), +] + + +@pytest.mark.parametrize(("variant", "value", "expected_wire"), CASES) +def test_compute_attributes_serialization(variant: str, value: ComputeAttributes, expected_wire: List[Any]) -> None: + """Every variant of the ComputeAttributes union serializes to its expected wire form.""" + body: NamespaceQueryParams = { + "compute_attributes": {"my_attr": value}, + "top_k": 10, + } + + transformed = cast(Dict[str, Any], maybe_transform(body, NamespaceQueryParams)) + serialized = transformed["compute_attributes"]["my_attr"] + + # The wire form is JSON, where tuples become arrays. Round-trip through JSON + # to normalize tuples -> lists, then assert structural equality. + assert json.loads(json.dumps(serialized)) == expected_wire, ( + f"variant {variant!r} serialized to {serialized!r}, expected {expected_wire!r}" + ) diff --git a/uv.lock b/uv.lock index 4c617a75..d5a7049c 100644 --- a/uv.lock +++ b/uv.lock @@ -2246,7 +2246,7 @@ wheels = [ [[package]] name = "turbopuffer" -version = "2.4.0" +version = "2.6.0" source = { editable = "." } dependencies = [ { name = "aiohttp" },