|
|
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')}"} |
|
|
|
|
|
|
|
|
if isinstance(result["content"], bytes): |
|
|
image_content = result["content"] |
|
|
else: |
|
|
|
|
|
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')}"} |
|
|
|
|
|
|
|
|
if isinstance(result["content"], bytes): |
|
|
pdf_content = result["content"] |
|
|
else: |
|
|
|
|
|
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} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
