Spaces:
Build error
Build error
File size: 4,053 Bytes
715112a |
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 112 113 114 115 116 117 118 119 |
import base64
import mimetypes
import os
import uuid
import requests
from dotenv import load_dotenv
from smolagents import tool
load_dotenv(override=True)
# Function to encode the image
def encode_image(image_path):
if image_path.startswith("http"):
user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
request_kwargs = {
"headers": {"User-Agent": user_agent},
"stream": True,
}
# Send a HTTP request to the URL
response = requests.get(image_path, **request_kwargs)
response.raise_for_status()
content_type = response.headers.get("content-type", "")
extension = mimetypes.guess_extension(content_type)
if extension is None:
extension = ".download"
fname = str(uuid.uuid4()) + extension
download_path = os.path.abspath(os.path.join("downloads", fname))
with open(download_path, "wb") as fh:
for chunk in response.iter_content(chunk_size=512):
fh.write(chunk)
image_path = download_path
with open(image_path, "rb") as image_file:
return base64.b64encode(image_file.read()).decode("utf-8")
@tool
def visualizer(image_path: str, question: str | None = None) -> str:
"""A tool that can answer questions about attached images.
Args:
image_path: The path to the image on which to answer the question. This should be a local path to downloaded image.
question: The question to answer.
"""
import mimetypes
import os
import requests
from .visual_qa import encode_image
add_note = False
if not question:
add_note = True
question = "Please write a detailed caption for this image."
if not isinstance(image_path, str):
raise Exception("You should provide at least `image_path` string argument to this tool!")
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
raise Exception("Google API key not found. Please set the GEMINI_API_KEY environment variable.")
mime_type, _ = mimetypes.guess_type(image_path)
base64_image = encode_image(image_path)
payload = {
"contents": [
{
"parts": [
{"text": question},
{
"inline_data": {
"mime_type": mime_type,
"data": base64_image,
}
},
],
}
],
"generationConfig": {"maxOutputTokens": 2048},
}
headers = {"Content-Type": "application/json"}
url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key={api_key}"
response = requests.post(url, headers=headers, json=payload)
if response.status_code != 200:
raise Exception(f"Request failed with status {response.status_code}: {response.text}")
print(response.json())
try:
response_data = response.json()
candidate = response_data["candidates"][0]
# Improved error handling for specific API responses
finish_reason = candidate.get("finishReason")
if finish_reason == "MAX_TOKENS":
raise Exception("The model's response was truncated because it reached the maximum token limit. The returned content may be incomplete.")
if "parts" not in candidate["content"]:
raise Exception(f"The model returned empty content. Finish Reason: {finish_reason}. Full response: {response_data}")
output = candidate["content"]["parts"][0]["text"]
except (KeyError, IndexError) as e:
# Fallback for any other unexpected format
raise Exception(f"Response format unexpected: {response.json()}") from e
if add_note:
output = f"You did not provide a particular question, so here is a detailed caption for the image: {output}"
return output
|