Update app.py

app.py

@@ -3,8 +3,8 @@ Blockchain Wallet Analyzer - A tool for analyzing Ethereum wallet contents and N
 
 This module provides a complete implementation of a blockchain wallet analysis tool
 with a Gradio web interface. It includes wallet analysis, NFT tracking, 
-interactive chat capabilities via the OpenAI API, 
-and now NFT image rendering from OpenSea (with user-provided API key).
+interactive chat capabilities using the OpenAI API, 
+and NFT image rendering from OpenSea via PIL objects.
 
 Author: Claude
 Date: January 2025
@@ -25,11 +25,11 @@ from decimal import Decimal
 from dataclasses import dataclass
 from pathlib import Path
 from io import BytesIO
-from PIL import Image
 
 import aiohttp
 import openai
 import gradio as gr
+from PIL import Image
 from tenacity import retry, stop_after_attempt, wait_exponential
 
 # Type variables
@@ -69,6 +69,7 @@ class Config:
     Your personality:
     - Friendly and enthusiastic
     - Explain findings in fun, simple ways
+    - Provide NFT images from OpenSea when possible
 
     Instructions:
     - You have access to detailed wallet data in your context
@@ -78,34 +79,36 @@ class Config:
     """
     ETHERSCAN_BASE_URL: str = "https://api.etherscan.io/api"
     ETHEREUM_ADDRESS_REGEX: str = r"0x[a-fA-F0-9]{40}"
-    RATE_LIMIT_DELAY: float = 0.2   # 5 requests/sec for free tier
+    RATE_LIMIT_DELAY: float = 0.2  # 5 requests per second
     MAX_RETRIES: int = 3
-    OPENAI_MODEL: str = "gpt-4o-mini"  # GPT model name
+    OPENAI_MODEL: str = "gpt-4o-mini"  # Example updated model
     MAX_TOKENS: int = 4000
     TEMPERATURE: float = 0.7
     HISTORY_LIMIT: int = 5
+#### Part 2: Wallet Analyzer (Etherscan)
+
 class WalletAnalyzer:
     """Analyzes Ethereum wallet contents using Etherscan API."""
     
     def __init__(self, api_key: str):
-        """Initialize the analyzer with API key."""
         self.api_key = api_key
         self.base_url = Config.ETHERSCAN_BASE_URL
-        self.last_request_time = 0
         self.session: Optional[aiohttp.ClientSession] = None
+        self.last_request_time = 0
     
     async def __aenter__(self) -> WalletAnalyzer:
-        """Create aiohttp session on context manager enter."""
         self.session = aiohttp.ClientSession()
         return self
     
     async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
-        """Close aiohttp session on context manager exit."""
         if self.session:
             await self.session.close()
             self.session = None
 
-    @retry(stop=stop_after_attempt(Config.MAX_RETRIES), wait=wait_exponential(multiplier=1, min=4, max=10))
+    @retry(
+        stop=stop_after_attempt(Config.MAX_RETRIES),
+        wait=wait_exponential(multiplier=1, min=4, max=10)
+    )
     async def _fetch_data(self, params: Dict[str, str]) -> Dict[str, Any]:
         """Fetch data from Etherscan with retry logic."""
         if not self.session:
@@ -113,17 +116,16 @@ class WalletAnalyzer:
         
         await self._rate_limit()
         params["apikey"] = self.api_key
-        
         try:
             async with self.session.get(self.base_url, params=params) as response:
                 if response.status != 200:
-                    raise APIError(f"API request failed: {response.status}")
+                    raise APIError(f"Etherscan request failed: {response.status}")
                 data = await response.json()
                 if data.get("status") == "0":
-                    err_msg = data.get("message", "Unknown Etherscan error")
-                    if "Max rate limit reached" in err_msg:
-                        raise APIError("Rate limit exceeded")
-                    raise APIError(f"API error: {err_msg}")
+                    err = data.get("message", "Unknown Etherscan error")
+                    if "Max rate limit reached" in err:
+                        raise APIError("Etherscan rate limit exceeded")
+                    raise APIError(f"Etherscan error: {err}")
                 return data
         except aiohttp.ClientError as e:
             raise APIError(f"Network error: {str(e)}")
@@ -131,7 +133,7 @@ class WalletAnalyzer:
             raise APIError(f"Unexpected error: {str(e)}")
 
     async def _rate_limit(self) -> None:
-        """Simple rate limit for Etherscan free tier requests."""
+        """Simple rate limiting for Etherscan free tier."""
         now = time.time()
         diff = now - self.last_request_time
         if diff < Config.RATE_LIMIT_DELAY:
@@ -145,30 +147,30 @@ class WalletAnalyzer:
 
     async def get_portfolio_data(self, address: str) -> WalletData:
         """
-        Get entire portfolio data:
-         - ETH balance
-         - ERC-20 tokens
-         - NFT collections (with contract + token_ids)
+        Return a dictionary with:
+         - address
+         - last_updated
+         - eth_balance
+         - tokens (list of ERC-20)
+         - nft_collections (with contract & token_id)
         """
         if not self._validate_address(address):
             raise ValidationError(f"Invalid Ethereum address: {address}")
 
         logger.info(f"Fetching portfolio data for {address}")
-        
         eth_balance = await self._get_eth_balance(address)
-        token_holdings = await self._get_token_holdings(address)
-        nft_collections = await self._get_nft_holdings(address)
-        
+        tokens = await self._get_token_holdings(address)
+        nft_colls = await self._get_nft_holdings(address)
+
         return {
             "address": address,
             "last_updated": datetime.now().isoformat(),
             "eth_balance": float(eth_balance),
-            "tokens": token_holdings,
-            "nft_collections": nft_collections
+            "tokens": tokens,
+            "nft_collections": nft_colls
         }
 
     async def _get_eth_balance(self, address: str) -> Decimal:
-        """Get ETH balance for address."""
         params = {
             "module": "account",
             "action": "balance",
@@ -179,10 +181,7 @@ class WalletAnalyzer:
         return Decimal(data["result"]) / Decimal("1e18")
 
     async def _get_token_holdings(self, address: str) -> List[Dict[str, Any]]:
-        """
-        Retrieve ERC-20 token balances for the address by
-        summing inbound/outbound transactions.
-        """
+        """Fetch ERC-20 tokens for address."""
         params = {
             "module": "account",
             "action": "tokentx",
@@ -193,19 +192,19 @@ class WalletAnalyzer:
         
         token_map: Dict[str, Dict[str, Any]] = {}
         for tx in data.get("result", []):
-            caddr = tx["contractAddress"]
-            if caddr not in token_map:
-                token_map[caddr] = {
+            contract = tx["contractAddress"]
+            if contract not in token_map:
+                token_map[contract] = {
                     "name": tx["tokenName"],
                     "symbol": tx["tokenSymbol"],
                     "decimals": int(tx["tokenDecimal"]),
                     "balance": Decimal(0)
                 }
-            amount = Decimal(tx["value"]) / Decimal(10 ** token_map[caddr]["decimals"])
+            amount = Decimal(tx["value"]) / Decimal(10 ** token_map[contract]["decimals"])
             if tx["to"].lower() == address.lower():
-                token_map[caddr]["balance"] += amount
+                token_map[contract]["balance"] += amount
             elif tx["from"].lower() == address.lower():
-                token_map[caddr]["balance"] -= amount
+                token_map[contract]["balance"] -= amount
 
         return [
             {
@@ -218,23 +217,8 @@ class WalletAnalyzer:
 
     async def _get_nft_holdings(self, address: str) -> Dict[str, Any]:
         """
-        Retrieve NFT holdings for address. 
-        We store them in a structure that includes contract addresses + token_ids
-        so we can fetch images from OpenSea easily.
-
-        Example structure:
-        {
-          "collections": [
-            {
-              "collection_name": "Bored Ape Yacht Club",
-              "items": [
-                 { "contract": "0xbc4ca0e...", "token_id": "1234" },
-                 ...
-              ]
-            },
-            ...
-          ]
-        }
+        Return { "collections": [ { "collection_name": str, "items": [{ contract, token_id }, ...] }, ... ] }
+        so we can fetch images from OpenSea by contract + token_id.
         """
         params = {
             "module": "account",
@@ -243,99 +227,92 @@ class WalletAnalyzer:
             "sort": "desc"
         }
         data = await self._fetch_data(params)
-        
+
         if data.get("status") != "1" or "result" not in data:
             return {"collections": []}
 
-        # We'll track them by (collection_name -> list of NFT {contract, token_id})
-        coll_map: Dict[str, List[Dict[str, str]]] = {}
-        # We'll also track what's currently owned
-        ownership = {}
-
+        # track final ownership
+        ownership_map = {}  # key = (contract_token), value = {contract, token_id, coll_name}
         for tx in data["result"]:
             contract = tx["contractAddress"]
             coll_name = tx.get("tokenName", "Unknown Collection")
             token_id = tx["tokenID"]
             key = f"{contract}_{token_id}"
 
-            # If received
             if tx["to"].lower() == address.lower():
-                ownership[key] = {
+                ownership_map[key] = {
                     "contract": contract,
-                    "collection_name": coll_name,
-                    "token_id": token_id
+                    "token_id": token_id,
+                    "collection_name": coll_name
                 }
-            # If sent out
             elif tx["from"].lower() == address.lower():
-                if key in ownership:
-                    ownership.pop(key, None)
-
-        # Group by collection
-        for entry in ownership.values():
-            coll = entry["collection_name"]
-            if coll not in coll_map:
-                coll_map[coll] = []
-            coll_map[coll].append({
-                "contract": entry["contract"],
-                "token_id": entry["token_id"]
+                if key in ownership_map:
+                    ownership_map.pop(key, None)
+
+        # group by collection_name
+        coll_dict: Dict[str, List[Dict[str, str]]] = {}
+        for item in ownership_map.values():
+            c_name = item["collection_name"]
+            if c_name not in coll_dict:
+                coll_dict[c_name] = []
+            coll_dict[c_name].append({
+                "contract": item["contract"],
+                "token_id": item["token_id"]
             })
-        
-        # Convert to list form
+
         collections_out = []
-        for c_name, items in coll_map.items():
+        for cname, items in coll_dict.items():
             collections_out.append({
-                "collection_name": c_name,
+                "collection_name": cname,
                 "items": items
             })
 
-        return {"collections": collections_out}
-#########################
-# OPENSEA IMAGE FETCHING
-#########################
-
-OPENSEA_API_BASE = "https://api.opensea.io/api/v2/chain/ethereum/contract"
+        return { "collections": collections_out }
+#### Part 3: OpenSea & PIL Conversion
 
 async def fetch_nft_metadata(opensea_key: str, contract: str, token_id: str) -> Dict[str, Any]:
     """
-    Fetch NFT metadata (name, image_url) from OpenSea.
-    Returns { "name": ..., "image_url": ... } or { "error": ... }
+    Fetch NFT metadata (including image_url) from OpenSea v2.
+    Returns { "name": str, "image_url": str } or { "error": str }
     """
-    url = f"{OPENSEA_API_BASE}/{contract}/nfts/{token_id}"
+    url = f"https://api.opensea.io/api/v2/chain/ethereum/contract/{contract}/nfts/{token_id}"
     headers = {"X-API-KEY": opensea_key} if opensea_key else {}
-
     async with aiohttp.ClientSession() as session:
-        async with session.get(url, headers=headers) as response:
-            if response.status == 403:
-                return {"error": "403 Forbidden: Invalid or restricted OpenSea API key."}
-            if response.status == 404:
+        async with session.get(url, headers=headers) as resp:
+            if resp.status == 403:
+                return {"error": "403 Forbidden: OpenSea API key issue"}
+            if resp.status == 404:
                 return {"error": f"404 Not Found: {contract} #{token_id}"}
-
             try:
-                data = await response.json()
+                data = await resp.json()
             except Exception as e:
-                return {"error": f"Invalid JSON from OpenSea: {str(e)}"}
-    
-    # parse data
+                return {"error": f"OpenSea JSON parse error: {str(e)}"}
+
     nft_obj = data.get("nft", {})
     name = nft_obj.get("name", f"NFT #{token_id}")
     image_url = nft_obj.get("image_url", "")
     return {"name": name, "image_url": image_url}
 
-def convert_image_to_base64(image_bytes: bytes) -> str:
+async def fetch_image_as_pil(url: str) -> Optional[Image.Image]:
     """
-    Convert raw image bytes to a base64 data URI that Gradio can display.
+    Download an image from a URL and return as a PIL Image.
     """
+    if not url:
+        return None
+    async with aiohttp.ClientSession() as session:
+        async with session.get(url) as resp:
+            if resp.status != 200:
+                return None
+            raw_bytes = await resp.read()
     try:
-        img = Image.open(BytesIO(image_bytes))
-        buffer = BytesIO()
-        img.save(buffer, format="PNG")
-        b64_str = base64.b64encode(buffer.getvalue()).decode()
-        return f"data:image/png;base64,{b64_str}"
+        return Image.open(BytesIO(raw_bytes))
     except Exception as e:
-        logger.error(f"Error converting image to base64: {e}")
-        return ""
+        logger.error(f"Error converting to PIL: {e}")
+        return None
+#### Part 4: ChatInterface
+
 class ChatInterface:
-    """Handles the chat + wallet + NFT images integration."""
+    """Handles chat logic with Etherscan (wallet data), OpenSea (NFT images), and OpenAI (chat)."""
     
     def __init__(self, openai_key: str, etherscan_key: str, opensea_key: str):
         self.openai_key = openai_key
@@ -347,8 +324,7 @@ class ChatInterface:
     @staticmethod
     def _validate_api_keys(openai_key: str, etherscan_key: str, opensea_key: str) -> Tuple[bool, str]:
         """
-        Validate OpenAI + Etherscan keys with sample calls, 
-        and check if OpenSea key is non-empty.
+        Validate all keys. We'll do a minimal check for OpenSea (non-empty).
         """
         try:
             # Check OpenAI
@@ -358,7 +334,7 @@ class ChatInterface:
                 messages=[{"role": "user", "content": "test"}],
                 max_tokens=1
             )
-            
+
             # Check Etherscan
             async def check_etherscan():
                 async with WalletAnalyzer(etherscan_key) as analyzer:
@@ -366,138 +342,120 @@ class ChatInterface:
                     await analyzer._fetch_data(params)
             asyncio.run(check_etherscan())
 
+            # Check OpenSea
             if not opensea_key.strip():
-                return False, "OpenSea API key cannot be empty."
+                return False, "OpenSea API key is empty!"
 
-            return True, "All API keys are valid!"
+            return True, "All API keys validated!"
         except Exception as e:
             return False, f"API key validation failed: {str(e)}"
 
     def _format_context_message(self) -> str:
-        """Format the wallet data in the system prompt for GPT context."""
+        """Format wallet data for GPT system prompt."""
         lines = []
         if not self.context:
             return ""
-        lines.append("Current Wallet Data:\n")
-        for addr, data in self.context.items():
-            lines.append(f"Wallet {addr[:8]}...{addr[-6:]}")
-            lines.append(f"  ETH Balance: {data['eth_balance']:.4f}")
-            lines.append(f"  ERC-20 Tokens: {len(data['tokens'])}")
+        lines.append("Current Wallet Data:")
+        for addr, wdata in self.context.items():
+            lines.append(f"Wallet {addr[:8]}...{addr[-6:]}:")
+            lines.append(f" ETH Balance: {wdata['eth_balance']:.4f}")
+            lines.append(f" # of Tokens: {len(wdata['tokens'])}")
             # NFT aggregator
-            coll_struct = data["nft_collections"]
-            if "collections" in coll_struct:
-                lines.append("  NFT Collections:")
-                for col in coll_struct["collections"]:
-                    lines.append(f"   * {col['collection_name']}: {len(col['items'])} NFT(s)")
+            nft_data = wdata["nft_collections"]
+            if "collections" in nft_data:
+                lines.append(" NFT Collections:")
+                for c in nft_data["collections"]:
+                    lines.append(f"  - {c['collection_name']}: {len(c['items'])} NFT(s)")
         return "\n".join(lines)
 
     async def process_message(
         self,
         message: str,
         history: Optional[ChatHistory] = None
-    ) -> Tuple[ChatHistory, Dict[str, Any], List[str]]:
+    ) -> Tuple[ChatHistory, Dict[str, Any], List[Image.Image]]:
         """
         1) Detect Ethereum address
-        2) Fetch wallet data from Etherscan
-        3) For each NFT, fetch from OpenSea + convert to base64
-        4) Return image data + chat response
+        2) Etherscan for wallet data
+        3) For each NFT, fetch from OpenSea, convert to PIL
+        4) Return images + chat response
         """
         if history is None:
             history = []
+
         if not message.strip():
             return history, self.context, []
 
-        # Attempt address detection
         match = re.search(Config.ETHEREUM_ADDRESS_REGEX, message)
-        nft_images_base64: List[str] = []
+        nft_images: List[Image.Image] = []
 
         if match:
             eth_address = match.group(0)
-            partial_info = f"Analyzing {eth_address}..."
-            history.append((message, partial_info))
+            # partial
+            partial_text = f"Analyzing {eth_address}..."
+            history.append((message, partial_text))
 
             try:
-                # Grab wallet data
+                # Etherscan
                 async with WalletAnalyzer(self.etherscan_key) as analyzer:
                     wallet_data = await analyzer.get_portfolio_data(eth_address)
-
                 self.context[eth_address] = wallet_data
 
                 # Summaries
-                eth_bal = wallet_data["eth_balance"]
-                token_count = len(wallet_data["tokens"])
-                nft_data = wallet_data["nft_collections"]
-                
-                # Summarize
                 lines = [
                     f"📊 Summary for {eth_address[:8]}...{eth_address[-6:]}",
-                    f"ETH: {eth_bal:.4f}",
-                    f"Tokens: {token_count}"
+                    f"ETH: {wallet_data['eth_balance']:.4f}",
+                    f"Tokens: {len(wallet_data['tokens'])}"
                 ]
-                total_nft_count = 0
-                # If we have "collections"
-                if "collections" in nft_data:
-                    for col in nft_data["collections"]:
-                        total_nft_count += len(col["items"])
-                    lines.append(f"NFTs: {total_nft_count}")
-
-                # Append summary
+                nft_info = wallet_data["nft_collections"]
+                total_nfts = 0
+                if "collections" in nft_info:
+                    for c in nft_info["collections"]:
+                        total_nfts += len(c["items"])
+                lines.append(f"NFTs: {total_nfts}")
+                # Add to chat
                 history.append((message, "\n".join(lines)))
 
-                # Fetch NFT images (limit for demonstration)
-                # For each collection, let's do up to 2 NFTs
-                # We'll do 2 collections max as well, to avoid heavy rate usage
-                if "collections" in nft_data:
-                    for col in nft_data["collections"][:2]:
-                        for nft_obj in col["items"][:2]:
-                            contract = nft_obj["contract"]
-                            token_id = nft_obj["token_id"]
-                            
-                            # 1) fetch metadata from OpenSea
-                            meta = await fetch_nft_metadata(self.opensea_key, contract, token_id)
+                # If we want to fetch images for, say, first 2 collections, 2 items each
+                if "collections" in nft_info:
+                    for coll in nft_info["collections"][:2]:
+                        for item in coll["items"][:2]:
+                            # 1) fetch metadata
+                            meta = await fetch_nft_metadata(self.opensea_key, item["contract"], item["token_id"])
                             if "error" in meta:
-                                logger.warning(f"Failed to fetch NFT: {meta['error']}")
+                                logger.warning(f"OpenSea metadata error: {meta['error']}")
                                 continue
                             image_url = meta["image_url"]
                             if not image_url:
-                                logger.info(f"No image for NFT {contract}#{token_id}")
+                                logger.info(f"No image for {meta['name']}")
                                 continue
+                            # 2) Download + convert to PIL
+                            pil_img = await fetch_image_as_pil(image_url)
+                            if pil_img:
+                                nft_images.append(pil_img)
+                                # Also mention in chat
+                                found_msg = f"Found NFT image: {meta['name']} (contract {item['contract'][:8]}...{item['contract'][-6:]}, id={item['token_id']})"
+                                history.append((message, found_msg))
 
-                            # 2) Download image & convert to base64
-                            async with aiohttp.ClientSession() as session:
-                                try:
-                                    async with session.get(image_url) as resp:
-                                        if resp.status == 200:
-                                            raw_img = await resp.read()
-                                            img_b64 = convert_image_to_base64(raw_img)
-                                            if img_b64:
-                                                nft_images_base64.append(img_b64)
-                                                # Also reflect in chat
-                                                found_msg = f"Found NFT: {meta['name']} (Contract {contract[:8]}...{contract[-6:]}, ID {token_id})"
-                                                history.append((message, found_msg))
-                                        else:
-                                            logger.warning(f"Image fetch failed: {resp.status}")
-                                except Exception as e:
-                                    logger.error(f"Error fetching NFT image: {e}")
             except Exception as e:
-                err_msg = f"Error analyzing wallet: {str(e)}"
-                logger.error(err_msg)
-                history.append((message, err_msg))
+                err = f"Error analyzing {eth_address}: {str(e)}"
+                logger.error(err)
+                history.append((message, err))
 
-        # Generate response from GPT
+        # Now do an OpenAI chat
         try:
             context_str = self._format_context_message()
-            # Convert local chat history -> OpenAI format
+            # Convert chat to OpenAI format
             limit = Config.HISTORY_LIMIT
-            truncated = history[-limit:]
+            short_hist = history[-limit:]
             openai_msgs = []
-            for umsg, amsg in truncated:
-                openai_msgs.append({"role": "user", "content": umsg})
-                openai_msgs.append({"role": "assistant", "content": amsg})
+            for usr, ans in short_hist:
+                openai_msgs.append({"role": "user", "content": usr})
+                openai_msgs.append({"role": "assistant", "content": ans})
 
+            # openai
             openai.api_key = self.openai_key
             client = openai.OpenAI(api_key=self.openai_key)
-            response = client.chat.completions.create(
+            resp = client.chat.completions.create(
                 model=Config.OPENAI_MODEL,
                 messages=[
                     {"role": "system", "content": Config.SYSTEM_PROMPT},
@@ -508,23 +466,25 @@ class ChatInterface:
                 temperature=Config.TEMPERATURE,
                 max_tokens=Config.MAX_TOKENS
             )
-            final_answer = response.choices[0].message.content
-            history.append((message, final_answer))
+            final_msg = resp.choices[0].message.content
+            history.append((message, final_msg))
+
+            return history, self.context, nft_images
 
-            # Return new chat, new context, and the collected base64 images
-            return history, self.context, nft_images_base64
         except Exception as e:
-            logger.error(f"OpenAI Error: {e}")
-            err_resp = f"OpenAI error: {str(e)}"
-            history.append((message, err_resp))
+            logger.error(f"OpenAI error: {e}")
+            err_chat = f"OpenAI error: {str(e)}"
+            history.append((message, err_chat))
             return history, self.context, []
 
     def clear_context(self) -> Tuple[Dict[str, Any], List[Tuple[str, str]]]:
-        """Clear wallet context + chat."""
+        """Clear the wallet context and chat."""
         self.context = {}
         return {}, []
+#### Part 5: Gradio UI & Launch
+
 class GradioInterface:
-    """Manages the Gradio UI for Hugging Face Space, with NFT images top-right."""
+    """Manages the Gradio UI for Hugging Face Space, with top-right NFT gallery as PIL images."""
     
     def __init__(self):
         self.chat_interface: Optional[ChatInterface] = None
@@ -533,12 +493,11 @@ class GradioInterface:
     def _create_interface(self) -> gr.Blocks:
         with gr.Blocks(theme=gr.themes.Soft()) as demo:
             gr.Markdown("""
-            # 🐕 LOSS DOG: Blockchain Wallet Analyzer (NFT Images at Top-Right)
-
-            **Instructions**:
-            - Enter your **OpenAI**, **Etherscan**, and **OpenSea** API keys below.
-            - Validate them, then chat with your Ethereum wallet address.
-            - NFT images (if any) will appear as **base64** images in the top-right gallery.
+            # 🐕 LOSS DOG: Blockchain Wallet Analyzer (NFT Images, top-right)
+            
+            - Enter your **OpenAI**, **Etherscan**, **OpenSea** keys below
+            - Validate, then chat with your Ethereum address
+            - NFT images will appear as **PIL** objects in the top-right
             """)
 
             with gr.Row():
@@ -558,55 +517,47 @@ class GradioInterface:
                     placeholder="Enter your OpenSea API key..."
                 )
 
-            validate_status = gr.Textbox(label="Validation Status", interactive=False)
-            validate_btn = gr.Button("Validate API Keys")
+            validation_status = gr.Textbox(label="Validation Status", interactive=False)
+            validate_btn = gr.Button("Validate API Keys", variant="primary")
 
             with gr.Row():
                 # Main Chat Column
                 with gr.Column(scale=2):
-                    chatbot = gr.Chatbot(
-                        label="Chat History",
-                        height=420,
-                        value=[]
-                    )
+                    chatbot = gr.Chatbot(label="Chat History", height=420, value=[])
                     with gr.Row():
                         msg_input = gr.Textbox(
                             label="Message",
-                            placeholder="Enter an Ethereum address or question..."
+                            placeholder="Enter an ETH address or question..."
                         )
-                        send_btn = gr.Button("Send")
-
-                # NFT Gallery on Top-Right
+                        send_btn = gr.Button("Send", variant="primary")
+                # Top-Right: NFT Gallery (PIL images)
                 with gr.Column(scale=1):
                     nft_gallery = gr.Gallery(
                         label="NFT Images (Top-Right)",
-                        columns=2,
-                        show_label=True
+                        columns=2
                     )
-                    # Then wallet context below
                     wallet_context = gr.JSON(
                         label="Active Wallet Context",
                         value={}
                     )
                     clear_btn = gr.Button("Clear Context")
 
-            # Initially disable chat until keys are validated
             msg_input.interactive = False
             send_btn.interactive = False
 
-            # Validation function
-            def validate_keys(openai_k: str, etherscan_k: str, opensea_k: str):
-                is_valid, message = ChatInterface._validate_api_keys(openai_k, etherscan_k, opensea_k)
+            def validate_keys(openai_k: str, etherscan_k: str, opensea_k: str) -> Tuple[str, gr.update, gr.update]:
+                """Validate user-provided keys; enable chat if all pass."""
+                is_valid, msg = ChatInterface._validate_api_keys(openai_k, etherscan_k, opensea_k)
                 if is_valid:
                     self.chat_interface = ChatInterface(openai_k, etherscan_k, opensea_k)
                     return (
-                        f"✅ {message}",
+                        f"✅ {msg}",
                         gr.update(interactive=True),
                         gr.update(interactive=True)
                     )
                 else:
                     return (
-                        f"❌ {message}",
+                        f"❌ {msg}",
                         gr.update(interactive=False),
                         gr.update(interactive=False)
                     )
@@ -614,11 +565,11 @@ class GradioInterface:
             validate_btn.click(
                 fn=validate_keys,
                 inputs=[openai_key, etherscan_key, opensea_key],
-                outputs=[validate_status, msg_input, send_btn]
+                outputs=[validation_status, msg_input, send_btn]
             )
 
-            # Clear context
             def clear_all():
+                """Clear wallet context and chat."""
                 if self.chat_interface:
                     return self.chat_interface.clear_context()
                 return {}, []
@@ -629,24 +580,26 @@ class GradioInterface:
                 outputs=[wallet_context, chatbot]
             )
 
-            # Async callback to handle chat
             async def handle_message(
                 message: str,
                 chat_hist: List[Tuple[str, str]],
                 context: Dict[str, Any]
-            ) -> Tuple[List[Tuple[str, str]], Dict[str, Any], List[str]]:
-                """Process user message. Return updated chat, context, base64 images."""
+            ) -> Tuple[List[Tuple[str, str]], Dict[str, Any], List[Image.Image]]:
+                """Process user input, return updated chat, context, list of PIL images."""
                 if not self.chat_interface:
                     return [], {}, []
+                
                 try:
-                    new_hist, new_ctx, nft_imgs_b64 = await self.chat_interface.process_message(message, chat_hist)
-                    return new_hist, new_ctx, nft_imgs_b64
+                    new_history, new_context, images = await self.chat_interface.process_message(message, chat_hist)
+                    return new_history, new_context, images
                 except Exception as e:
                     logger.error(f"Error in handle_message: {e}")
+                    if chat_hist is None:
+                        chat_hist = []
                     chat_hist.append((message, f"Error: {str(e)}"))
                     return chat_hist, context, []
 
-            # Submit callback
+            # Chat flow
             msg_input.submit(
                 fn=handle_message,
                 inputs=[msg_input, chatbot, wallet_context],
@@ -657,7 +610,6 @@ class GradioInterface:
                 [msg_input]
             )
 
-            # Send button callback
             send_btn.click(
                 fn=handle_message,
                 inputs=[msg_input, chatbot, wallet_context],
@@ -675,10 +627,8 @@ class GradioInterface:
         self.demo.launch()
 
 def main():
-    """
-    Main entry point for the Hugging Face Space.
-    """
-    logger.info("Launching LOSS DOG with NFT image display (top-right).")
+    """Main entry point for Hugging Face Space."""
+    logger.info("Launching LOSS DOG with PIL-based NFT images (top-right).")
     interface = GradioInterface()
     interface.launch()