llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/remote/vector_io/weaviate/weaviate.py

@@ -0,0 +1,393 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+import json
+from typing import Any
+
+import weaviate
+import weaviate.classes as wvc
+from numpy.typing import NDArray
+from weaviate.classes.init import Auth
+from weaviate.classes.query import Filter, HybridFusion
+
+from llama_stack.apis.common.content_types import InterleavedContent
+from llama_stack.apis.common.errors import VectorStoreNotFoundError
+from llama_stack.apis.files import Files
+from llama_stack.apis.inference import Inference
+from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO
+from llama_stack.apis.vector_stores import VectorStore
+from llama_stack.core.request_headers import NeedsRequestProviderData
+from llama_stack.log import get_logger
+from llama_stack.providers.datatypes import VectorStoresProtocolPrivate
+from llama_stack.providers.utils.kvstore import kvstore_impl
+from llama_stack.providers.utils.kvstore.api import KVStore
+from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
+from llama_stack.providers.utils.memory.vector_store import (
+    RERANKER_TYPE_RRF,
+    ChunkForDeletion,
+    EmbeddingIndex,
+    VectorStoreWithIndex,
+)
+from llama_stack.providers.utils.vector_io.vector_utils import sanitize_collection_name
+
+from .config import WeaviateVectorIOConfig
+
+log = get_logger(name=__name__, category="vector_io::weaviate")
+
+VERSION = "v3"
+VECTOR_DBS_PREFIX = f"vector_stores:weaviate:{VERSION}::"
+VECTOR_INDEX_PREFIX = f"vector_index:weaviate:{VERSION}::"
+OPENAI_VECTOR_STORES_PREFIX = f"openai_vector_stores:weaviate:{VERSION}::"
+OPENAI_VECTOR_STORES_FILES_PREFIX = f"openai_vector_stores_files:weaviate:{VERSION}::"
+OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX = f"openai_vector_stores_files_contents:weaviate:{VERSION}::"
+
+
+class WeaviateIndex(EmbeddingIndex):
+    def __init__(self, client: weaviate.WeaviateClient, collection_name: str, kvstore: KVStore | None = None):
+        self.client = client
+        self.collection_name = sanitize_collection_name(collection_name, weaviate_format=True)
+        self.kvstore = kvstore
+
+    async def initialize(self):
+        pass
+
+    async def add_chunks(self, chunks: list[Chunk], embeddings: NDArray):
+        assert len(chunks) == len(embeddings), (
+            f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
+        )
+
+        data_objects = []
+        for chunk, embedding in zip(chunks, embeddings, strict=False):
+            data_objects.append(
+                wvc.data.DataObject(
+                    properties={
+                        "chunk_id": chunk.chunk_id,
+                        "chunk_content": chunk.model_dump_json(),
+                    },
+                    vector=embedding.tolist(),
+                )
+            )
+
+        # Inserting chunks into a prespecified Weaviate collection
+        collection = self.client.collections.get(self.collection_name)
+
+        # TODO: make this async friendly
+        collection.data.insert_many(data_objects)
+
+    async def delete_chunks(self, chunks_for_deletion: list[ChunkForDeletion]) -> None:
+        sanitized_collection_name = sanitize_collection_name(self.collection_name, weaviate_format=True)
+        collection = self.client.collections.get(sanitized_collection_name)
+        chunk_ids = [chunk.chunk_id for chunk in chunks_for_deletion]
+        collection.data.delete_many(where=Filter.by_property("chunk_id").contains_any(chunk_ids))
+
+    async def query_vector(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse:
+        """
+        Performs vector search using Weaviate's built-in vector search.
+        Args:
+            embedding: The query embedding vector
+            k: Limit of number of results to return
+            score_threshold: Minimum similarity score threshold
+        Returns:
+            QueryChunksResponse with chunks and scores.
+        """
+        log.debug(
+            f"WEAVIATE VECTOR SEARCH CALLED: embedding_shape={embedding.shape}, k={k}, threshold={score_threshold}"
+        )
+        sanitized_collection_name = sanitize_collection_name(self.collection_name, weaviate_format=True)
+        collection = self.client.collections.get(sanitized_collection_name)
+
+        try:
+            results = collection.query.near_vector(
+                near_vector=embedding.tolist(), limit=k, return_metadata=wvc.query.MetadataQuery(distance=True)
+            )
+        except Exception as e:
+            log.error(f"Weaviate client vector search failed: {e}")
+            raise
+
+        chunks = []
+        scores = []
+        for doc in results.objects:
+            chunk_json = doc.properties["chunk_content"]
+            try:
+                chunk_dict = json.loads(chunk_json)
+                chunk = Chunk(**chunk_dict)
+            except Exception:
+                log.exception(f"Failed to parse document: {chunk_json}")
+                continue
+
+            if doc.metadata.distance is None:
+                continue
+            # Convert cosine distance ∈ [0,2] -> normalized cosine similarity ∈ [0,1]
+            score = 1.0 - (float(doc.metadata.distance) / 2.0)
+            if score < score_threshold:
+                continue
+
+            chunks.append(chunk)
+            scores.append(score)
+
+        log.debug(f"WEAVIATE VECTOR SEARCH RESULTS: Found {len(chunks)} chunks with scores {scores}")
+        return QueryChunksResponse(chunks=chunks, scores=scores)
+
+    async def delete(self, chunk_ids: list[str] | None = None) -> None:
+        """
+        Delete chunks by IDs if provided, otherwise drop the entire collection.
+        """
+        sanitized_collection_name = sanitize_collection_name(self.collection_name, weaviate_format=True)
+        if chunk_ids is None:
+            # Drop entire collection if it exists
+            if self.client.collections.exists(sanitized_collection_name):
+                self.client.collections.delete(sanitized_collection_name)
+            return
+        collection = self.client.collections.get(sanitized_collection_name)
+        collection.data.delete_many(where=Filter.by_property("id").contains_any(chunk_ids))
+
+    async def query_keyword(self, query_string: str, k: int, score_threshold: float) -> QueryChunksResponse:
+        """
+        Performs BM25-based keyword search using Weaviate's built-in full-text search.
+        Args:
+            query_string: The text query for keyword search
+            k: Limit of number of results to return
+            score_threshold: Minimum similarity score threshold
+        Returns:
+            QueryChunksResponse with chunks and scores
+        """
+        log.debug(f"WEAVIATE KEYWORD SEARCH CALLED: query='{query_string}', k={k}, threshold={score_threshold}")
+        sanitized_collection_name = sanitize_collection_name(self.collection_name, weaviate_format=True)
+        collection = self.client.collections.get(sanitized_collection_name)
+
+        # Perform BM25 keyword search on chunk_content field
+        try:
+            results = collection.query.bm25(
+                query=query_string, limit=k, return_metadata=wvc.query.MetadataQuery(score=True)
+            )
+        except Exception as e:
+            log.error(f"Weaviate client keyword search failed: {e}")
+            raise
+
+        chunks = []
+        scores = []
+        for doc in results.objects:
+            chunk_json = doc.properties["chunk_content"]
+            try:
+                chunk_dict = json.loads(chunk_json)
+                chunk = Chunk(**chunk_dict)
+            except Exception:
+                log.exception(f"Failed to parse document: {chunk_json}")
+                continue
+
+            score = doc.metadata.score if doc.metadata.score is not None else 0.0
+            if score < score_threshold:
+                continue
+
+            chunks.append(chunk)
+            scores.append(score)
+
+        log.debug(f"WEAVIATE KEYWORD SEARCH RESULTS: Found {len(chunks)} chunks with scores {scores}.")
+        return QueryChunksResponse(chunks=chunks, scores=scores)
+
+    async def query_hybrid(
+        self,
+        embedding: NDArray,
+        query_string: str,
+        k: int,
+        score_threshold: float,
+        reranker_type: str,
+        reranker_params: dict[str, Any] | None = None,
+    ) -> QueryChunksResponse:
+        """
+        Hybrid search combining vector similarity and keyword search using Weaviate's native hybrid search.
+        Args:
+            embedding: The query embedding vector
+            query_string: The text query for keyword search
+            k: Limit of number of results to return
+            score_threshold: Minimum similarity score threshold
+            reranker_type: Type of reranker to use ("rrf" or "normalized")
+            reranker_params: Parameters for the reranker
+        Returns:
+            QueryChunksResponse with combined results
+        """
+        log.debug(
+            f"WEAVIATE HYBRID SEARCH CALLED: query='{query_string}', embedding_shape={embedding.shape}, k={k}, threshold={score_threshold}, reranker={reranker_type}"
+        )
+        sanitized_collection_name = sanitize_collection_name(self.collection_name, weaviate_format=True)
+        collection = self.client.collections.get(sanitized_collection_name)
+
+        # Ranked (RRF) reranker fusion type
+        if reranker_type == RERANKER_TYPE_RRF:
+            rerank = HybridFusion.RANKED
+        # Relative score (Normalized) reranker fusion type
+        else:
+            rerank = HybridFusion.RELATIVE_SCORE
+
+        # Perform hybrid search using Weaviate's native hybrid search
+        try:
+            results = collection.query.hybrid(
+                query=query_string,
+                alpha=0.5,  # Range <0, 1>, where 0.5 will equally favor vector and keyword search
+                vector=embedding.tolist(),
+                limit=k,
+                fusion_type=rerank,
+                return_metadata=wvc.query.MetadataQuery(score=True),
+            )
+        except Exception as e:
+            log.error(f"Weaviate client hybrid search failed: {e}")
+            raise
+
+        chunks = []
+        scores = []
+        for doc in results.objects:
+            chunk_json = doc.properties["chunk_content"]
+            try:
+                chunk_dict = json.loads(chunk_json)
+                chunk = Chunk(**chunk_dict)
+            except Exception:
+                log.exception(f"Failed to parse document: {chunk_json}")
+                continue
+
+            score = doc.metadata.score if doc.metadata.score is not None else 0.0
+            if score < score_threshold:
+                continue
+
+            chunks.append(chunk)
+            scores.append(score)
+
+        log.debug(f"WEAVIATE HYBRID SEARCH RESULTS: Found {len(chunks)} chunks with scores {scores}")
+        return QueryChunksResponse(chunks=chunks, scores=scores)
+
+
+class WeaviateVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, NeedsRequestProviderData, VectorStoresProtocolPrivate):
+    def __init__(self, config: WeaviateVectorIOConfig, inference_api: Inference, files_api: Files | None) -> None:
+        super().__init__(files_api=files_api, kvstore=None)
+        self.config = config
+        self.inference_api = inference_api
+        self.client_cache = {}
+        self.cache = {}
+        self.vector_store_table = None
+        self.metadata_collection_name = "openai_vector_stores_metadata"
+
+    def _get_client(self) -> weaviate.WeaviateClient:
+        if "localhost" in self.config.weaviate_cluster_url:
+            log.info("Using Weaviate locally in container")
+            host, port = self.config.weaviate_cluster_url.split(":")
+            key = "local_test"
+            client = weaviate.connect_to_local(host=host, port=port)
+        else:
+            log.info("Using Weaviate remote cluster with URL")
+            key = f"{self.config.weaviate_cluster_url}::{self.config.weaviate_api_key}"
+            if key in self.client_cache:
+                return self.client_cache[key]
+            client = weaviate.connect_to_weaviate_cloud(
+                cluster_url=self.config.weaviate_cluster_url,
+                auth_credentials=Auth.api_key(self.config.weaviate_api_key),
+            )
+        self.client_cache[key] = client
+        return client
+
+    async def initialize(self) -> None:
+        """Set up KV store and load existing vector DBs and OpenAI vector stores."""
+        # Initialize KV store for metadata if configured
+        if self.config.persistence is not None:
+            self.kvstore = await kvstore_impl(self.config.persistence)
+        else:
+            self.kvstore = None
+            log.info("No kvstore configured, registry will not persist across restarts")
+
+        # Load existing vector DB definitions
+        if self.kvstore is not None:
+            start_key = VECTOR_DBS_PREFIX
+            end_key = f"{VECTOR_DBS_PREFIX}\xff"
+            stored = await self.kvstore.values_in_range(start_key, end_key)
+            for raw in stored:
+                vector_store = VectorStore.model_validate_json(raw)
+                client = self._get_client()
+                idx = WeaviateIndex(client=client, collection_name=vector_store.identifier, kvstore=self.kvstore)
+                self.cache[vector_store.identifier] = VectorStoreWithIndex(
+                    vector_store=vector_store, index=idx, inference_api=self.inference_api
+                )
+
+            # Load OpenAI vector stores metadata into cache
+            await self.initialize_openai_vector_stores()
+
+    async def shutdown(self) -> None:
+        for client in self.client_cache.values():
+            client.close()
+        # Clean up mixin resources (file batch tasks)
+        await super().shutdown()
+
+    async def register_vector_store(self, vector_store: VectorStore) -> None:
+        client = self._get_client()
+        sanitized_collection_name = sanitize_collection_name(vector_store.identifier, weaviate_format=True)
+        # Create collection if it doesn't exist
+        if not client.collections.exists(sanitized_collection_name):
+            client.collections.create(
+                name=sanitized_collection_name,
+                vectorizer_config=wvc.config.Configure.Vectorizer.none(),
+                properties=[
+                    wvc.config.Property(name="chunk_content", data_type=wvc.config.DataType.TEXT),
+                ],
+            )
+
+        self.cache[vector_store.identifier] = VectorStoreWithIndex(
+            vector_store, WeaviateIndex(client=client, collection_name=sanitized_collection_name), self.inference_api
+        )
+
+    async def unregister_vector_store(self, vector_store_id: str) -> None:
+        client = self._get_client()
+        sanitized_collection_name = sanitize_collection_name(vector_store_id, weaviate_format=True)
+        if vector_store_id not in self.cache or client.collections.exists(sanitized_collection_name) is False:
+            return
+        client.collections.delete(sanitized_collection_name)
+        await self.cache[vector_store_id].index.delete()
+        del self.cache[vector_store_id]
+
+    async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> VectorStoreWithIndex | None:
+        if vector_store_id in self.cache:
+            return self.cache[vector_store_id]
+
+        # Try to load from kvstore
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before using vector stores.")
+
+        key = f"{VECTOR_DBS_PREFIX}{vector_store_id}"
+        vector_store_data = await self.kvstore.get(key)
+        if not vector_store_data:
+            raise VectorStoreNotFoundError(vector_store_id)
+
+        vector_store = VectorStore.model_validate_json(vector_store_data)
+        client = self._get_client()
+        sanitized_collection_name = sanitize_collection_name(vector_store.identifier, weaviate_format=True)
+        if not client.collections.exists(sanitized_collection_name):
+            raise ValueError(f"Collection with name `{sanitized_collection_name}` not found")
+
+        index = VectorStoreWithIndex(
+            vector_store=vector_store,
+            index=WeaviateIndex(client=client, collection_name=vector_store.identifier),
+            inference_api=self.inference_api,
+        )
+        self.cache[vector_store_id] = index
+        return index
+
+    async def insert_chunks(self, vector_db_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None:
+        index = await self._get_and_cache_vector_store_index(vector_db_id)
+        if not index:
+            raise VectorStoreNotFoundError(vector_db_id)
+
+        await index.insert_chunks(chunks)
+
+    async def query_chunks(
+        self, vector_db_id: str, query: InterleavedContent, params: dict[str, Any] | None = None
+    ) -> QueryChunksResponse:
+        index = await self._get_and_cache_vector_store_index(vector_db_id)
+        if not index:
+            raise VectorStoreNotFoundError(vector_db_id)
+
+        return await index.query_chunks(query, params)
+
+    async def delete_chunks(self, store_id: str, chunks_for_deletion: list[ChunkForDeletion]) -> None:
+        index = await self._get_and_cache_vector_store_index(store_id)
+        if not index:
+            raise ValueError(f"Vector DB {store_id} not found")
+
+        await index.index.delete_chunks(chunks_for_deletion)

llama_stack/providers/utils/bedrock/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

llama_stack/providers/utils/bedrock/client.py

@@ -0,0 +1,74 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+
+import boto3
+from botocore.client import BaseClient
+from botocore.config import Config
+
+from llama_stack.providers.utils.bedrock.config import BedrockBaseConfig
+from llama_stack.providers.utils.bedrock.refreshable_boto_session import (
+    RefreshableBotoSession,
+)
+
+
+def create_bedrock_client(config: BedrockBaseConfig, service_name: str = "bedrock-runtime") -> BaseClient:
+    """Creates a boto3 client for Bedrock services with the given configuration.
+
+    Args:
+        config: The Bedrock configuration containing AWS credentials and settings
+        service_name: The AWS service name to create client for (default: "bedrock-runtime")
+
+    Returns:
+        A configured boto3 client
+    """
+    if config.aws_access_key_id and config.aws_secret_access_key:
+        retries_config = {
+            k: v
+            for k, v in dict(
+                total_max_attempts=config.total_max_attempts,
+                mode=config.retry_mode,
+            ).items()
+            if v is not None
+        }
+
+        config_args = {
+            k: v
+            for k, v in dict(
+                region_name=config.region_name,
+                retries=retries_config if retries_config else None,
+                connect_timeout=config.connect_timeout,
+                read_timeout=config.read_timeout,
+            ).items()
+            if v is not None
+        }
+
+        boto3_config = Config(**config_args)
+
+        session_args = {
+            "aws_access_key_id": config.aws_access_key_id,
+            "aws_secret_access_key": config.aws_secret_access_key,
+            "aws_session_token": config.aws_session_token,
+            "region_name": config.region_name,
+            "profile_name": config.profile_name,
+            "session_ttl": config.session_ttl,
+        }
+
+        # Remove None values
+        session_args = {k: v for k, v in session_args.items() if v is not None}
+
+        boto3_session = boto3.session.Session(**session_args)
+        return boto3_session.client(service_name, config=boto3_config)
+    else:
+        return (
+            RefreshableBotoSession(
+                region_name=config.region_name,
+                profile_name=config.profile_name,
+                session_ttl=config.session_ttl,
+            )
+            .refreshable_session()
+            .client(service_name)
+        )

llama_stack/providers/utils/bedrock/config.py

@@ -0,0 +1,64 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+
+from pydantic import Field
+
+from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
+
+
+class BedrockBaseConfig(RemoteInferenceProviderConfig):
+    auth_credential: None = Field(default=None, exclude=True)
+    aws_access_key_id: str | None = Field(
+        default_factory=lambda: os.getenv("AWS_ACCESS_KEY_ID"),
+        description="The AWS access key to use. Default use environment variable: AWS_ACCESS_KEY_ID",
+    )
+    aws_secret_access_key: str | None = Field(
+        default_factory=lambda: os.getenv("AWS_SECRET_ACCESS_KEY"),
+        description="The AWS secret access key to use. Default use environment variable: AWS_SECRET_ACCESS_KEY",
+    )
+    aws_session_token: str | None = Field(
+        default_factory=lambda: os.getenv("AWS_SESSION_TOKEN"),
+        description="The AWS session token to use. Default use environment variable: AWS_SESSION_TOKEN",
+    )
+    region_name: str | None = Field(
+        default_factory=lambda: os.getenv("AWS_DEFAULT_REGION"),
+        description="The default AWS Region to use, for example, us-west-1 or us-west-2."
+        "Default use environment variable: AWS_DEFAULT_REGION",
+    )
+    profile_name: str | None = Field(
+        default_factory=lambda: os.getenv("AWS_PROFILE"),
+        description="The profile name that contains credentials to use.Default use environment variable: AWS_PROFILE",
+    )
+    total_max_attempts: int | None = Field(
+        default_factory=lambda: int(val) if (val := os.getenv("AWS_MAX_ATTEMPTS")) else None,
+        description="An integer representing the maximum number of attempts that will be made for a single request, "
+        "including the initial attempt. Default use environment variable: AWS_MAX_ATTEMPTS",
+    )
+    retry_mode: str | None = Field(
+        default_factory=lambda: os.getenv("AWS_RETRY_MODE"),
+        description="A string representing the type of retries Boto3 will perform."
+        "Default use environment variable: AWS_RETRY_MODE",
+    )
+    connect_timeout: float | None = Field(
+        default_factory=lambda: float(os.getenv("AWS_CONNECT_TIMEOUT", "60")),
+        description="The time in seconds till a timeout exception is thrown when attempting to make a connection. "
+        "The default is 60 seconds.",
+    )
+    read_timeout: float | None = Field(
+        default_factory=lambda: float(os.getenv("AWS_READ_TIMEOUT", "60")),
+        description="The time in seconds till a timeout exception is thrown when attempting to read from a connection."
+        "The default is 60 seconds.",
+    )
+    session_ttl: int | None = Field(
+        default_factory=lambda: int(os.getenv("AWS_SESSION_TTL", "3600")),
+        description="The time in seconds till a session expires. The default is 3600 seconds (1 hour).",
+    )
+
+    @classmethod
+    def sample_run_config(cls, **kwargs):
+        return {}

llama_stack/providers/utils/bedrock/refreshable_boto_session.py

@@ -0,0 +1,112 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import datetime
+from time import time
+from uuid import uuid4
+
+from boto3 import Session
+from botocore.credentials import RefreshableCredentials
+from botocore.session import get_session
+
+
+class RefreshableBotoSession:
+    """
+    Boto Helper class which lets us create a refreshable session so that we can cache the client or resource.
+
+    Usage
+    -----
+    session = RefreshableBotoSession().refreshable_session()
+
+    client = session.client("s3") # we now can cache this client object without worrying about expiring credentials
+    """
+
+    def __init__(
+        self,
+        region_name: str = None,
+        profile_name: str = None,
+        sts_arn: str = None,
+        session_name: str = None,
+        session_ttl: int = 30000,
+    ):
+        """
+        Initialize `RefreshableBotoSession`
+
+        Parameters
+        ----------
+        region_name : str (optional)
+            Default region when creating a new connection.
+
+        profile_name : str (optional)
+            The name of a profile to use.
+
+        sts_arn : str (optional)
+            The role arn to sts before creating a session.
+
+        session_name : str (optional)
+            An identifier for the assumed role session. (required when `sts_arn` is given)
+
+        session_ttl : int (optional)
+            An integer number to set the TTL for each session. Beyond this session, it will renew the token.
+            50 minutes by default which is before the default role expiration of 1 hour
+        """
+
+        self.region_name = region_name
+        self.profile_name = profile_name
+        self.sts_arn = sts_arn
+        self.session_name = session_name or uuid4().hex
+        self.session_ttl = session_ttl
+
+    def __get_session_credentials(self):
+        """
+        Get session credentials
+        """
+        session = Session(region_name=self.region_name, profile_name=self.profile_name)
+
+        # if sts_arn is given, get credential by assuming the given role
+        if self.sts_arn:
+            sts_client = session.client(service_name="sts", region_name=self.region_name)
+            response = sts_client.assume_role(
+                RoleArn=self.sts_arn,
+                RoleSessionName=self.session_name,
+                DurationSeconds=self.session_ttl,
+            ).get("Credentials")
+
+            credentials = {
+                "access_key": response.get("AccessKeyId"),
+                "secret_key": response.get("SecretAccessKey"),
+                "token": response.get("SessionToken"),
+                "expiry_time": response.get("Expiration").isoformat(),
+            }
+        else:
+            session_credentials = session.get_credentials().get_frozen_credentials()
+            credentials = {
+                "access_key": session_credentials.access_key,
+                "secret_key": session_credentials.secret_key,
+                "token": session_credentials.token,
+                "expiry_time": datetime.datetime.fromtimestamp(time() + self.session_ttl, datetime.UTC).isoformat(),
+            }
+
+        return credentials
+
+    def refreshable_session(self) -> Session:
+        """
+        Get refreshable boto3 session.
+        """
+        # Get refreshable credentials
+        refreshable_credentials = RefreshableCredentials.create_from_metadata(
+            metadata=self.__get_session_credentials(),
+            refresh_using=self.__get_session_credentials,
+            method="sts-assume-role",
+        )
+
+        # attach refreshable credentials current session
+        session = get_session()
+        session._credentials = refreshable_credentials
+        session.set_config_variable("region", self.region_name)
+        autorefresh_session = Session(botocore_session=session)
+
+        return autorefresh_session

llama_stack/providers/utils/common/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.