mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
test(integration): bedrock cohere.embed-english-v3 embeddings accept encoding_format and dimensions (Pylon #7963)
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
989d7b87b2
commit
8e846d4cad
2 changed files with 61 additions and 0 deletions
|
|
@ -128,6 +128,9 @@
|
|||
"tests/integration/providers/test_bedrock_auth_wire.py::test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload": [
|
||||
"other.provider_wire.bedrock.bearer_db_yaml_survives_reload"
|
||||
],
|
||||
"tests/integration/providers/test_bedrock_embedding_wire.py::test_cohere_embed_english_v3_accepts_encoding_format_and_dimensions": [
|
||||
"other.provider_wire.bedrock.cohere_embed_english_v3_accepts_encoding_format"
|
||||
],
|
||||
"tests/integration/streaming/test_stream_contracts.py::test_generated_tcp_partitions_preserve_unicode_text_identity_and_final_usage": [
|
||||
"other.streaming.byte_partitions.preserve_text_identity_and_usage"
|
||||
],
|
||||
|
|
|
|||
58
tests/integration/providers/test_bedrock_embedding_wire.py
Normal file
58
tests/integration/providers/test_bedrock_embedding_wire.py
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
MODEL: Final = "bedrock/cohere.embed-english-v3"
|
||||
TOKEN: Final = "synthetic-bedrock-bearer"
|
||||
INPUT: Final = "hello world"
|
||||
VECTOR: Final = [0.1, 0.2, 0.3]
|
||||
RESPONSE: Final = json.dumps(
|
||||
{
|
||||
"embeddings": {"float": [VECTOR]},
|
||||
"id": "synthetic-cohere-embed",
|
||||
"response_type": "embeddings_by_type",
|
||||
"texts": [INPUT],
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
def cohere_english_v3_peer(request: Request) -> Reply:
|
||||
assert request.method == "POST" and request.target == "/model/cohere.embed-english-v3/invoke"
|
||||
assert request.headers["authorization"] == f"Bearer {TOKEN}"
|
||||
assert json.loads(request.body) == {
|
||||
"texts": [INPUT],
|
||||
"input_type": "search_document",
|
||||
"embedding_types": ["float"],
|
||||
"output_dimension": 512,
|
||||
}
|
||||
return Reply(body=RESPONSE)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.provider_wire.bedrock.cohere_embed_english_v3_accepts_encoding_format")
|
||||
def test_cohere_embed_english_v3_accepts_encoding_format_and_dimensions(gateway: Gateway) -> None:
|
||||
with wire_server(cohere_english_v3_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=MODEL,
|
||||
api_key=TOKEN,
|
||||
api_base=wire.url,
|
||||
aws_region_name="us-east-1",
|
||||
)
|
||||
for encoding_format in ("float", "base64"):
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/embeddings",
|
||||
{
|
||||
"model": model,
|
||||
"input": INPUT,
|
||||
"encoding_format": encoding_format,
|
||||
"dimensions": 512,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, f"encoding_format={encoding_format}: {response.text}"
|
||||
assert response.json()["data"] == [
|
||||
{"object": "embedding", "index": 0, "embedding": VECTOR, "type": "float"},
|
||||
], response.text
|
||||
assert len(wire.drain()) == 1, f"encoding_format={encoding_format} never reached Bedrock"
|
||||
Loading…
Add table
Reference in a new issue