pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-23 15:06:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

Feat: bedrock iam authentication (#12020)

Browse Source 
### What problem does this PR solve?

Feat: bedrock iam authentication #12008 

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Magicbook1108
2025-12-18 17:13:09 +08:00
committed by GitHub
parent 151480dc85
commit e84d5412bc
7 changed files with 28 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
rag/llm/chat_model.py
View File

@ -1626,11 +1626,13 @@ class LiteLLMBase(ABC):
elif self.provider == SupportedLiteLLMProvider.Bedrock:
completion_args.pop("api_key", None)
completion_args.pop("api_base", None)
bedrock_credentials = { "aws_region_name": self.bedrock_region }
if self.bedrock_ak and self.bedrock_sk:
bedrock_credentials["aws_access_key_id"] = self.bedrock_ak
bedrock_credentials["aws_secret_access_key"] = self.bedrock_sk
completion_args.update(
{
"aws_access_key_id": self.bedrock_ak,
"aws_secret_access_key": self.bedrock_sk,
"aws_region_name": self.bedrock_region,
"bedrock_credentials": bedrock_credentials,
}
)
elif self.provider == SupportedLiteLLMProvider.OpenRouter:

7
rag/llm/embedding_model.py
View File

@ -471,9 +471,10 @@ class BedrockEmbed(Base):
self.is_amazon = self.model_name.split(".")[0] == "amazon"
self.is_cohere = self.model_name.split(".")[0] == "cohere"
if self.bedrock_ak == "" or self.bedrock_sk == "" or self.bedrock_region == "":
# Try to create a client using the default credentials (AWS_PROFILE, AWS_DEFAULT_REGION, etc.)
self.client = boto3.client("bedrock-runtime")
if self.bedrock_ak == "" or self.bedrock_sk == "":
# Try to create a client using the default credentials if ak/sk are not provided.
# Must provide a region.
self.client = boto3.client("bedrock-runtime", region_name=self.bedrock_region)
else:
self.client = boto3.client(service_name="bedrock-runtime", region_name=self.bedrock_region, aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk)