Akjava/open_Deep-Research-DuckDuckGo

1import base642import json3import mimetypes4import os5import uuid6from io import BytesIO7from typing import Optional8 9import requests10from dotenv import load_dotenv11from huggingface_hub import InferenceClient12from PIL import Image13from transformers import AutoProcessor14 15from smolagents import Tool, tool16 17 18load_dotenv(override=True)19 20idefics_processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b-chatty")21 22 23def process_images_and_text(image_path, query, client):24    messages = [25        {26            "role": "user",27            "content": [28                {"type": "image"},29                {"type": "text", "text": query},30            ],31        },32    ]33 34    prompt_with_template = idefics_processor.apply_chat_template(messages, add_generation_prompt=True)35 36    # load images from local directory37 38    # encode images to strings which can be sent to the endpoint39    def encode_local_image(image_path):40        # load image41        image = Image.open(image_path).convert("RGB")42 43        # Convert the image to a base64 string44        buffer = BytesIO()45        image.save(buffer, format="JPEG")  # Use the appropriate format (e.g., JPEG, PNG)46        base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")47 48        # add string formatting required by the endpoint49        image_string = f"data:image/jpeg;base64,{base64_image}"50 51        return image_string52 53    image_string = encode_local_image(image_path)54    prompt_with_images = prompt_with_template.replace("<image>", "![]({}) ").format(image_string)55 56    payload = {57        "inputs": prompt_with_images,58        "parameters": {59            "return_full_text": False,60            "max_new_tokens": 200,61        },62    }63 64    return json.loads(client.post(json=payload).decode())[0]65 66 67# Function to encode the image68def encode_image(image_path):69    if image_path.startswith("http"):70        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"71        request_kwargs = {72            "headers": {"User-Agent": user_agent},73            "stream": True,74        }75 76        # Send a HTTP request to the URL77        response = requests.get(image_path, **request_kwargs)78        response.raise_for_status()79        content_type = response.headers.get("content-type", "")80 81        extension = mimetypes.guess_extension(content_type)82        if extension is None:83            extension = ".download"84 85        fname = str(uuid.uuid4()) + extension86        download_path = os.path.abspath(os.path.join("downloads", fname))87 88        with open(download_path, "wb") as fh:89            for chunk in response.iter_content(chunk_size=512):90                fh.write(chunk)91 92        image_path = download_path93 94    with open(image_path, "rb") as image_file:95        return base64.b64encode(image_file.read()).decode("utf-8")96 97 98headers = {"Content-Type": "application/json", "Authorization": f"Bearer {os.getenv('OPENAI_API_KEY')}"}99 100 101def resize_image(image_path):102    img = Image.open(image_path)103    width, height = img.size104    img = img.resize((int(width / 2), int(height / 2)))105    new_image_path = f"resized_{image_path}"106    img.save(new_image_path)107    return new_image_path108 109 110class VisualQATool(Tool):111    name = "visualizer"112    description = "A tool that can answer questions about attached images."113    inputs = {114        "image_path": {115            "description": "The path to the image on which to answer the question",116            "type": "string",117        },118        "question": {"description": "the question to answer", "type": "string", "nullable": True},119    }120    output_type = "string"121 122    client = InferenceClient("HuggingFaceM4/idefics2-8b-chatty")123 124    def forward(self, image_path: str, question: Optional[str] = None) -> str:125        output = ""126        add_note = False127        if not question:128            add_note = True129            question = "Please write a detailed caption for this image."130        try:131            output = process_images_and_text(image_path, question, self.client)132        except Exception as e:133            print(e)134            if "Payload Too Large" in str(e):135                new_image_path = resize_image(image_path)136                output = process_images_and_text(new_image_path, question, self.client)137 138        if add_note:139            output = (140                f"You did not provide a particular question, so here is a detailed caption for the image: {output}"141            )142 143        return output144 145 146@tool147def visualizer(image_path: str, question: Optional[str] = None) -> str:148    """A tool that can answer questions about attached images.149 150    Args:151        image_path: The path to the image on which to answer the question. This should be a local path to downloaded image.152        question: The question to answer.153    """154 155    add_note = False156    if not question:157        add_note = True158        question = "Please write a detailed caption for this image."159    if not isinstance(image_path, str):160        raise Exception("You should provide at least `image_path` string argument to this tool!")161 162    mime_type, _ = mimetypes.guess_type(image_path)163    base64_image = encode_image(image_path)164 165    payload = {166        "model": "gpt-4o",167        "messages": [168            {169                "role": "user",170                "content": [171                    {"type": "text", "text": question},172                    {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{base64_image}"}},173                ],174            }175        ],176        "max_tokens": 1000,177    }178    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)179    try:180        output = response.json()["choices"][0]["message"]["content"]181    except Exception:182        raise Exception(f"Response format unexpected: {response.json()}")183 184    if add_note:185        output = f"You did not provide a particular question, so here is a detailed caption for the image: {output}"186 187    return output188