153 lines
5.0 KiB
Python
153 lines
5.0 KiB
Python
"""Cohere Embed v4 (`cohere.embed-v4:0`).
|
|
|
|
Request/response shape taken from the AWS Bedrock user guide page "Cohere
|
|
Embed v4" (read 2026-08-03):
|
|
|
|
request {"input_type": "search_document|search_query|classification|
|
|
clustering",
|
|
"texts": [str], # max 96 per call
|
|
"embedding_types": ["float"|"int8"|"uint8"|"binary"|"ubinary"],
|
|
"output_dimension": 256|512|1024|1536,
|
|
"truncate": "NONE|LEFT|RIGHT"}
|
|
|
|
The response has two documented shapes and this adapter accepts both. Asking
|
|
for one or more `embedding_types` returns
|
|
`{"response_type": "embeddings_by_type", "embeddings": {"float": [[...]]}}`;
|
|
omitting the field returns
|
|
`{"response_type": "embeddings_floats", "embeddings": [[...]]}`.
|
|
|
|
Three defaults are chosen rather than inherited:
|
|
|
|
- `output_dimension` is set explicitly. The documented default is 1536, and a
|
|
collection built at one width cannot absorb vectors of another.
|
|
- `input_type` is derived from the caller's input kind. This is the model
|
|
whose asymmetry the `ports` contract exists for: corpus records go in as
|
|
`search_document`, queries as `search_query`.
|
|
- `truncate` is `NONE`, which makes an over-length input an error instead of a
|
|
silently shortened one. A dosing section that lost its tail and embedded
|
|
anyway is exactly the failure this project's rules are written against.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, List, Sequence
|
|
|
|
from .bedrock_runtime import BedrockInvoker
|
|
from .ports import (
|
|
INPUT_DOCUMENT,
|
|
INPUT_QUERY,
|
|
EmbeddingProvider,
|
|
EmbeddingVector,
|
|
text_digest,
|
|
)
|
|
|
|
MODEL_ID = "cohere.embed-v4:0"
|
|
PROVIDER_NAME = "cohere-v4"
|
|
|
|
SUPPORTED_DIMENSIONS = (256, 512, 1024, 1536)
|
|
|
|
# The documented per-request ceiling for `texts`.
|
|
MAX_TEXTS_PER_REQUEST = 96
|
|
|
|
COHERE_INPUT_TYPES = {
|
|
INPUT_DOCUMENT: "search_document",
|
|
INPUT_QUERY: "search_query",
|
|
}
|
|
|
|
# Cohere's docs do not state whether float vectors are unit-length, so this
|
|
# stays unset rather than being asserted either way.
|
|
NORMALIZED_UNKNOWN = None
|
|
|
|
|
|
class CohereEmbedV4(EmbeddingProvider):
|
|
def __init__(
|
|
self,
|
|
invoker: BedrockInvoker,
|
|
dimensions: int = 1024,
|
|
truncate: str = "NONE",
|
|
batch_size: int = MAX_TEXTS_PER_REQUEST,
|
|
):
|
|
if dimensions not in SUPPORTED_DIMENSIONS:
|
|
raise ValueError(
|
|
f"{MODEL_ID} supports {SUPPORTED_DIMENSIONS}, got {dimensions}"
|
|
)
|
|
if not 1 <= batch_size <= MAX_TEXTS_PER_REQUEST:
|
|
raise ValueError(
|
|
f"batch_size must be 1..{MAX_TEXTS_PER_REQUEST}, got {batch_size}"
|
|
)
|
|
self._invoker = invoker
|
|
self._dimensions = dimensions
|
|
self._truncate = truncate
|
|
self._batch_size = batch_size
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return PROVIDER_NAME
|
|
|
|
@property
|
|
def model_id(self) -> str:
|
|
return MODEL_ID
|
|
|
|
@property
|
|
def dimensions(self) -> int:
|
|
return self._dimensions
|
|
|
|
@property
|
|
def max_batch_size(self) -> int:
|
|
return self._batch_size
|
|
|
|
def _embed_batch(
|
|
self, texts: Sequence[str], input_kind: str
|
|
) -> List[EmbeddingVector]:
|
|
body = self._invoker.invoke_json(
|
|
MODEL_ID,
|
|
{
|
|
"texts": list(texts),
|
|
"input_type": COHERE_INPUT_TYPES[input_kind],
|
|
"embedding_types": ["float"],
|
|
"output_dimension": self._dimensions,
|
|
"truncate": self._truncate,
|
|
},
|
|
# The AWS code example for this model sends `*/*`.
|
|
accept="*/*",
|
|
)
|
|
rows = _float_rows(body)
|
|
if len(rows) != len(texts):
|
|
raise ValueError(
|
|
f"{MODEL_ID} returned {len(rows)} vectors for {len(texts)} texts"
|
|
)
|
|
|
|
vectors: List[EmbeddingVector] = []
|
|
for text, values in zip(texts, rows, strict=True):
|
|
self._check_dimensions(values)
|
|
vectors.append(
|
|
EmbeddingVector(
|
|
values=list(values),
|
|
text_sha256=text_digest(text),
|
|
provider=PROVIDER_NAME,
|
|
model_id=MODEL_ID,
|
|
dimensions=self._dimensions,
|
|
input_kind=input_kind,
|
|
normalized=NORMALIZED_UNKNOWN,
|
|
)
|
|
)
|
|
return vectors
|
|
|
|
|
|
def _float_rows(body: Any) -> List[List[float]]:
|
|
"""Pull the float vectors out of either documented response shape."""
|
|
embeddings = body.get("embeddings")
|
|
if embeddings is None:
|
|
raise ValueError(
|
|
f"{MODEL_ID} response has no 'embeddings' field; "
|
|
f"keys were {sorted(body)}"
|
|
)
|
|
if isinstance(embeddings, dict):
|
|
rows = embeddings.get("float")
|
|
if rows is None:
|
|
raise ValueError(
|
|
f"{MODEL_ID} returned no float embeddings; "
|
|
f"types present: {sorted(embeddings)}"
|
|
)
|
|
return rows
|
|
return embeddings
|