diff --git a/python/packages/kagent-adk/src/kagent/adk/_memory_service.py b/python/packages/kagent-adk/src/kagent/adk/_memory_service.py index 797c9ef3f..b5fdcef89 100644 --- a/python/packages/kagent-adk/src/kagent/adk/_memory_service.py +++ b/python/packages/kagent-adk/src/kagent/adk/_memory_service.py @@ -368,6 +368,8 @@ async def _call_embedding_provider( return await self._embed_ollama(model_name, texts, api_base) if provider in ("vertex_ai", "gemini"): return await self._embed_google(provider, model_name, texts) + if provider == "bedrock": + return await self._embed_bedrock(model_name, texts) # Unknown provider — try OpenAI-compatible as a fallback logger.warning("Unknown embedding provider '%s'; attempting OpenAI-compatible call.", provider) return await self._embed_openai("openai", model_name, texts, api_base) @@ -437,6 +439,39 @@ async def _embed_google( ) return [list(emb.values) for emb in response.embeddings] + async def _embed_bedrock( + self, + model_name: str, + texts: List[str], + ) -> List[List[float]]: + """Embed using the AWS Bedrock Titan Embedding API via boto3. + + Uses the same credential chain (env vars, IRSA, instance profile) as + KAgentBedrockLlm. Each text is embedded individually because the + Titan Embedding API accepts a single ``inputText`` per invocation. + """ + import os + + import boto3 + + region = os.environ.get("AWS_DEFAULT_REGION") or os.environ.get("AWS_REGION") or "us-east-1" + client = boto3.client("bedrock-runtime", region_name=region) + + async def _invoke_single(text: str) -> List[float]: + body = json.dumps({"inputText": text}) + response = await asyncio.to_thread( + client.invoke_model, + modelId=model_name, + body=body, + contentType="application/json", + accept="application/json", + ) + result = json.loads(response["body"].read()) + return result["embedding"] + + embeddings = await asyncio.gather(*[_invoke_single(t) for t in texts]) + return list(embeddings) + async def _summarize_session_content_async( self, content: str, diff --git a/python/packages/kagent-adk/tests/unittests/test_embedding.py b/python/packages/kagent-adk/tests/unittests/test_embedding.py index 54464c253..ffeffd78a 100644 --- a/python/packages/kagent-adk/tests/unittests/test_embedding.py +++ b/python/packages/kagent-adk/tests/unittests/test_embedding.py @@ -150,3 +150,43 @@ async def test_embedding_shorter_than_768_rejected(self): mock_cls.return_value = instance result = await svc._generate_embedding_async("test") assert result == [] + + @pytest.mark.asyncio + async def test_bedrock_embed(self): + svc = make_service(provider="bedrock", model="amazon.titan-embed-text-v1") + vec = [0.1] * 1536 + mock_client = mock.MagicMock() + mock_client.invoke_model = mock.MagicMock( + return_value={"body": mock.MagicMock(read=lambda: b'{"embedding": ' + str(vec).encode() + b"}")} + ) + + with mock.patch("boto3.client", return_value=mock_client): + result = await svc._generate_embedding_async("hello world") + assert len(result) == 768 + mock_client.invoke_model.assert_called_once() + + @pytest.mark.asyncio + async def test_bedrock_embed_uses_region_from_env(self): + svc = make_service(provider="bedrock", model="amazon.titan-embed-text-v1") + vec = [0.5] * 1536 + mock_client = mock.MagicMock() + mock_client.invoke_model = mock.MagicMock( + return_value={"body": mock.MagicMock(read=lambda: b'{"embedding": ' + str(vec).encode() + b"}")} + ) + + with ( + mock.patch.dict("os.environ", {"AWS_REGION": "eu-west-1", "AWS_DEFAULT_REGION": ""}), + mock.patch("boto3.client", return_value=mock_client) as mock_boto, + ): + await svc._generate_embedding_async("test") + mock_boto.assert_called_once_with("bedrock-runtime", region_name="eu-west-1") + + @pytest.mark.asyncio + async def test_bedrock_embed_error_returns_empty(self): + svc = make_service(provider="bedrock", model="amazon.titan-embed-text-v1") + mock_client = mock.MagicMock() + mock_client.invoke_model = mock.MagicMock(side_effect=Exception("Bedrock API error")) + + with mock.patch("boto3.client", return_value=mock_client): + result = await svc._generate_embedding_async("test") + assert result == []