tools/image_tavily.py

# tools/image_tavily.py

import os
from typing import List, Dict


class TavilyImageSearch:
    """
    Tavily image search API wrapper.
    """

    def __init__(self):
        api_key = os.getenv("TAVILY_API_KEY")
        if not api_key:
            print("⚠️ WARNING: TAVILY_API_KEY not found - image search disabled")
            self.client = None
        else:
            from tavily import TavilyClient
            self.client = TavilyClient(api_key=api_key)

    def search(self, query: str, count: int = 6) -> List[Dict]:
        """
        Fetch images for a query.
        Returns list of dicts with 'url' and 'title' keys for Streamlit compatibility.
        """
        if not self.client:
            print("❌ Image search skipped - no Tavily client")
            return []

        try:
            print(f"🖼️ Searching images for: {query[:30]}...")
            resp = self.client.search(
                query=query,
                max_results=count,
                include_images=True,
                include_answer=False
            )
        except Exception as e:
            print(f"❌ Tavily image search error: {e}")
            return []

        images = []
        raw_images = resp.get("images", [])
        
        for item in raw_images:
            if isinstance(item, dict):
                # Dict format from API
                img_url = item.get("url") or item.get("thumbnail") or item.get("content_url", "")
                images.append({
                    "url": img_url,
                    "thumbnail_url": img_url,
                    "title": item.get("title", item.get("description", "")),
                })
            elif isinstance(item, str):
                # String URL format
                images.append({
                    "url": item,
                    "thumbnail_url": item,
                    "title": "",
                })
        
        return images