50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
import os
|
|
import anthropic
|
|
import json
|
|
import base64
|
|
from ..config_loader import get_config
|
|
|
|
def extract(image_bytes: bytes, prompt: str):
|
|
config = get_config()
|
|
api_key = config.get("ai", {}).get("claude_api_key")
|
|
if not api_key:
|
|
return None
|
|
|
|
client = anthropic.Anthropic(api_key=api_key)
|
|
base64_image = base64.b64encode(image_bytes).decode('utf-8')
|
|
|
|
# 2026 Target models
|
|
models_to_try = ["claude-3-5-haiku-latest", "claude-3-5-sonnet-latest"]
|
|
|
|
for model_name in models_to_try:
|
|
try:
|
|
print(f"Claude Attempt: {model_name}")
|
|
message = client.messages.create(
|
|
model=model_name,
|
|
max_tokens=1024,
|
|
messages=[{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "image",
|
|
"source": {
|
|
"type": "base64",
|
|
"media_type": "image/jpeg",
|
|
"data": base64_image,
|
|
},
|
|
},
|
|
{"type": "text", "text": prompt}
|
|
],
|
|
}]
|
|
)
|
|
text = message.content[0].text.strip()
|
|
|
|
if "```json" in text:
|
|
text = text.split("```json")[1].split("```")[0].strip()
|
|
|
|
return json.loads(text)
|
|
except Exception as e:
|
|
print(f"Claude {model_name} failed: {e}")
|
|
continue
|
|
return None
|