Spaces:
Running
Running
| import json | |
| import boto3 | |
| from src.adapters.interfaces.embeddings import EmbeddingsClient | |
| from src.config import config | |
| class BedrockEmbeddingsClient(EmbeddingsClient): | |
| def __init__(self): | |
| session = boto3.Session( | |
| aws_access_key_id=config.bedrock.access_key_id, | |
| aws_secret_access_key=config.bedrock.secret_access_key, | |
| aws_session_token=config.bedrock.session_token, | |
| region_name=config.bedrock.region, | |
| ) | |
| self.client = session.client("bedrock-runtime") | |
| self.model = config.bedrock.embedding_model | |
| def _embed_one(self, text: str) -> list[float]: | |
| body = json.dumps({"inputText": text}) | |
| response = self.client.invoke_model( | |
| modelId=self.model, | |
| body=body, | |
| contentType="application/json", | |
| accept="application/json", | |
| ) | |
| payload = json.loads(response["body"].read()) | |
| return payload["embedding"] | |
| def embed_text(self, text: str) -> list[float]: | |
| return self._embed_one(text) | |
| def embed_texts(self, texts: list[str]) -> list[list[float]]: | |
| return [self._embed_one(text) for text in texts] | |
| _embeddings_bedrock_client = BedrockEmbeddingsClient() | |