diff --git a/nextcloud_mcp_server/config.py b/nextcloud_mcp_server/config.py index 0a1730dce..af613d74f 100644 --- a/nextcloud_mcp_server/config.py +++ b/nextcloud_mcp_server/config.py @@ -110,6 +110,8 @@ "aws_secret_access_key": None, "bedrock_embedding_model": None, "bedrock_generation_model": None, + "bedrock_image_embedding_model": None, + "bedrock_image_output_dim": 1024, # Mistral "mistral_api_key": None, "mistral_embedding_model": "mistral-embed", @@ -614,6 +616,8 @@ class Settings: aws_secret_access_key: str | None = None bedrock_embedding_model: str | None = None bedrock_generation_model: str | None = None + bedrock_image_embedding_model: str | None = None + bedrock_image_output_dim: int = 1024 # Mistral settings (embeddings only) mistral_api_key: str | None = None @@ -793,6 +797,7 @@ def get_embedding_model_name(self) -> str: self.aws_region or self.bedrock_embedding_model or self.bedrock_generation_model + or self.bedrock_image_embedding_model ): return self.bedrock_embedding_model or "bedrock-default" @@ -1097,6 +1102,8 @@ def get_settings() -> Settings: "aws_secret_access_key": "AWS_SECRET_ACCESS_KEY", "bedrock_embedding_model": "BEDROCK_EMBEDDING_MODEL", "bedrock_generation_model": "BEDROCK_GENERATION_MODEL", + "bedrock_image_embedding_model": "BEDROCK_IMAGE_EMBEDDING_MODEL", + "bedrock_image_output_dim": "BEDROCK_IMAGE_OUTPUT_DIM", # Mistral settings "mistral_api_key": "MISTRAL_API_KEY", "mistral_embedding_model": "MISTRAL_EMBEDDING_MODEL", diff --git a/nextcloud_mcp_server/providers/base.py b/nextcloud_mcp_server/providers/base.py index 3ea958271..c690e5f33 100644 --- a/nextcloud_mcp_server/providers/base.py +++ b/nextcloud_mcp_server/providers/base.py @@ -85,6 +85,72 @@ async def generate(self, prompt: str, max_tokens: int = 500) -> str: """ pass + @property + def supports_image_embeddings(self) -> bool: + """Whether this provider can embed images into a joint text-image vector space. + + Default: False. Providers that support multimodal embedding models + (e.g. Bedrock Titan Multimodal G1, Cohere Embed v4) override this. + """ + return False + + async def embed_image( + self, image: bytes, mime_type: str = "image/jpeg" + ) -> list[float]: + """Generate an embedding for an image. + + Args: + image: Raw image bytes (JPEG/PNG/GIF/WebP) + mime_type: Image MIME type (used for providers that require a data URI) + + Returns: + Vector embedding in the joint text-image space + + Raises: + NotImplementedError: If provider doesn't support image embeddings + """ + raise NotImplementedError("Image embeddings not supported by this provider") + + async def embed_image_batch( + self, images: list[bytes], mime_type: str = "image/jpeg" + ) -> list[list[float]]: + """Generate embeddings for multiple images. + + Default implementation calls :meth:`embed_image` sequentially; providers + with a native batch endpoint should override. + + Args: + images: List of raw image byte payloads + mime_type: Image MIME type for all entries + + Returns: + List of vector embeddings, one per image + + Raises: + NotImplementedError: If provider doesn't support image embeddings + """ + return [await self.embed_image(img, mime_type) for img in images] + + async def embed_for_image_space(self, text: str) -> list[float]: + """Embed a text query into the *image* embedding space. + + Distinct from :meth:`embed` because the image embedding model and the + text-document embedding model may be different and produce + incompatible spaces. Callers performing text→image search must use this. + + Raises: + NotImplementedError: If provider doesn't support image embeddings + """ + raise NotImplementedError("Image embeddings not supported by this provider") + + def get_image_dimension(self) -> int: + """Vector dimension of the image embedding space. + + Raises: + NotImplementedError: If provider doesn't support image embeddings + """ + raise NotImplementedError("Image embeddings not supported by this provider") + @abstractmethod async def close(self) -> None: """Close the provider and release resources.""" diff --git a/nextcloud_mcp_server/providers/bedrock.py b/nextcloud_mcp_server/providers/bedrock.py index a1b9c1793..b9f8f3f8d 100644 --- a/nextcloud_mcp_server/providers/bedrock.py +++ b/nextcloud_mcp_server/providers/bedrock.py @@ -1,5 +1,6 @@ """Amazon Bedrock provider for embeddings and text generation.""" +import base64 import json import logging from typing import Any @@ -16,6 +17,10 @@ logger = logging.getLogger(__name__) +# Cohere Embed v4 documents up to 96 images per /v2/embed call; chunk well +# under that to leave headroom for serialization and avoid 400s on edge sizes. +_COHERE_IMAGE_BATCH_SIZE = 64 + class BedrockProvider(Provider): """ @@ -36,6 +41,8 @@ def __init__( region_name: str | None = None, embedding_model: str | None = None, generation_model: str | None = None, + image_embedding_model: str | None = None, + image_output_dim: int = 1024, aws_access_key_id: str | None = None, aws_secret_access_key: str | None = None, ): @@ -44,10 +51,15 @@ def __init__( Args: region_name: AWS region (e.g., "us-east-1"). Defaults to AWS_REGION env var. - embedding_model: Model ID for embeddings (e.g., "amazon.titan-embed-text-v2:0"). - None disables embeddings. + embedding_model: Model ID for text embeddings (e.g., "amazon.titan-embed-text-v2:0"). + None disables text embeddings. generation_model: Model ID for text generation (e.g., "anthropic.claude-3-sonnet-20240229-v1:0"). None disables generation. + image_embedding_model: Model ID for joint text-image embeddings + (e.g., "amazon.titan-embed-image-v1", "cohere.embed-v4:0"). + None disables image embeddings. + image_output_dim: Output dimension for Titan Multimodal G1 (256, 384, or 1024). + Ignored by Cohere and other models. aws_access_key_id: AWS access key (optional, uses default credential chain if not provided) aws_secret_access_key: AWS secret key (optional, uses default credential chain if not provided) @@ -61,7 +73,10 @@ def __init__( self.embedding_model = embedding_model self.generation_model = generation_model + self.image_embedding_model = image_embedding_model + self.image_output_dim = image_output_dim self._dimension: int | None = None # Detected dynamically + self._image_dimension: int | None = None # Detected on first image embed # Initialize bedrock-runtime client client_kwargs: dict[str, Any] = {} @@ -75,10 +90,12 @@ def __init__( self.client = boto3.client("bedrock-runtime", **client_kwargs) logger.info( - "Initialized Bedrock provider in region %s (embedding_model=%s, generation_model=%s)", + "Initialized Bedrock provider in region %s " + "(embedding_model=%s, generation_model=%s, image_embedding_model=%s)", region_name or "default", embedding_model, generation_model, + image_embedding_model, ) @property @@ -255,6 +272,173 @@ def get_dimension(self) -> int: ) return self._dimension + @property + def supports_image_embeddings(self) -> bool: + return self.image_embedding_model is not None + + def _create_image_embedding_request( + self, + *, + image_b64s: list[str] | None = None, + text: str | None = None, + mime_type: str = "image/jpeg", + cohere_input_type: str = "search_document", + ) -> dict[str, Any]: + """Build the Bedrock invoke_model body for the configured image model. + + For Cohere, pass 1+ images via ``image_b64s`` (the API natively batches). + For Titan G1, ``image_b64s`` must have length ≤1 — it only accepts a + single image per call. + """ + if not self.image_embedding_model: + raise NotImplementedError( + "Image embeddings not supported - no image_embedding_model configured" + ) + + if self.image_embedding_model.startswith("amazon.titan-embed-image"): + if image_b64s and len(image_b64s) > 1: + raise ValueError( + "Titan Multimodal G1 accepts only one image per call; " + "callers must iterate via embed_image()" + ) + body: dict[str, Any] = { + "embeddingConfig": {"outputEmbeddingLength": self.image_output_dim}, + } + if image_b64s: + body["inputImage"] = image_b64s[0] + if text is not None: + body["inputText"] = text + return body + + if self.image_embedding_model.startswith("cohere.embed"): + body = { + "input_type": cohere_input_type, + "embedding_types": ["float"], + } + if image_b64s: + body["images"] = [f"data:{mime_type};base64,{b}" for b in image_b64s] + if text is not None: + body["texts"] = [text] + return body + + raise ValueError( + f"Unsupported image embedding model: {self.image_embedding_model}" + ) + + def _parse_image_embedding_response( + self, response: dict[str, Any] + ) -> list[list[float]]: + """Return the list of vectors from a multimodal embedding response. + + Always returns a list (length 1 for single-input models like Titan). + """ + model = self.image_embedding_model or "" + if model.startswith("amazon.titan-embed-image"): + if response.get("message"): + raise RuntimeError( + f"Titan multimodal embedding error: {response['message']}" + ) + return [response["embedding"]] + if model.startswith("cohere.embed"): + return response["embeddings"]["float"] + raise ValueError(f"Unsupported image embedding model: {model}") + + def _invoke_image_model(self, body: dict[str, Any]) -> dict[str, Any]: + if not self.image_embedding_model: + raise NotImplementedError( + "Image embeddings not supported - no image_embedding_model configured" + ) + try: + response = self.client.invoke_model( + modelId=self.image_embedding_model, + body=json.dumps(body), + accept="application/json", + contentType="application/json", + ) + return json.loads(response["body"].read()) + except (BotoCoreError, ClientError) as e: + logger.error("Bedrock image embedding error: %s", e) + raise + + def _remember_image_dim(self, vector: list[float]) -> None: + if self._image_dimension is None: + self._image_dimension = len(vector) + logger.info( + "Detected image embedding dimension: %s for model %s", + self._image_dimension, + self.image_embedding_model, + ) + + async def embed_image( + self, image: bytes, mime_type: str = "image/jpeg" + ) -> list[float]: + if not self.supports_image_embeddings: + raise NotImplementedError( + "Image embeddings not supported - no image_embedding_model configured" + ) + b64 = base64.b64encode(image).decode() + body = self._create_image_embedding_request( + image_b64s=[b64], mime_type=mime_type + ) + vectors = self._parse_image_embedding_response(self._invoke_image_model(body)) + self._remember_image_dim(vectors[0]) + return vectors[0] + + async def embed_image_batch( + self, images: list[bytes], mime_type: str = "image/jpeg" + ) -> list[list[float]]: + if not self.supports_image_embeddings: + raise NotImplementedError( + "Image embeddings not supported - no image_embedding_model configured" + ) + if not images: + return [] + + model = self.image_embedding_model or "" + + if model.startswith("cohere.embed"): + results: list[list[float]] = [] + for i in range(0, len(images), _COHERE_IMAGE_BATCH_SIZE): + chunk = images[i : i + _COHERE_IMAGE_BATCH_SIZE] + b64s = [base64.b64encode(img).decode() for img in chunk] + body = self._create_image_embedding_request( + image_b64s=b64s, mime_type=mime_type + ) + vectors = self._parse_image_embedding_response( + self._invoke_image_model(body) + ) + results.extend(vectors) + self._remember_image_dim(results[0]) + return results + + # Titan and unknown models: sequential fallback + return [await self.embed_image(img, mime_type) for img in images] + + async def embed_for_image_space(self, text: str) -> list[float]: + if not self.supports_image_embeddings: + raise NotImplementedError( + "Image embeddings not supported - no image_embedding_model configured" + ) + body = self._create_image_embedding_request( + text=text, cohere_input_type="search_query" + ) + vectors = self._parse_image_embedding_response(self._invoke_image_model(body)) + self._remember_image_dim(vectors[0]) + return vectors[0] + + def get_image_dimension(self) -> int: + if not self.supports_image_embeddings: + raise NotImplementedError( + "Image embeddings not supported - no image_embedding_model configured" + ) + if self._image_dimension is None: + raise RuntimeError( + f"Image embedding dimension not detected yet for model " + f"{self.image_embedding_model}. Call embed_image() or " + f"embed_for_image_space() first." + ) + return self._image_dimension + def _create_generation_request( self, prompt: str, max_tokens: int ) -> dict[str, Any]: diff --git a/nextcloud_mcp_server/providers/registry.py b/nextcloud_mcp_server/providers/registry.py index 4768f9989..3e9ccbb24 100644 --- a/nextcloud_mcp_server/providers/registry.py +++ b/nextcloud_mcp_server/providers/registry.py @@ -55,18 +55,22 @@ def create_provider() -> Provider: settings.aws_region or settings.bedrock_embedding_model or settings.bedrock_generation_model + or settings.bedrock_image_embedding_model ): logger.info( "Using Bedrock provider: region=%s, embedding_model=%s, " - "generation_model=%s", + "generation_model=%s, image_embedding_model=%s", settings.aws_region, settings.bedrock_embedding_model, settings.bedrock_generation_model, + settings.bedrock_image_embedding_model, ) return BedrockProvider( region_name=settings.aws_region, embedding_model=settings.bedrock_embedding_model, generation_model=settings.bedrock_generation_model, + image_embedding_model=settings.bedrock_image_embedding_model, + image_output_dim=settings.bedrock_image_output_dim, aws_access_key_id=settings.aws_access_key_id, aws_secret_access_key=settings.aws_secret_access_key, ) diff --git a/tests/unit/providers/test_bedrock.py b/tests/unit/providers/test_bedrock.py index e8ac0df6a..a3db3ecb6 100644 --- a/tests/unit/providers/test_bedrock.py +++ b/tests/unit/providers/test_bedrock.py @@ -1,5 +1,6 @@ """Unit tests for Bedrock provider.""" +import base64 import json from unittest.mock import MagicMock @@ -278,3 +279,195 @@ async def test_bedrock_cohere_embedding(mock_bedrock_client): assert embedding == [0.1, 0.2, 0.3] body = json.loads(mock_bedrock_client.invoke_model.call_args.kwargs["body"]) assert body == {"texts": ["test text"], "input_type": "search_document"} + + +def _mock_body(payload: dict) -> dict: + return { + "body": MagicMock(read=MagicMock(return_value=json.dumps(payload).encode())) + } + + +@pytest.mark.unit +async def test_bedrock_image_embed_titan(mock_bedrock_client): + """Titan multimodal: image bytes → vector via inputImage + outputEmbeddingLength.""" + mock_bedrock_client.invoke_model.return_value = _mock_body( + {"embedding": [0.1] * 1024} + ) + + provider = BedrockProvider( + region_name="us-east-1", + image_embedding_model="amazon.titan-embed-image-v1", + image_output_dim=1024, + ) + assert provider.supports_image_embeddings is True + + vec = await provider.embed_image(b"\xff\xd8\xff\xe0fake-jpeg") + + assert len(vec) == 1024 + assert provider.get_image_dimension() == 1024 + call = mock_bedrock_client.invoke_model.call_args + assert call.kwargs["modelId"] == "amazon.titan-embed-image-v1" + body = json.loads(call.kwargs["body"]) + assert body["inputImage"] == base64.b64encode(b"\xff\xd8\xff\xe0fake-jpeg").decode() + assert body["embeddingConfig"] == {"outputEmbeddingLength": 1024} + assert "inputText" not in body + + +@pytest.mark.unit +async def test_bedrock_image_embed_titan_returns_error_message(mock_bedrock_client): + """Titan returns errors via a `message` field — must surface as RuntimeError.""" + mock_bedrock_client.invoke_model.return_value = _mock_body( + {"embedding": [], "message": "Image too small"} + ) + provider = BedrockProvider( + region_name="us-east-1", + image_embedding_model="amazon.titan-embed-image-v1", + ) + + with pytest.raises(RuntimeError, match="Image too small"): + await provider.embed_image(b"tiny") + + +@pytest.mark.unit +async def test_bedrock_embed_for_image_space_titan(mock_bedrock_client): + """Text→image-space query: Titan uses inputText against the same image model.""" + mock_bedrock_client.invoke_model.return_value = _mock_body( + {"embedding": [0.5] * 1024} + ) + provider = BedrockProvider( + region_name="us-east-1", + image_embedding_model="amazon.titan-embed-image-v1", + image_output_dim=1024, + ) + + vec = await provider.embed_for_image_space("a coast at sunset") + + assert len(vec) == 1024 + body = json.loads(mock_bedrock_client.invoke_model.call_args.kwargs["body"]) + assert body["inputText"] == "a coast at sunset" + assert "inputImage" not in body + assert body["embeddingConfig"] == {"outputEmbeddingLength": 1024} + + +@pytest.mark.unit +async def test_bedrock_image_embed_cohere_batch(mock_bedrock_client): + """Cohere v4: batch image embedding in a single invoke_model call.""" + mock_bedrock_client.invoke_model.return_value = _mock_body( + {"embeddings": {"float": [[0.1] * 1536, [0.2] * 1536, [0.3] * 1536]}} + ) + + provider = BedrockProvider( + region_name="us-east-1", + image_embedding_model="cohere.embed-v4:0", + ) + + imgs = [b"img1", b"img2", b"img3"] + vecs = await provider.embed_image_batch(imgs, mime_type="image/png") + + assert len(vecs) == 3 + assert all(len(v) == 1536 for v in vecs) + assert mock_bedrock_client.invoke_model.call_count == 1 # batched + body = json.loads(mock_bedrock_client.invoke_model.call_args.kwargs["body"]) + assert body["input_type"] == "search_document" + assert body["embedding_types"] == ["float"] + assert len(body["images"]) == 3 + assert body["images"][0].startswith("data:image/png;base64,") + + +@pytest.mark.unit +async def test_bedrock_embed_for_image_space_cohere(mock_bedrock_client): + """Cohere v4 text→image-space: input_type=search_query, single vector returned.""" + mock_bedrock_client.invoke_model.return_value = _mock_body( + {"embeddings": {"float": [[0.7] * 1536]}} + ) + provider = BedrockProvider( + region_name="us-east-1", + image_embedding_model="cohere.embed-v4:0", + ) + + vec = await provider.embed_for_image_space("hummingbird") + + assert len(vec) == 1536 + body = json.loads(mock_bedrock_client.invoke_model.call_args.kwargs["body"]) + assert body["texts"] == ["hummingbird"] + assert body["input_type"] == "search_query" + assert "images" not in body + + +@pytest.mark.unit +async def test_bedrock_image_embeddings_disabled(): + """No image_embedding_model → capability False, all image methods raise.""" + provider = BedrockProvider( + region_name="us-east-1", + embedding_model="amazon.titan-embed-text-v2:0", + ) + assert provider.supports_image_embeddings is False + + with pytest.raises(NotImplementedError, match="no image_embedding_model"): + await provider.embed_image(b"x") + with pytest.raises(NotImplementedError, match="no image_embedding_model"): + await provider.embed_image_batch([b"x"]) + with pytest.raises(NotImplementedError, match="no image_embedding_model"): + await provider.embed_for_image_space("q") + with pytest.raises(NotImplementedError, match="no image_embedding_model"): + provider.get_image_dimension() + + +@pytest.mark.unit +async def test_bedrock_image_dimension_not_detected_yet(): + """get_image_dimension before any embed call raises RuntimeError.""" + provider = BedrockProvider( + region_name="us-east-1", + image_embedding_model="amazon.titan-embed-image-v1", + ) + with pytest.raises(RuntimeError, match="not detected yet"): + provider.get_image_dimension() + + +@pytest.mark.unit +async def test_bedrock_image_embed_cohere_chunks_over_cap(mock_bedrock_client): + """Cohere batch >64 images chunks into multiple invoke_model calls to stay + under the per-request cap (96).""" + chunk1_vecs = [[0.1] * 8 for _ in range(64)] + chunk2_vecs = [[0.2] * 8 for _ in range(36)] + responses = [ + _mock_body({"embeddings": {"float": chunk1_vecs}}), + _mock_body({"embeddings": {"float": chunk2_vecs}}), + ] + mock_bedrock_client.invoke_model.side_effect = responses + + provider = BedrockProvider( + region_name="us-east-1", + image_embedding_model="cohere.embed-v4:0", + ) + + images = [f"img{i}".encode() for i in range(100)] + vecs = await provider.embed_image_batch(images) + + assert len(vecs) == 100 + assert mock_bedrock_client.invoke_model.call_count == 2 + body1 = json.loads( + mock_bedrock_client.invoke_model.call_args_list[0].kwargs["body"] + ) + body2 = json.loads( + mock_bedrock_client.invoke_model.call_args_list[1].kwargs["body"] + ) + assert len(body1["images"]) == 64 + assert len(body2["images"]) == 36 + + +@pytest.mark.unit +async def test_bedrock_image_embed_batch_titan_sequential(mock_bedrock_client): + """Titan has no batch endpoint — embed_image_batch falls back to sequential calls.""" + mock_bedrock_client.invoke_model.return_value = _mock_body( + {"embedding": [0.1] * 1024} + ) + provider = BedrockProvider( + region_name="us-east-1", + image_embedding_model="amazon.titan-embed-image-v1", + ) + + vecs = await provider.embed_image_batch([b"a", b"b"]) + + assert len(vecs) == 2 + assert mock_bedrock_client.invoke_model.call_count == 2