Update app.py

app.py

@@ -1,10 +1,12 @@
 """
-Blockchain Wallet Analyzer - A tool for analyzing Ethereum wallet contents and NFT holdings.
+Enhanced Blockchain Wallet Analyzer with Persistent State Management
+and Improved NFT Interactions
 
-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 using the OpenAI API, 
-and NFT image rendering from OpenSea via PIL objects.
+This version includes:
+- Persistent state management for wallet data
+- Enhanced NFT tracking and caching
+- Improved chat context management
+- Better error handling and recovery
 
 Author: Claude
 Date: January 2025
@@ -18,24 +20,28 @@ import json
 import time
 import logging
 import asyncio
+import pickle
 import base64
-from typing import List, Dict, Tuple, Any, Optional, TypeVar
-from datetime import datetime
+from typing import List, Dict, Tuple, Any, Optional, TypeVar, Set
+from datetime import datetime, timedelta
 from decimal import Decimal
-from dataclasses import dataclass
+from dataclasses import dataclass, asdict
 from pathlib import Path
 from io import BytesIO
+from threading import Lock
 
 import aiohttp
 import openai
 import gradio as gr
 from PIL import Image
 from tenacity import retry, stop_after_attempt, wait_exponential
+from cachetools import TTLCache, LRUCache
 
 # Type variables
 T = TypeVar('T')
 WalletData = Dict[str, Any]
 ChatHistory = List[Tuple[str, str]]
+NFTData = Dict[str, Any]
 
 # Configure logging
 logging.basicConfig(
@@ -48,6 +54,7 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Custom Exceptions
 class ConfigError(Exception):
     """Raised when there's an error in configuration."""
     pass
@@ -60,6 +67,10 @@ class ValidationError(Exception):
     """Raised when there's an error in input validation."""
     pass
 
+class CacheError(Exception):
+    """Raised when there's an error with the cache."""
+    pass
+
 @dataclass
 class Config:
     """Application configuration settings."""
@@ -69,32 +80,170 @@ class Config:
     Your personality:
     - Friendly and enthusiastic
     - Explain findings in fun, simple ways
-    - Provide NFT images from OpenSea when possible
+    - Track and remember NFTs across conversations
+    - Maintain context about wallets and collections
 
     Instructions:
-    - You have access to detailed wallet data in your context
-    - Use this data to provide specific answers about holdings
-    - Reference exact numbers and collections when discussing NFTs
-    - Compare wallets if multiple are available
+    - Use cached wallet data when available
+    - Reference specific NFTs by name and ID
+    - Compare wallets and collections over time
+    - Highlight interesting patterns or changes
     """
     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 per second
     MAX_RETRIES: int = 3
-    OPENAI_MODEL: str = "gpt-4o-mini"  # Example updated model
+    OPENAI_MODEL: str = "gpt-4o-mini"  # Updated model
     MAX_TOKENS: int = 4000
     TEMPERATURE: float = 0.7
-    HISTORY_LIMIT: int = 5
-#### Part 2: Wallet Analyzer (Etherscan)
+    HISTORY_LIMIT: int = 10  # Increased history
+    CACHE_DIR: str = ".blockchain_cache"
+    NFT_CACHE_TTL: int = 3600  # 1 hour
+    WALLET_CACHE_TTL: int = 1800  # 30 minutes
+    MAX_CHAT_SESSIONS: int = 50
+    MAX_NFTS_PER_COLLECTION: int = 5  # Increased from 2
+
+@dataclass
+class NFTMetadata:
+    """Stores NFT metadata with cache management."""
+    contract: str
+    token_id: str
+    name: str
+    image_url: str
+    collection_name: str
+    last_updated: datetime
+    image_data: Optional[bytes] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "contract": self.contract,
+            "token_id": self.token_id,
+            "name": self.name,
+            "image_url": self.image_url,
+            "collection_name": self.collection_name,
+            "last_updated": self.last_updated.isoformat(),
+            "image_data": base64.b64encode(self.image_data).decode() if self.image_data else None
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'NFTMetadata':
+        """Create from dictionary after deserialization."""
+        if data.get("image_data"):
+            data["image_data"] = base64.b64decode(data["image_data"])
+        data["last_updated"] = datetime.fromisoformat(data["last_updated"])
+        return cls(**data)
+
+class StateManager:
+    """Manages persistent state for the application."""
+    
+    def __init__(self, cache_dir: str = Config.CACHE_DIR):
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(exist_ok=True)
+        
+        # In-memory caches with TTL
+        self.nft_cache = TTLCache(maxsize=1000, ttl=Config.NFT_CACHE_TTL)
+        self.wallet_cache = TTLCache(maxsize=100, ttl=Config.WALLET_CACHE_TTL)
+        self.chat_history_cache = LRUCache(maxsize=Config.MAX_CHAT_SESSIONS)
+        
+        # Thread safety for cache operations
+        self._lock = Lock()
+        
+        # Load existing cache data
+        self._load_cache()
+    
+    def _load_cache(self):
+        """Load cached data from disk with error handling."""
+        try:
+            # Load NFT cache
+            nft_cache_file = self.cache_dir / "nft_cache.pkl"
+            if nft_cache_file.exists():
+                with open(nft_cache_file, "rb") as f:
+                    cached_data = pickle.load(f)
+                    now = datetime.now()
+                    for key, value in cached_data.items():
+                        if isinstance(value, NFTMetadata):
+                            if now - value.last_updated < timedelta(hours=1):
+                                self.nft_cache[key] = value
+            
+            # Load wallet cache
+            wallet_cache_file = self.cache_dir / "wallet_cache.pkl"
+            if wallet_cache_file.exists():
+                with open(wallet_cache_file, "rb") as f:
+                    cached_data = pickle.load(f)
+                    now = datetime.now()
+                    for key, value in cached_data.items():
+                        if now - value.get("last_updated", now) < timedelta(minutes=30):
+                            self.wallet_cache[key] = value
+        except Exception as e:
+            logger.error(f"Error loading cache: {e}")
+            # Continue with empty cache if load fails
+    
+    def save_cache(self):
+        """Save current cache to disk with error handling."""
+        try:
+            with self._lock:
+                # Save NFT cache
+                nft_cache_file = self.cache_dir / "nft_cache.pkl"
+                with open(nft_cache_file, "wb") as f:
+                    pickle.dump(dict(self.nft_cache), f)
+                
+                # Save wallet cache
+                wallet_cache_file = self.cache_dir / "wallet_cache.pkl"
+                with open(wallet_cache_file, "wb") as f:
+                    pickle.dump(dict(self.wallet_cache), f)
+        except Exception as e:
+            logger.error(f"Error saving cache: {e}")
+            raise CacheError(f"Failed to save cache: {str(e)}")
+
+    def get_nft_metadata(self, contract: str, token_id: str) -> Optional[NFTMetadata]:
+        """Get NFT metadata from cache if available."""
+        cache_key = f"{contract}_{token_id}"
+        return self.nft_cache.get(cache_key)
+
+    def store_nft_metadata(self, metadata: NFTMetadata):
+        """Store NFT metadata in cache with immediate disk sync."""
+        cache_key = f"{metadata.contract}_{metadata.token_id}"
+        with self._lock:
+            self.nft_cache[cache_key] = metadata
+            self.save_cache()
+
+    def get_wallet_data(self, address: str) -> Optional[WalletData]:
+        """Get wallet data from cache if available."""
+        return self.wallet_cache.get(address.lower())
+
+    def store_wallet_data(self, address: str, data: WalletData):
+        """Store wallet data in cache with timestamp."""
+        with self._lock:
+            data["last_updated"] = datetime.now()
+            self.wallet_cache[address.lower()] = data
+            self.save_cache()
+
+    def get_chat_history(self, session_id: str) -> ChatHistory:
+        """Get chat history for a session."""
+        return self.chat_history_cache.get(session_id, [])
+
+    def store_chat_history(self, session_id: str, history: ChatHistory):
+        """Store chat history for a session."""
+        with self._lock:
+            self.chat_history_cache[session_id] = history
+
+    def clear_session(self, session_id: str):
+        """Clear all data for a specific session."""
+        with self._lock:
+            if session_id in self.chat_history_cache:
+                del self.chat_history_cache[session_id]
 
 class WalletAnalyzer:
-    """Analyzes Ethereum wallet contents using Etherscan API."""
+    """Analyzes Ethereum wallet contents with improved state management."""
     
-    def __init__(self, api_key: str):
+    def __init__(self, api_key: str, state_manager: StateManager):
         self.api_key = api_key
+        self.state_manager = state_manager
         self.base_url = Config.ETHERSCAN_BASE_URL
         self.session: Optional[aiohttp.ClientSession] = None
         self.last_request_time = 0
+        self.active_collections: Set[str] = set()
     
     async def __aenter__(self) -> WalletAnalyzer:
         self.session = aiohttp.ClientSession()
@@ -110,78 +259,125 @@ class WalletAnalyzer:
         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."""
+        """Fetch data from Etherscan with enhanced error handling."""
         if not self.session:
             raise APIError("No active session. Use context manager.")
         
         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 == 429:  # Too Many Requests
+                    raise APIError("Rate limit exceeded. Please wait.")
                 if response.status != 200:
                     raise APIError(f"Etherscan request failed: {response.status}")
+                
                 data = await response.json()
                 if data.get("status") == "0":
-                    err = data.get("message", "Unknown Etherscan error")
-                    if "Max rate limit reached" in err:
+                    err_msg = data.get("message", "Unknown Etherscan error")
+                    if "Max rate limit reached" in err_msg:
                         raise APIError("Etherscan rate limit exceeded")
-                    raise APIError(f"Etherscan error: {err}")
+                    if "Invalid API Key" in err_msg:
+                        raise ConfigError("Invalid Etherscan API key")
+                    raise APIError(f"Etherscan error: {err_msg}")
+                
                 return data
+                
         except aiohttp.ClientError as e:
             raise APIError(f"Network error: {str(e)}")
+        except json.JSONDecodeError as e:
+            raise APIError(f"Invalid JSON response: {str(e)}")
         except Exception as e:
             raise APIError(f"Unexpected error: {str(e)}")
 
     async def _rate_limit(self) -> None:
-        """Simple rate limiting for Etherscan free tier."""
+        """Enhanced rate limiting with dynamic backoff."""
         now = time.time()
         diff = now - self.last_request_time
         if diff < Config.RATE_LIMIT_DELAY:
-            await asyncio.sleep(Config.RATE_LIMIT_DELAY - diff)
+            delay = Config.RATE_LIMIT_DELAY - diff
+            logger.debug(f"Rate limiting: waiting {delay:.2f}s")
+            await asyncio.sleep(delay)
         self.last_request_time = time.time()
 
     @staticmethod
     def _validate_address(address: str) -> bool:
-        """Validate Ethereum address format."""
-        return bool(re.match(Config.ETHEREUM_ADDRESS_REGEX, address))
+        """Validate Ethereum address format with checksum."""
+        if not re.match(Config.ETHEREUM_ADDRESS_REGEX, address):
+            return False
+        # Additional validation could be added here
+        return True
 
     async def get_portfolio_data(self, address: str) -> WalletData:
         """
-        Return a dictionary with:
+        Get comprehensive wallet portfolio data with caching.
+        
+        Returns a dictionary with:
          - address
          - last_updated
          - eth_balance
          - tokens (list of ERC-20)
          - nft_collections (with contract & token_id)
+         - analytics (additional metrics)
         """
         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)
-        tokens = await self._get_token_holdings(address)
-        nft_colls = await self._get_nft_holdings(address)
+        # Check cache first
+        cached_data = self.state_manager.get_wallet_data(address)
+        if cached_data:
+            logger.info(f"Using cached data for {address}")
+            return cached_data
 
-        return {
-            "address": address,
-            "last_updated": datetime.now().isoformat(),
-            "eth_balance": float(eth_balance),
-            "tokens": tokens,
-            "nft_collections": nft_colls
-        }
+        logger.info(f"Fetching fresh portfolio data for {address}")
+        
+        try:
+            # Parallel fetching of data
+            eth_balance, tokens, nft_colls = await asyncio.gather(
+                self._get_eth_balance(address),
+                self._get_token_holdings(address),
+                self._get_nft_holdings(address)
+            )
+            
+            # Calculate additional analytics
+            analytics = self._calculate_analytics(eth_balance, tokens, nft_colls)
+            
+            portfolio_data = {
+                "address": address,
+                "last_updated": datetime.now().isoformat(),
+                "eth_balance": float(eth_balance),
+                "tokens": tokens,
+                "nft_collections": nft_colls,
+                "analytics": analytics
+            }
+            
+            # Cache the results
+            self.state_manager.store_wallet_data(address, portfolio_data)
+            
+            return portfolio_data
+            
+        except Exception as e:
+            logger.error(f"Error fetching portfolio data: {e}")
+            raise
 
     async def _get_eth_balance(self, address: str) -> Decimal:
+        """Get ETH balance with proper decimal handling."""
         params = {
             "module": "account",
             "action": "balance",
             "address": address,
             "tag": "latest"
         }
-        data = await self._fetch_data(params)
-        return Decimal(data["result"]) / Decimal("1e18")
+        try:
+            data = await self._fetch_data(params)
+            balance_wei = Decimal(data["result"])
+            return balance_wei / Decimal("1e18")  # Convert to ETH
+        except (KeyError, ValueError) as e:
+            raise APIError(f"Invalid balance data: {str(e)}")
 
     async def _get_token_holdings(self, address: str) -> List[Dict[str, Any]]:
-        """Fetch ERC-20 tokens for address."""
+        """Get ERC-20 tokens with improved tracking."""
         params = {
             "module": "account",
             "action": "tokentx",
@@ -191,34 +387,52 @@ class WalletAnalyzer:
         data = await self._fetch_data(params)
         
         token_map: Dict[str, Dict[str, Any]] = {}
-        for tx in data.get("result", []):
-            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)
+        try:
+            for tx in data.get("result", []):
+                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),
+                        "last_tx": datetime.fromtimestamp(int(tx["timeStamp"]))
+                    }
+                
+                amount = Decimal(tx["value"]) / Decimal(10 ** token_map[contract]["decimals"])
+                if tx["to"].lower() == address.lower():
+                    token_map[contract]["balance"] += amount
+                else:
+                    token_map[contract]["balance"] -= amount
+
+            # Filter and format results
+            return [
+                {
+                    "name": v["name"],
+                    "symbol": v["symbol"],
+                    "balance": float(v["balance"]),
+                    "last_activity": v["last_tx"].isoformat()
                 }
-            amount = Decimal(tx["value"]) / Decimal(10 ** token_map[contract]["decimals"])
-            if tx["to"].lower() == address.lower():
-                token_map[contract]["balance"] += amount
-            elif tx["from"].lower() == address.lower():
-                token_map[contract]["balance"] -= amount
-
-        return [
-            {
-                "name": v["name"],
-                "symbol": v["symbol"],
-                "balance": float(v["balance"])
-            }
-            for v in token_map.values() if v["balance"] > 0
-        ]
+                for v in token_map.values()
+                if v["balance"] > 0
+            ]
+            
+        except Exception as e:
+            logger.error(f"Error processing token data: {e}")
+            raise APIError(f"Token processing error: {str(e)}")
 
     async def _get_nft_holdings(self, address: str) -> Dict[str, Any]:
         """
-        Return { "collections": [ { "collection_name": str, "items": [{ contract, token_id }, ...] }, ... ] }
-        so we can fetch images from OpenSea by contract + token_id.
+        Enhanced NFT holdings tracking with collection stats.
+        Returns {
+            "collections": [
+                {
+                    "collection_name": str,
+                    "items": [{"contract": str, "token_id": str}],
+                    "stats": {"count": int, "first_acquired": datetime}
+                }
+            ]
+        }
         """
         params = {
             "module": "account",
@@ -231,375 +445,718 @@ class WalletAnalyzer:
         if data.get("status") != "1" or "result" not in data:
             return {"collections": []}
 
-        # 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 tx["to"].lower() == address.lower():
-                ownership_map[key] = {
-                    "contract": contract,
-                    "token_id": token_id,
-                    "collection_name": coll_name
-                }
-            elif tx["from"].lower() == address.lower():
-                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"]
-            })
-
-        collections_out = []
-        for cname, items in coll_dict.items():
-            collections_out.append({
-                "collection_name": cname,
-                "items": items
-            })
-
-        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 (including image_url) from OpenSea v2.
-    Returns { "name": str, "image_url": str } or { "error": str }
-    """
-    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 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 resp.json()
-            except Exception as e:
-                return {"error": f"OpenSea JSON parse error: {str(e)}"}
+        try:
+            # Track ownership and collection stats
+            ownership_map = {}  # (contract, token_id) -> ownership_data
+            collection_stats = {}  # collection_name -> stats
+            
+            for tx in data["result"]:
+                contract = tx["contractAddress"]
+                coll_name = tx.get("tokenName", "Unknown Collection")
+                token_id = tx["tokenID"]
+                timestamp = datetime.fromtimestamp(int(tx["timeStamp"]))
+                key = (contract, token_id)
+
+                # Update collection stats
+                if coll_name not in collection_stats:
+                    collection_stats[coll_name] = {
+                        "count": 0,
+                        "first_acquired": timestamp,
+                        "last_activity": timestamp
+                    }
+
+                if tx["to"].lower() == address.lower():
+                    ownership_map[key] = {
+                        "contract": contract,
+                        "token_id": token_id,
+                        "collection_name": coll_name,
+                        "acquired": timestamp
+                    }
+                    collection_stats[coll_name]["count"] += 1
+                    if timestamp < collection_stats[coll_name]["first_acquired"]:
+                        collection_stats[coll_name]["first_acquired"] = timestamp
+                elif tx["from"].lower() == address.lower():
+                    if key in ownership_map:
+                        del ownership_map[key]
+                        collection_stats[coll_name]["count"] -= 1
+
+                collection_stats[coll_name]["last_activity"] = timestamp
+
+            # Format output
+            collections_out = []
+            for cname, stats in collection_stats.items():
+                if stats["count"] > 0:
+                    items = [
+                        {"contract": data["contract"], "token_id": data["token_id"]}
+                        for data in ownership_map.values()
+                        if data["collection_name"] == cname
+                    ]
+                    collections_out.append({
+                        "collection_name": cname,
+                        "items": items,
+                        "stats": {
+                            "count": stats["count"],
+                            "first_acquired": stats["first_acquired"].isoformat(),
+                            "last_activity": stats["last_activity"].isoformat()
+                        }
+                    })
+
+            return {"collections": collections_out}
 
-    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}
+        except Exception as e:
+            logger.error(f"Error processing NFT data: {e}")
+            raise APIError(f"NFT processing error: {str(e)}")
 
-async def fetch_image_as_pil(url: str) -> Optional[Image.Image]:
-    """
-    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:
-        return Image.open(BytesIO(raw_bytes))
-    except Exception as e:
-        logger.error(f"Error converting to PIL: {e}")
-        return None
-#### Part 4: ChatInterface
+    def _calculate_analytics(
+        self,
+        eth_balance: Decimal,
+        tokens: List[Dict[str, Any]],
+        nft_collections: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Calculate additional portfolio analytics."""
+        total_nfts = sum(
+            len(c["items"]) 
+            for c in nft_collections.get("collections", [])
+        )
+        
+        token_count = len(tokens)
+        collection_count = len(nft_collections.get("collections", []))
+        
+        return {
+            "total_assets": {
+                "eth": float(eth_balance),
+                "tokens": token_count,
+                "nfts": total_nfts,
+                "collections": collection_count
+            },
+            "last_calculated": datetime.now().isoformat()
+        }
 
-class ChatInterface:
-    """Handles chat logic with Etherscan (wallet data), OpenSea (NFT images), and OpenAI (chat)."""
+class NFTProcessor:
+    """Enhanced NFT processing with caching and rate limiting."""
     
-    def __init__(self, openai_key: str, etherscan_key: str, opensea_key: str):
-        self.openai_key = openai_key
-        self.etherscan_key = etherscan_key
+    def __init__(self, opensea_key: str, state_manager: StateManager):
         self.opensea_key = opensea_key
-        self.context: Dict[str, Any] = {}
-        openai.api_key = openai_key
+        self.state_manager = state_manager
+        self.session: Optional[aiohttp.ClientSession] = None
+        self.last_request_time = 0
+        self.rate_limit_delay = 0.25  # 4 requests per second for OpenSea
+        self._lock = Lock()
+        
+    async def __aenter__(self) -> 'NFTProcessor':
+        self.session = aiohttp.ClientSession()
+        return self
+        
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self.session:
+            await self.session.close()
+            self.session = None
 
-    @staticmethod
-    def _validate_api_keys(openai_key: str, etherscan_key: str, opensea_key: str) -> Tuple[bool, str]:
+    async def _rate_limit(self) -> None:
+        """OpenSea-specific rate limiting."""
+        now = time.time()
+        async with self._lock:
+            diff = now - self.last_request_time
+            if diff < self.rate_limit_delay:
+                await asyncio.sleep(self.rate_limit_delay - diff)
+            self.last_request_time = time.time()
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10)
+    )
+    async def fetch_nft_metadata(
+        self,
+        contract: str,
+        token_id: str,
+        collection_name: str
+    ) -> Optional[NFTMetadata]:
         """
-        Validate all keys. We'll do a minimal check for OpenSea (non-empty).
+        Fetch NFT metadata from OpenSea with caching and retry logic.
+        Returns cached data if available and fresh.
         """
+        # Check cache first
+        cached = self.state_manager.get_nft_metadata(contract, token_id)
+        if cached and (datetime.now() - cached.last_updated < timedelta(hours=1)):
+            logger.debug(f"Using cached metadata for {contract}_{token_id}")
+            return cached
+            
+        if not self.session:
+            raise APIError("No active session. Use context manager.")
+
+        await self._rate_limit()
+        
+        url = f"https://api.opensea.io/api/v2/chain/ethereum/contract/{contract}/nfts/{token_id}"
+        headers = {"X-API-KEY": self.opensea_key} if self.opensea_key else {}
+        
         try:
-            # Check OpenAI
-            client = openai.OpenAI(api_key=openai_key)
-            client.chat.completions.create(
-                model=Config.OPENAI_MODEL,
-                messages=[{"role": "user", "content": "test"}],
-                max_tokens=1
+            async with self.session.get(url, headers=headers) as resp:
+                if resp.status == 403:
+                    raise APIError("OpenSea API key invalid or expired")
+                if resp.status == 404:
+                    logger.warning(f"NFT not found: {contract} #{token_id}")
+                    return None
+                if resp.status == 429:
+                    raise APIError("OpenSea rate limit exceeded")
+                if resp.status != 200:
+                    raise APIError(f"OpenSea API error: {resp.status}")
+                    
+                data = await resp.json()
+                
+            nft_obj = data.get("nft", {})
+            if not nft_obj:
+                logger.warning(f"Empty NFT data for {contract} #{token_id}")
+                return None
+                
+            # Create metadata object
+            metadata = NFTMetadata(
+                contract=contract,
+                token_id=token_id,
+                name=nft_obj.get("name", f"NFT #{token_id}"),
+                image_url=nft_obj.get("image_url", ""),
+                collection_name=collection_name,
+                last_updated=datetime.now()
             )
+            
+            # Fetch and store image if available
+            if metadata.image_url:
+                image_data = await self._fetch_and_process_image(metadata.image_url)
+                metadata.image_data = image_data
+            
+            # Cache the metadata
+            self.state_manager.store_nft_metadata(metadata)
+            
+            return metadata
+            
+        except Exception as e:
+            logger.error(f"Error fetching NFT metadata: {e}")
+            return None
 
-            # Check Etherscan
-            async def check_etherscan():
-                async with WalletAnalyzer(etherscan_key) as analyzer:
-                    params = {"module": "stats", "action": "ethsupply"}
-                    await analyzer._fetch_data(params)
-            asyncio.run(check_etherscan())
+    async def _fetch_and_process_image(self, url: str) -> Optional[bytes]:
+        """
+        Fetch and process NFT image with enhanced error handling.
+        Returns image bytes suitable for PIL conversion.
+        """
+        if not url:
+            return None
+            
+        try:
+            async with self.session.get(url) as resp:
+                if resp.status != 200:
+                    logger.warning(f"Failed to fetch image: {url} ({resp.status})")
+                    return None
+                    
+                content_type = resp.headers.get("content-type", "")
+                if not content_type.startswith(("image/", "application/octet-stream")):
+                    logger.warning(f"Invalid content type: {content_type}")
+                    return None
+                    
+                raw_bytes = await resp.read()
+                
+                # Process with PIL to validate and optimize
+                try:
+                    img = Image.open(BytesIO(raw_bytes))
+                    
+                    # Convert to RGB if necessary
+                    if img.mode not in ("RGB", "RGBA"):
+                        img = img.convert("RGB")
+                    
+                    # Resize if too large while maintaining aspect ratio
+                    max_size = (800, 800)
+                    if img.size[0] > max_size[0] or img.size[1] > max_size[1]:
+                        img.thumbnail(max_size, Image.Resampling.LANCZOS)
+                    
+                    # Save as PNG
+                    output = BytesIO()
+                    img.save(output, format="PNG", optimize=True)
+                    return output.getvalue()
+                    
+                except Exception as e:
+                    logger.error(f"Error processing image: {e}")
+                    return None
+                    
+        except Exception as e:
+            logger.error(f"Error fetching image: {e}")
+            return None
 
-            # Check OpenSea
-            if not opensea_key.strip():
-                return False, "OpenSea API key is empty!"
+    async def batch_process_nfts(
+        self,
+        nft_collections: List[Dict[str, Any]],
+        max_per_collection: int = Config.MAX_NFTS_PER_COLLECTION
+    ) -> List[NFTMetadata]:
+        """
+        Process multiple NFTs in parallel with rate limiting.
+        Returns list of successfully processed NFT metadata.
+        """
+        all_tasks = []
+        processed_nfts = []
+        
+        async with self:  # Use context manager for session
+            for collection in nft_collections:
+                coll_name = collection["collection_name"]
+                items = collection["items"][:max_per_collection]  # Limit items per collection
+                
+                # Create tasks for parallel processing
+                tasks = [
+                    self.fetch_nft_metadata(
+                        item["contract"],
+                        item["token_id"],
+                        coll_name
+                    )
+                    for item in items
+                ]
+                
+                all_tasks.extend(tasks)
+                
+                # Process in batches to respect rate limits
+                batch_size = 4  # Process 4 at a time
+                for i in range(0, len(tasks), batch_size):
+                    batch = tasks[i:i + batch_size]
+                    results = await asyncio.gather(*batch, return_exceptions=True)
+                    
+                    # Filter successful results
+                    for result in results:
+                        if isinstance(result, NFTMetadata):
+                            processed_nfts.append(result)
+                        elif isinstance(result, Exception):
+                            logger.error(f"Error processing NFT: {result}")
+                
+        return processed_nfts
 
-            return True, "All API keys validated!"
-        except Exception as e:
-            return False, f"API key validation failed: {str(e)}"
+    def get_cached_images(self, nft_list: List[NFTMetadata]) -> List[Image.Image]:
+        """Convert cached image data to PIL Images."""
+        images = []
+        for nft in nft_list:
+            if nft.image_data:
+                try:
+                    img = Image.open(BytesIO(nft.image_data))
+                    images.append(img)
+                except Exception as e:
+                    logger.error(f"Error converting image data: {e}")
+        return images
+
+class ChatInterface:
+    """Enhanced chat interface with persistent state and improved NFT handling."""
+    
+    def __init__(self, openai_key: str, etherscan_key: str, opensea_key: str):
+        self.openai_key = openai_key
+        self.etherscan_key = etherscan_key
+        self.opensea_key = opensea_key
+        self.state_manager = StateManager()
+        self.nft_processor = None
+        self.current_session_id = str(int(time.time()))
+        
+        # Track conversation context
+        self.active_wallets: Set[str] = set()
+        self.active_nfts: Set[str] = set()  # contract_tokenid
+        self.mentioned_collections: Set[str] = set()
+        
+        # Initialize OpenAI client
+        openai.api_key = openai_key
+        self.client = openai.OpenAI(api_key=openai_key)
 
     def _format_context_message(self) -> str:
-        """Format wallet data for GPT system prompt."""
+        """Format enhanced context for GPT prompt."""
         lines = []
-        if not self.context:
-            return ""
-        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
-            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)")
+        
+        # Active wallets section
+        if self.active_wallets:
+            lines.append("\nActive Wallets:")
+            for addr in self.active_wallets:
+                wallet_data = self.state_manager.get_wallet_data(addr)
+                if wallet_data:
+                    lines.append(f"\nWallet {addr[:8]}...{addr[-6:]}:")
+                    lines.append(f"- ETH Balance: {wallet_data['eth_balance']:.4f}")
+                    lines.append(f"- Token Count: {len(wallet_data['tokens'])}")
+                    
+                    # Token details
+                    if wallet_data['tokens']:
+                        lines.append("- Notable Tokens:")
+                        for token in wallet_data['tokens'][:3]:  # Top 3 tokens
+                            lines.append(f"  * {token['name']} ({token['symbol']}): {token['balance']}")
+                    
+                    # NFT Collections
+                    nft_data = wallet_data["nft_collections"]
+                    if "collections" in nft_data and nft_data["collections"]:
+                        lines.append("- NFT Collections:")
+                        for coll in nft_data["collections"]:
+                            count = len(coll["items"])
+                            lines.append(f"  * {coll['collection_name']}: {count} NFT(s)")
+                            if "stats" in coll:
+                                lines.append(f"    First acquired: {coll['stats']['first_acquired']}")
+
+        # Active NFTs section
+        if self.active_nfts:
+            lines.append("\nRecently Discussed NFTs:")
+            for nft_key in list(self.active_nfts)[-5:]:  # Last 5 NFTs
+                contract, token_id = nft_key.split("_")
+                nft_meta = self.state_manager.get_nft_metadata(contract, token_id)
+                if nft_meta:
+                    lines.append(f"- {nft_meta.name} from {nft_meta.collection_name}")
+
         return "\n".join(lines)
 
+    def _prepare_chat_messages(
+        self,
+        user_message: str,
+        history: List[Tuple[str, str]]
+    ) -> List[Dict[str, str]]:
+        """Prepare messages for OpenAI chat completion."""
+        messages = [
+            {
+                "role": "system",
+                "content": Config.SYSTEM_PROMPT
+            },
+            {
+                "role": "system",
+                "content": self._format_context_message()
+            }
+        ]
+
+        # Add relevant history with context
+        limit = Config.HISTORY_LIMIT
+        for usr, ast in history[-limit:]:
+            messages.append({"role": "user", "content": usr})
+            messages.append({"role": "assistant", "content": ast})
+
+        # Add current message
+        messages.append({"role": "user", "content": user_message})
+        
+        return messages
+
     async def process_message(
         self,
         message: str,
         history: Optional[ChatHistory] = None
     ) -> Tuple[ChatHistory, Dict[str, Any], List[Image.Image]]:
         """
-        1) Detect Ethereum address
-        2) Etherscan for wallet data
-        3) For each NFT, fetch from OpenSea, convert to PIL
-        4) Return images + chat response
+        Enhanced message processing with better context management:
+        1) Track wallet addresses and NFTs mentioned
+        2) Maintain conversation context
+        3) Use cached data when appropriate
+        4) Return relevant NFT images
         """
-        if history is None:
-            history = []
-
+        history = self.state_manager.get_chat_history(self.current_session_id) if history is None else history
+        
         if not message.strip():
-            return history, self.context, []
+            return history, {}, []
 
-        match = re.search(Config.ETHEREUM_ADDRESS_REGEX, message)
-        nft_images: List[Image.Image] = []
-
-        if match:
-            eth_address = match.group(0)
-            # partial
-            partial_text = f"Analyzing {eth_address}..."
-            history.append((message, partial_text))
+        try:
+            # Initialize NFT processor if needed
+            if self.nft_processor is None:
+                self.nft_processor = NFTProcessor(self.opensea_key, self.state_manager)
 
+            # Look for Ethereum addresses
+            eth_addresses = re.findall(Config.ETHEREUM_ADDRESS_REGEX, message)
+            nft_images: List[Image.Image] = []
+            
+            # Process new addresses
+            for eth_address in eth_addresses:
+                if eth_address not in self.active_wallets:
+                    # Add to tracking
+                    self.active_wallets.add(eth_address)
+                    
+                    # Fetch wallet data
+                    try:
+                        partial_msg = f"📊 Analyzing {eth_address[:8]}...{eth_address[-6:]}"
+                        history.append((message, partial_msg))
+                        
+                        async with WalletAnalyzer(self.etherscan_key, self.state_manager) as analyzer:
+                            wallet_data = await analyzer.get_portfolio_data(eth_address)
+                        
+                        # Process NFTs for new wallet
+                        if "collections" in wallet_data["nft_collections"]:
+                            collections = wallet_data["nft_collections"]["collections"]
+                            if collections:
+                                # Process NFTs in parallel
+                                processed_nfts = await self.nft_processor.batch_process_nfts(
+                                    collections,
+                                    max_per_collection=Config.MAX_NFTS_PER_COLLECTION
+                                )
+                                
+                                # Update tracking
+                                for nft in processed_nfts:
+                                    self.active_nfts.add(f"{nft.contract}_{nft.token_id}")
+                                    self.mentioned_collections.add(nft.collection_name)
+                                
+                                # Get images for display
+                                nft_images.extend(
+                                    self.nft_processor.get_cached_images(processed_nfts)
+                                )
+                        
+                        # Generate wallet summary
+                        summary = [
+                            f"📊 Summary for {eth_address[:8]}...{eth_address[-6:]}",
+                            f"ETH: {wallet_data['eth_balance']:.4f}",
+                            f"Tokens: {len(wallet_data['tokens'])}",
+                            f"NFTs: {sum(len(c['items']) for c in wallet_data['nft_collections'].get('collections', []))}",
+                            "\nNotable Collections:"
+                        ]
+                        
+                        for coll in wallet_data["nft_collections"].get("collections", [])[:3]:
+                            summary.append(f"- {coll['collection_name']}: {len(coll['items'])} items")
+                        
+                        history.append((message, "\n".join(summary)))
+                        
+                    except Exception as e:
+                        err_msg = f"Error analyzing {eth_address}: {str(e)}"
+                        logger.error(err_msg)
+                        history.append((message, err_msg))
+
+            # Prepare and send chat completion
             try:
-                # 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
-                lines = [
-                    f"📊 Summary for {eth_address[:8]}...{eth_address[-6:]}",
-                    f"ETH: {wallet_data['eth_balance']:.4f}",
-                    f"Tokens: {len(wallet_data['tokens'])}"
-                ]
-                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)))
-
-                # 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"OpenSea metadata error: {meta['error']}")
-                                continue
-                            image_url = meta["image_url"]
-                            if not image_url:
-                                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))
-
+                messages = self._prepare_chat_messages(message, history)
+                
+                response = await openai.chat.completions.acreate(
+                    model=Config.OPENAI_MODEL,
+                    messages=messages,
+                    temperature=Config.TEMPERATURE,
+                    max_tokens=Config.MAX_TOKENS
+                )
+                
+                ai_message = response.choices[0].message.content
+                history.append((message, ai_message))
+                
             except Exception as e:
-                err = f"Error analyzing {eth_address}: {str(e)}"
-                logger.error(err)
-                history.append((message, err))
-
-        # Now do an OpenAI chat
-        try:
-            context_str = self._format_context_message()
-            # Convert chat to OpenAI format
-            limit = Config.HISTORY_LIMIT
-            short_hist = history[-limit:]
-            openai_msgs = []
-            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)
-            resp = client.chat.completions.create(
-                model=Config.OPENAI_MODEL,
-                messages=[
-                    {"role": "system", "content": Config.SYSTEM_PROMPT},
-                    {"role": "system", "content": context_str},
-                    *openai_msgs,
-                    {"role": "user", "content": message}
-                ],
-                temperature=Config.TEMPERATURE,
-                max_tokens=Config.MAX_TOKENS
-            )
-            final_msg = resp.choices[0].message.content
-            history.append((message, final_msg))
+                logger.error(f"OpenAI error: {e}")
+                error_msg = f"Chat error: {str(e)}"
+                history.append((message, error_msg))
 
-            return history, self.context, nft_images
+            # Update chat history in state manager
+            self.state_manager.store_chat_history(self.current_session_id, history)
+            
+            # Prepare context for response
+            context = {
+                "active_wallets": list(self.active_wallets),
+                "active_nfts": list(self.active_nfts),
+                "mentioned_collections": list(self.mentioned_collections)
+            }
+            
+            return history, context, nft_images
 
         except Exception as e:
-            logger.error(f"OpenAI error: {e}")
-            err_chat = f"OpenAI error: {str(e)}"
-            history.append((message, err_chat))
-            return history, self.context, []
+            logger.error(f"Error in process_message: {e}")
+            error_msg = f"Error processing message: {str(e)}"
+            history.append((message, error_msg))
+            return history, {}, []
 
     def clear_context(self) -> Tuple[Dict[str, Any], List[Tuple[str, str]]]:
-        """Clear the wallet context and chat."""
-        self.context = {}
+        """Clear all conversation context and tracking."""
+        self.active_wallets.clear()
+        self.active_nfts.clear()
+        self.mentioned_collections.clear()
+        self.state_manager.clear_session(self.current_session_id)
         return {}, []
-#### Part 5: Gradio UI & Launch
 
 class GradioInterface:
-    """Manages the Gradio UI for Hugging Face Space, with top-right NFT gallery as PIL images."""
+    """Enhanced Gradio UI with improved state management and user experience."""
     
     def __init__(self):
         self.chat_interface: Optional[ChatInterface] = None
         self.demo = self._create_interface()
 
     def _create_interface(self) -> gr.Blocks:
-        with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        """Create enhanced Gradio interface with better state visualization."""
+        
+        with gr.Blocks(
+            theme=gr.themes.Soft(),
+            css="""
+                .container { max-width: 1200px; margin: auto; }
+                .status-box { padding: 10px; border-radius: 5px; }
+                .success { background-color: #d4edda; color: #155724; }
+                .error { background-color: #f8d7da; color: #721c24; }
+            """
+        ) as demo:
             gr.Markdown("""
-            # 🐕 LOSS DOG: Blockchain Wallet Analyzer (NFT Images, top-right)
+            # 🐕 LOSS DOG: Enhanced Blockchain Wallet Analyzer
             
-            - 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
+            ### Features:
+            - **Smart Context**: Maintains conversation memory about wallets and NFTs
+            - **Caching**: Faster responses for previously seen data
+            - **Enhanced NFT Display**: View NFT images with metadata
+            - **Intelligent Analysis**: Tracks patterns across wallets
             """)
 
             with gr.Row():
-                openai_key = gr.Textbox(
-                    label="OpenAI API Key",
-                    type="password",
-                    placeholder="Enter your OpenAI API key..."
-                )
-                etherscan_key = gr.Textbox(
-                    label="Etherscan API Key",
-                    type="password",
-                    placeholder="Enter your Etherscan API key..."
-                )
-                opensea_key = gr.Textbox(
-                    label="OpenSea API Key",
-                    type="password",
-                    placeholder="Enter your OpenSea API key..."
-                )
-
-            validation_status = gr.Textbox(label="Validation Status", interactive=False)
-            validate_btn = gr.Button("Validate API Keys", variant="primary")
+                with gr.Column(scale=1):
+                    # API Configuration
+                    with gr.Group(label="API Configuration"):
+                        openai_key = gr.Textbox(
+                            label="OpenAI API Key",
+                            type="password",
+                            placeholder="Enter your OpenAI API key...",
+                            container=False
+                        )
+                        etherscan_key = gr.Textbox(
+                            label="Etherscan API Key",
+                            type="password",
+                            placeholder="Enter your Etherscan API key...",
+                            container=False
+                        )
+                        opensea_key = gr.Textbox(
+                            label="OpenSea API Key",
+                            type="password",
+                            placeholder="Enter your OpenSea API key...",
+                            container=False
+                        )
+                        
+                    # Validation Status
+                    validation_status = gr.Textbox(
+                        label="Status",
+                        interactive=False,
+                        container=True
+                    )
+                    validate_btn = gr.Button(
+                        "🔑 Validate API Keys",
+                        variant="primary",
+                        size="sm"
+                    )
 
             with gr.Row():
-                # Main Chat Column
+                # Main Chat Area
                 with gr.Column(scale=2):
-                    chatbot = gr.Chatbot(label="Chat History", height=420, value=[])
+                    chatbot = gr.Chatbot(
+                        label="Conversation",
+                        height=500,
+                        value=[],
+                        elem_classes="chatbot"
+                    )
+                    
                     with gr.Row():
                         msg_input = gr.Textbox(
                             label="Message",
-                            placeholder="Enter an ETH address or question..."
+                            placeholder="Enter an ETH address or ask about NFTs...",
+                            scale=4
+                        )
+                        send_btn = gr.Button(
+                            "🚀 Send",
+                            variant="primary",
+                            scale=1
+                        )
+                        clear_btn = gr.Button(
+                            "🔄 Clear",
+                            variant="secondary",
+                            scale=1
                         )
-                        send_btn = gr.Button("Send", variant="primary")
-                # Top-Right: NFT Gallery (PIL images)
+
+                # Right Column: NFTs and Context
                 with gr.Column(scale=1):
+                    # NFT Gallery
                     nft_gallery = gr.Gallery(
-                        label="NFT Images (Top-Right)",
-                        columns=2
-                    )
-                    wallet_context = gr.JSON(
-                        label="Active Wallet Context",
-                        value={}
+                        label="📸 NFT Gallery",
+                        columns=2,
+                        height=300,
+                        object_fit="contain"
                     )
-                    clear_btn = gr.Button("Clear Context")
+                    
+                    # Context Display
+                    with gr.Accordion("🔍 Active Context", open=True):
+                        wallet_context = gr.JSON(
+                            label="Tracked Data",
+                            value={}
+                        )
+                        
+                        context_info = gr.Markdown("""
+                        - **Wallets**: Currently tracked addresses
+                        - **NFTs**: Recently discussed tokens
+                        - **Collections**: Active collections
+                        """)
+
+            # Footer
+            gr.Markdown("""
+            ---
+            ### 📝 Usage Tips:
+            1. Enter your API keys and validate
+            2. Type an Ethereum address to analyze
+            3. Ask about specific NFTs or collections
+            4. Images appear in the gallery automatically
+            """)
 
+            # Initialize state
             msg_input.interactive = False
             send_btn.interactive = False
 
-            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"✅ {msg}",
-                        gr.update(interactive=True),
-                        gr.update(interactive=True)
+            def validate_keys(
+                openai_k: str,
+                etherscan_k: str,
+                opensea_k: str
+            ) -> Tuple[str, gr.update, gr.update]:
+                """Validate API keys with enhanced feedback."""
+                try:
+                    is_valid, msg = ChatInterface._validate_api_keys(
+                        openai_k, etherscan_k, opensea_k
                     )
-                else:
+                    if is_valid:
+                        self.chat_interface = ChatInterface(
+                            openai_k, etherscan_k, opensea_k
+                        )
+                        return (
+                            "✅ All API keys validated successfully!",
+                            gr.update(interactive=True),
+                            gr.update(interactive=True)
+                        )
+                    else:
+                        return (
+                            f"❌ Validation failed: {msg}",
+                            gr.update(interactive=False),
+                            gr.update(interactive=False)
+                        )
+                except Exception as e:
                     return (
-                        f"❌ {msg}",
+                        f"❌ Error during validation: {str(e)}",
                         gr.update(interactive=False),
                         gr.update(interactive=False)
                     )
 
-            validate_btn.click(
-                fn=validate_keys,
-                inputs=[openai_key, etherscan_key, opensea_key],
-                outputs=[validation_status, msg_input, send_btn]
-            )
-
-            def clear_all():
-                """Clear wallet context and chat."""
+            def clear_all() -> Tuple[Dict[str, Any], List[Tuple[str, str]], List[Image.Image]]:
+                """Clear all state with visual feedback."""
                 if self.chat_interface:
-                    return self.chat_interface.clear_context()
-                return {}, []
-
-            clear_btn.click(
-                fn=clear_all,
-                inputs=[],
-                outputs=[wallet_context, chatbot]
-            )
+                    context, history = self.chat_interface.clear_context()
+                    return context, history, []
+                return {}, [], []
 
             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[Image.Image]]:
-                """Process user input, return updated chat, context, list of PIL images."""
+                """Process messages with enhanced error handling."""
                 if not self.chat_interface:
                     return [], {}, []
                 
                 try:
-                    new_history, new_context, images = await self.chat_interface.process_message(message, chat_hist)
+                    # Show typing indicator
+                    yield [*chat_hist, (message, "🤔 Processing...")], context, []
+                    
+                    # Process message
+                    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}")
+                    error_msg = f"🔴 Error: {str(e)}"
                     if chat_hist is None:
                         chat_hist = []
-                    chat_hist.append((message, f"Error: {str(e)}"))
-                    return chat_hist, context, []
+                    return [*chat_hist, (message, error_msg)], context, []
 
-            # Chat flow
+            # Wire up events
+            validate_btn.click(
+                fn=validate_keys,
+                inputs=[openai_key, etherscan_key, opensea_key],
+                outputs=[validation_status, msg_input, send_btn]
+            )
+
+            clear_btn.click(
+                fn=clear_all,
+                inputs=[],
+                outputs=[wallet_context, chatbot, nft_gallery]
+            )
+
+            # Message handling
             msg_input.submit(
                 fn=handle_message,
                 inputs=[msg_input, chatbot, wallet_context],
@@ -622,15 +1179,25 @@ class GradioInterface:
 
         return demo
 
-    def launch(self):
-        self.demo.queue()
-        self.demo.launch()
+    def launch(self, **kwargs):
+        """Launch the interface with configurable options."""
+        self.demo.queue(concurrency_count=3)
+        self.demo.launch(**kwargs)
 
 def main():
-    """Main entry point for Hugging Face Space."""
-    logger.info("Launching LOSS DOG with PIL-based NFT images (top-right).")
-    interface = GradioInterface()
-    interface.launch()
+    """Enhanced main entry point with error handling."""
+    try:
+        logger.info("Launching Enhanced LOSS DOG Blockchain Analyzer")
+        interface = GradioInterface()
+        interface.launch(
+            share=False,
+            server_name="0.0.0.0",
+            server_port=7860,
+            show_error=True
+        )
+    except Exception as e:
+        logger.error(f"Failed to launch interface: {e}")
+        raise
 
 if __name__ == "__main__":
     main()