File size: 4,416 Bytes
5374a2d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
import requests
import base64
from typing import Dict, Optional, List
from ...tool import Tool
from ...storage_handler import FileStorageHandler, LocalStorageHandler
class ImageAnalysisTool(Tool):
name: str = "image_analysis"
description: str = (
"Analyze and understand images and PDF documents using a multimodal LLM (via OpenRouter). "
"Supports image URLs, local image files, and local PDF files."
)
inputs: Dict[str, Dict[str, str]] = {
"prompt": {"type": "string", "description": "Question or instruction for image/PDF analysis."},
"image_url": {"type": "string", "description": "URL of the image (optional)."},
"image_path": {"type": "string", "description": "Local image file path (optional)."},
"pdf_path": {"type": "string", "description": "Local PDF file path (optional)."},
}
required: Optional[List[str]] = ["prompt"]
def __init__(self, api_key, model="openai/gpt-4o", storage_handler: Optional[FileStorageHandler] = None):
super().__init__()
self.api_key = api_key
self.model = model
self.storage_handler = storage_handler or LocalStorageHandler()
def __call__(
self,
prompt: str,
image_url: str = None,
image_path: str = None,
pdf_path: str = None,
):
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": prompt}
]
}
]
if image_url:
messages[0]["content"].append({
"type": "image_url",
"image_url": {"url": image_url}
})
elif image_path:
try:
result = self.storage_handler.read(image_path)
if not result["success"]:
return {"error": f"Failed to read image: {result.get('error', 'Unknown error')}"}
# Get image content as bytes
if isinstance(result["content"], bytes):
image_content = result["content"]
else:
# If content is not bytes, convert to bytes
image_content = str(result["content"]).encode('utf-8')
base64_image = base64.b64encode(image_content).decode("utf-8")
except Exception as e:
return {"error": f"Failed to read image: {e}"}
data_url = f"data:image/jpeg;base64,{base64_image}"
messages[0]["content"].append({
"type": "image_url",
"image_url": {"url": data_url}
})
elif pdf_path:
try:
result = self.storage_handler.read(pdf_path)
if not result["success"]:
return {"error": f"Failed to read PDF: {result.get('error', 'Unknown error')}"}
# Get PDF content as bytes
if isinstance(result["content"], bytes):
pdf_content = result["content"]
else:
# If content is not bytes, convert to bytes
pdf_content = str(result["content"]).encode('utf-8')
base64_pdf = base64.b64encode(pdf_content).decode("utf-8")
except Exception as e:
return {"error": f"Failed to read PDF: {e}"}
data_url = f"data:application/pdf;base64,{base64_pdf}"
messages[0]["content"].append({
"type": "file",
"file": {"filename": pdf_path.split("/")[-1], "file_data": data_url}
})
payload = {"model": self.model, "messages": messages}
url = "https://openrouter.ai/api/v1/chat/completions"
headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
response = requests.post(url, headers=headers, json=payload)
try:
data = response.json()
result = {
"content": data.get("choices", [{}])[0].get("message", {}).get("content", ""),
"usage": data.get("usage", {})
}
return result
except Exception as e:
return {"error": f"Failed to parse OpenRouter response: {e}", "raw": response.text}
## ImageAnalysisToolkit moved to toolkit.py to consolidate toolkits
|