124 lines
4.3 KiB
Python
124 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from rag.ports import QueryEmbeddingUnavailable
|
|
|
|
BEDROCK_RUNTIME_SERVICE = "bedrock-runtime"
|
|
COHERE_EMBED_V4 = "cohere.embed-v4:0"
|
|
|
|
# Cohere embeds corpus records and queries into different subspaces. Sending
|
|
# the corpus value here raises no error — recall just drops silently — so this
|
|
# constant exists to make the asymmetry visible rather than incidental.
|
|
COHERE_QUERY_INPUT_TYPE = "search_query"
|
|
|
|
|
|
def _provider_error_types() -> tuple[type[BaseException], ...]:
|
|
"""botocore's error classes, or none when botocore is absent.
|
|
|
|
Resolved lazily and tolerantly so a test injecting a stub client — and a
|
|
deployment that never enables a cloud provider — neither imports botocore
|
|
nor depends on it being installed.
|
|
"""
|
|
try:
|
|
from botocore.exceptions import BotoCoreError, ClientError
|
|
except ImportError:
|
|
return ()
|
|
return (BotoCoreError, ClientError)
|
|
|
|
|
|
class BedrockCohereQueryEmbedder:
|
|
"""Embeds a query with the same model the collection was built from.
|
|
|
|
A collection loaded with Cohere vectors and queried with any other embedder
|
|
still returns results and still raises nothing — the hits are simply
|
|
meaningless. That failure is silent, which is why the corpus manifest
|
|
records `model_id` and why this adapter names the model explicitly.
|
|
|
|
The request shape is duplicated from `ingestion/embed/bedrock_cohere.py`
|
|
rather than imported: `ingestion` is a separate deployable and importing it
|
|
here would couple the API to the batch pipeline. The duplication is one
|
|
JSON body and is deliberate.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dimensions: int,
|
|
region: str = "us-east-1",
|
|
client: Any | None = None,
|
|
model_id: str = COHERE_EMBED_V4,
|
|
) -> None:
|
|
if dimensions <= 0:
|
|
raise ValueError("dimensions must be positive")
|
|
self._dimensions = dimensions
|
|
self._region = region
|
|
self._client = client
|
|
self._model_id = model_id
|
|
|
|
@property
|
|
def dimensions(self) -> int:
|
|
return self._dimensions
|
|
|
|
@property
|
|
def model_id(self) -> str:
|
|
return self._model_id
|
|
|
|
def _runtime(self) -> Any:
|
|
if self._client is None:
|
|
import boto3
|
|
from botocore.config import Config
|
|
|
|
self._client = boto3.client(
|
|
BEDROCK_RUNTIME_SERVICE,
|
|
region_name=self._region,
|
|
config=Config(
|
|
connect_timeout=10,
|
|
read_timeout=30,
|
|
# "adaptive" was tried 2026-08-07 and reverted same day — see
|
|
# the matching comment in adapters/bedrock_converse.py for the
|
|
# measured 1-5-minute regression it caused. "standard" kept.
|
|
retries={"max_attempts": 4, "mode": "standard"},
|
|
),
|
|
)
|
|
return self._client
|
|
|
|
def embed_query(self, text: str) -> list[float]:
|
|
import json
|
|
|
|
try:
|
|
response = self._runtime().invoke_model(
|
|
modelId=self._model_id,
|
|
body=json.dumps(
|
|
{
|
|
"texts": [text],
|
|
"input_type": COHERE_QUERY_INPUT_TYPE,
|
|
"embedding_types": ["float"],
|
|
"output_dimension": self._dimensions,
|
|
}
|
|
),
|
|
accept="*/*",
|
|
contentType="application/json",
|
|
)
|
|
except _provider_error_types() as error:
|
|
raise QueryEmbeddingUnavailable(
|
|
f"{self._model_id} could not be invoked: {type(error).__name__}"
|
|
) from error
|
|
body = json.loads(response["body"].read())
|
|
embeddings = body.get("embeddings")
|
|
if isinstance(embeddings, dict):
|
|
rows = embeddings.get("float")
|
|
else:
|
|
rows = embeddings
|
|
if not rows:
|
|
raise ValueError(
|
|
f"{self._model_id} returned no float embedding; "
|
|
f"response keys were {sorted(body)}"
|
|
)
|
|
values = list(rows[0])
|
|
if len(values) != self._dimensions:
|
|
raise ValueError(
|
|
f"{self._model_id} returned {len(values)} dimensions, "
|
|
f"expected {self._dimensions}"
|
|
)
|
|
return values
|