mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-20 22:51:06 +08:00
Feat: add BedrockCV for vision/image2text inference via LiteLLM (#14705)
## Summary - `CvModel["Bedrock"]` was absent from `rag/llm/cv_model.py`, causing `model_instance()` to return `None` when a Bedrock model was used as a PDF parser — even after correct model resolution. - This PR adds `BedrockCV`, enabling Bedrock vision models (e.g. `amazon.nova-pro-v1:0`, `anthropic.claude-3-5-sonnet`) to be used as PDF parsers. ## What problem does this PR solve? When a Bedrock model is selected as the PDF parser in a knowledge base, ingestion failed with: ``` 'LiteLLMBase' object has no attribute 'describe_with_prompt' ``` The root cause: `LiteLLMBase` (the Bedrock chat implementation) was the only registered handler for the Bedrock factory. It does not implement `describe_with_prompt`. `CvModel` had no Bedrock entry, so `model_instance()` returned `None` for `image2text` requests. ## Type of change - [x] New Feature (non-breaking change which adds functionality) ## Changes **`rag/llm/cv_model.py`** Adds `BedrockCV(Base)` with `_FACTORY_NAME = "Bedrock"`: - Uses `litellm.completion` with the `bedrock/` prefix (consistent with `LiteLLMBase`) - Parses AWS credentials from the JSON key assembled by `add_llm` (`auth_mode`, `bedrock_ak`, `bedrock_sk`, `bedrock_region`, `aws_role_arn`) - Supports three auth modes: `access_key_secret`, `iam_role` (via STS `assume_role`), and default credential chain (IRSA, instance profile) - Implements `describe_with_prompt` and `describe` ## Test plan - [ ] Configure a Bedrock vision model (e.g. `amazon.nova-pro-v1:0`) with valid AWS credentials - [ ] Select it as PDF parser in a knowledge base - [ ] Verify ingestion of a PDF document completes without errors - [ ] Verify `CvModel["Bedrock"]` resolves to `BedrockCV` 🤖 Generated with [Claude Code](https://claude.ai/claude-code) --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -1276,14 +1276,67 @@ class RAGconCV(GptV4):
|
||||
_FACTORY_NAME = "RAGcon"
|
||||
|
||||
def __init__(self, key, model_name, lang="Chinese", base_url="", **kwargs):
|
||||
|
||||
|
||||
if not base_url:
|
||||
base_url = "https://connect.ragcon.com/v1"
|
||||
|
||||
|
||||
# Initialize client
|
||||
self.client = OpenAI(api_key=key, base_url=base_url)
|
||||
self.async_client = AsyncOpenAI(api_key=key, base_url=base_url)
|
||||
self.model_name = model_name
|
||||
self.lang = lang
|
||||
|
||||
Base.__init__(self, **kwargs)
|
||||
|
||||
Base.__init__(self, **kwargs)
|
||||
|
||||
|
||||
class BedrockCV(Base):
|
||||
_FACTORY_NAME = "Bedrock"
|
||||
|
||||
def __init__(self, key, model_name, lang="Chinese", **kwargs):
|
||||
self.model_name = f"bedrock/{model_name}"
|
||||
self.lang = lang
|
||||
self._parse_credentials(key)
|
||||
Base.__init__(self, **kwargs)
|
||||
|
||||
def _parse_credentials(self, key):
|
||||
bedrock_key = json.loads(key)
|
||||
self.auth_mode = bedrock_key.get("auth_mode", "")
|
||||
self.aws_region = bedrock_key.get("bedrock_region", "us-east-1")
|
||||
self.aws_ak = bedrock_key.get("bedrock_ak", "")
|
||||
self.aws_sk = bedrock_key.get("bedrock_sk", "")
|
||||
self.aws_role_arn = bedrock_key.get("aws_role_arn", "")
|
||||
|
||||
def _get_aws_creds(self):
|
||||
if self.auth_mode == "access_key_secret":
|
||||
return {
|
||||
"aws_region_name": self.aws_region,
|
||||
"aws_access_key_id": self.aws_ak,
|
||||
"aws_secret_access_key": self.aws_sk,
|
||||
}
|
||||
elif self.auth_mode == "iam_role":
|
||||
import boto3
|
||||
sts_client = boto3.client("sts", region_name=self.aws_region)
|
||||
resp = sts_client.assume_role(RoleArn=self.aws_role_arn, RoleSessionName="BedrockCVSession")
|
||||
creds = resp["Credentials"]
|
||||
return {
|
||||
"aws_region_name": self.aws_region,
|
||||
"aws_access_key_id": creds["AccessKeyId"],
|
||||
"aws_secret_access_key": creds["SecretAccessKey"],
|
||||
"aws_session_token": creds["SessionToken"],
|
||||
}
|
||||
else:
|
||||
return {"aws_region_name": self.aws_region}
|
||||
|
||||
def describe_with_prompt(self, image, prompt=None):
|
||||
import litellm
|
||||
b64 = self.image2base64(image)
|
||||
messages = self.vision_llm_prompt(b64, prompt)
|
||||
res = litellm.completion(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
**self._get_aws_creds(),
|
||||
)
|
||||
return res.choices[0].message.content.strip(), total_token_count_from_response(res)
|
||||
|
||||
def describe(self, image):
|
||||
return self.describe_with_prompt(image)
|
||||
Reference in New Issue
Block a user