Update app.py

app.py

@@ -2,16 +2,14 @@
 Blockchain Wallet Analyzer - A tool for analyzing Ethereum wallet contents and NFT holdings.
 
 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.
+with a Gradio web interface. It includes improved NFT tracking, natural AI communication,
+and enhanced error handling.
 
 Author: Claude
 Date: January 2025
 """
 
 from __future__ import annotations
-
 import os
 import re
 import json
@@ -19,10 +17,10 @@ import time
 import logging
 import asyncio
 import base64
-from typing import List, Dict, Tuple, Any, Optional, TypeVar
+from typing import List, Dict, Tuple, Any, Optional, TypeVar, Set
 from datetime import datetime
 from decimal import Decimal
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 from io import BytesIO
 
@@ -36,6 +34,7 @@ from tenacity import retry, stop_after_attempt, wait_exponential
 T = TypeVar('T')
 WalletData = Dict[str, Any]
 ChatHistory = List[Tuple[str, str]]
+NFTIdentifier = Tuple[str, str]  # (contract, token_id)
 
 # Configure logging
 logging.basicConfig(
@@ -48,45 +47,99 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-class ConfigError(Exception):
-    """Raised when there's an error in configuration."""
+class BlockchainError(Exception):
+    """Base exception for blockchain-related errors."""
+    pass
+
+class ConfigError(BlockchainError):
+    """Configuration error."""
+    pass
+
+class APIError(BlockchainError):
+    """API call error."""
     pass
 
-class APIError(Exception):
-    """Raised when there's an error in API calls."""
+class ValidationError(BlockchainError):
+    """Input validation error."""
     pass
 
-class ValidationError(Exception):
-    """Raised when there's an error in input validation."""
+class NFTError(BlockchainError):
+    """NFT processing error."""
     pass
 
+@dataclass
+class NFTMetadata:
+    """Structure for NFT metadata."""
+    contract: str
+    token_id: str
+    name: str
+    collection_name: str
+    image_url: Optional[str] = None
+    description: Optional[str] = None
+    traits: List[Dict[str, Any]] = field(default_factory=list)
+    last_updated: datetime = field(default_factory=datetime.now)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary format."""
+        return {
+            "contract": self.contract,
+            "token_id": self.token_id,
+            "name": self.name,
+            "collection_name": self.collection_name,
+            "image_url": self.image_url,
+            "description": self.description,
+            "traits": self.traits,
+            "last_updated": self.last_updated.isoformat()
+        }
+
 @dataclass
 class Config:
-    """Application configuration settings."""
+    """Enhanced application configuration."""
+    
+    # Improved system prompt for more natural responses
     SYSTEM_PROMPT: str = """
     You are LOSS DOG 🐕 (Learning & Observing Smart Systems Digital Output Generator), 
-    an adorable blockchain-sniffing puppy!
-    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
-    - Use this data to provide specific answers about holdings
-    - Reference exact numbers and collections when discussing NFTs
-    - Compare wallets if multiple are available
+    a friendly blockchain analysis assistant!
+
+    Your conversation style:
+    - Maintain natural dialogue flow
+    - Avoid repetitive responses
+    - Use context from previous messages
+    - Track which NFTs have been discussed
+    - Provide insights when relevant
+
+    When analyzing wallets:
+    - Reference specific NFTs by collection and ID
+    - Highlight interesting patterns
+    - Compare wallets if multiple available
+    - Monitor NFT rarity and traits
     """
+    
+    # API Configuration
     ETHERSCAN_BASE_URL: str = "https://api.etherscan.io/api"
-    ETHEREUM_ADDRESS_REGEX: str = r"0x[a-fA-F0-9]{40}"
+    OPENSEA_BASE_URL: str = "https://api.opensea.io/api/v2"
+    ETHEREUM_ADDRESS_REGEX: str = r"^0x[a-fA-F0-9]{40}$"
+
+    # Rate Limiting & Retries
     RATE_LIMIT_DELAY: float = 0.2  # 5 requests per second
     MAX_RETRIES: int = 3
-    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)
+    RETRY_BASE_DELAY: float = 1.0
+    RETRY_MAX_DELAY: float = 10.0
 
+    # API Timeouts
+    REQUEST_TIMEOUT: float = 30.0
+    NFT_FETCH_TIMEOUT: int = 10
+    
+    # NFT Processing
+    MAX_NFTS_PER_COLLECTION: int = 50
+    NFT_BATCH_SIZE: int = 5
+    MAX_IMAGE_SIZE: Tuple[int, int] = (800, 800)
+    
+    # OpenAI Settings
+    OPENAI_MODEL: str = "gpt-4o-mini"
+    MAX_TOKENS: int = 8000
+    TEMPERATURE: float = 0.7
+    HISTORY_LIMIT: int = 10
 class WalletAnalyzer:
     """Analyzes Ethereum wallet contents using Etherscan API."""
     
@@ -95,9 +148,12 @@ class WalletAnalyzer:
         self.base_url = Config.ETHERSCAN_BASE_URL
         self.session: Optional[aiohttp.ClientSession] = None
         self.last_request_time = 0
+        self._rate_limit_semaphore = asyncio.Semaphore(1)
     
     async def __aenter__(self) -> WalletAnalyzer:
-        self.session = aiohttp.ClientSession()
+        self.session = aiohttp.ClientSession(
+            timeout=aiohttp.ClientTimeout(total=Config.REQUEST_TIMEOUT)
+        )
         return self
     
     async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
@@ -105,83 +161,108 @@ class WalletAnalyzer:
             await self.session.close()
             self.session = None
 
+    async def _rate_limit(self) -> None:
+        """Thread-safe rate limiting."""
+        async with self._rate_limit_semaphore:
+            now = time.time()
+            time_since_last = now - self.last_request_time
+            if time_since_last < Config.RATE_LIMIT_DELAY:
+                await asyncio.sleep(Config.RATE_LIMIT_DELAY - time_since_last)
+            self.last_request_time = time.time()
+
     @retry(
         stop=stop_after_attempt(Config.MAX_RETRIES),
-        wait=wait_exponential(multiplier=1, min=4, max=10)
+        wait=wait_exponential(
+            multiplier=Config.RETRY_BASE_DELAY,
+            min=Config.RETRY_BASE_DELAY,
+            max=Config.RETRY_MAX_DELAY
+        )
     )
     async def _fetch_data(self, params: Dict[str, str]) -> Dict[str, Any]:
-        """Fetch data from Etherscan with retry logic."""
+        """Fetch data from Etherscan with improved 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:
+                    raise APIError("Rate limit exceeded")
                 if response.status != 200:
                     raise APIError(f"Etherscan request failed: {response.status}")
-                data = await response.json()
+                    
+                try:
+                    data = await response.json()
+                except json.JSONDecodeError as e:
+                    raise APIError(f"Invalid JSON response: {e}")
+                    
                 if data.get("status") == "0":
-                    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}")
+                    message = data.get("message", "Unknown error")
+                    result = data.get("result", "")
+                    if "Max rate limit reached" in message:
+                        raise APIError(f"Rate limit exceeded: {result}")
+                    raise APIError(f"Etherscan API error: {message}")
+                    
                 return data
+                
+        except asyncio.TimeoutError:
+            raise APIError("Request timed out")
         except aiohttp.ClientError as e:
             raise APIError(f"Network error: {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."""
-        now = time.time()
-        diff = now - self.last_request_time
-        if diff < Config.RATE_LIMIT_DELAY:
-            await asyncio.sleep(Config.RATE_LIMIT_DELAY - diff)
-        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))
 
-    async def get_portfolio_data(self, address: str) -> WalletData:
+    async def get_wallet_data(self, address: str) -> WalletData:
         """
-        Return a dictionary with:
-         - address
-         - last_updated
-         - eth_balance
-         - tokens (list of ERC-20)
-         - nft_collections (with contract & token_id)
+        Get comprehensive wallet data including:
+        - ETH balance
+        - ERC-20 tokens
+        - NFT collections
         """
         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)
-
-        return {
-            "address": address,
-            "last_updated": datetime.now().isoformat(),
-            "eth_balance": float(eth_balance),
-            "tokens": tokens,
-            "nft_collections": nft_colls
-        }
+        logger.info(f"Fetching data for wallet: {address}")
+        try:
+            eth_balance = await self._get_eth_balance(address)
+            tokens = await self._get_token_holdings(address)
+            nfts = await self._get_nft_holdings(address)
+
+            return {
+                "address": address,
+                "last_updated": datetime.now().isoformat(),
+                "eth_balance": float(eth_balance),
+                "tokens": tokens,
+                "nft_collections": nfts
+            }
+        except Exception as e:
+            logger.error(f"Error fetching wallet data: {e}")
+            raise
 
     async def _get_eth_balance(self, address: str) -> Decimal:
+        """Get ETH balance for address."""
         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 Wei to ETH
+        except (KeyError, ValueError) as e:
+            raise APIError(f"Error parsing ETH balance: {e}")
 
     async def _get_token_holdings(self, address: str) -> List[Dict[str, Any]]:
-        """Fetch ERC-20 tokens for address."""
+        """Get ERC-20 token holdings with improved accuracy."""
         params = {
             "module": "account",
             "action": "tokentx",
@@ -192,33 +273,39 @@ class WalletAnalyzer:
         
         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)
-                }
-            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
-
+            try:
+                contract = tx["contractAddress"].lower()
+                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[contract]["decimals"])
+                if tx["to"].lower() == address.lower():
+                    token_map[contract]["balance"] += amount
+                elif tx["from"].lower() == address.lower():
+                    token_map[contract]["balance"] -= amount
+            except (KeyError, ValueError, decimal.InvalidOperation) as e:
+                logger.warning(f"Error processing token transaction: {e}")
+                continue
+
+        # Filter tokens with non-zero balances
         return [
             {
-                "name": v["name"],
-                "symbol": v["symbol"],
-                "balance": float(v["balance"])
+                "name": token["name"],
+                "symbol": token["symbol"],
+                "balance": float(token["balance"])
             }
-            for v in token_map.values() if v["balance"] > 0
+            for token in token_map.values()
+            if token["balance"] > 0
         ]
-
-    async def _get_nft_holdings(self, address: str) -> Dict[str, Any]:
+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.
+        Get NFT holdings grouped by collection.
+        Returns: {"collections": [{"collection_name": str, "items": [{"contract", "token_id"}, ...]}]}
         """
         params = {
             "module": "account",
@@ -226,412 +313,359 @@ class WalletAnalyzer:
             "address": address,
             "sort": "desc"
         }
-        data = await self._fetch_data(params)
+        
+        try:
+            data = await self._fetch_data(params)
+            if data.get("status") != "1" or "result" not in data:
+                return {"collections": []}
 
-        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)}"}
+            # Track NFT ownership by collection
+            collections: Dict[str, List[Dict[str, str]]] = {}
+            ownership: Dict[str, Dict[str, Any]] = {}  # contract_tokenId -> ownership info
+            
+            for tx in data["result"]:
+                contract = tx["contractAddress"].lower()
+                token_id = tx["tokenID"]
+                key = f"{contract}_{token_id}"
+                
+                # Get collection name, fallback with contract if needed
+                collection_name = tx.get("tokenName", "Unknown Collection")
+                if collection_name == "Unknown Collection" and contract in collections:
+                    collection_name = next(iter(collections[contract]))
+                
+                # Track ownership changes
+                if tx["to"].lower() == address.lower():
+                    ownership[key] = {
+                        "contract": contract,
+                        "token_id": token_id,
+                        "collection_name": collection_name
+                    }
+                elif tx["from"].lower() == address.lower():
+                    ownership.pop(key, None)
+
+            # Group owned NFTs by collection
+            for nft in ownership.values():
+                coll_name = nft["collection_name"]
+                if coll_name not in collections:
+                    collections[coll_name] = []
+                collections[coll_name].append({
+                    "contract": nft["contract"],
+                    "token_id": nft["token_id"]
+                })
+
+            return {
+                "collections": [
+                    {
+                        "collection_name": name,
+                        "items": items[:Config.MAX_NFTS_PER_COLLECTION]
+                    }
+                    for name, items in collections.items()
+                ]
+            }
+            
+        except Exception as e:
+            logger.error(f"Error fetching NFT holdings: {e}")
+            raise APIError(f"Failed to fetch NFT holdings: {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}
 
-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
+class NFTProcessor:
+    """Handles NFT metadata and image processing."""
+    
+    def __init__(self, opensea_key: str):
+        self.opensea_key = opensea_key
+        self._session: Optional[aiohttp.ClientSession] = None
+        self.last_request_time = 0
+        self._rate_limit_semaphore = asyncio.Semaphore(1)
+        self._metadata_cache: Dict[str, NFTMetadata] = {}
+        self._image_cache: Dict[str, Image.Image] = {}
+        
+    async def __aenter__(self) -> 'NFTProcessor':
+        self._session = aiohttp.ClientSession(
+            timeout=aiohttp.ClientTimeout(total=Config.NFT_FETCH_TIMEOUT)
+        )
+        return self
+        
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self._session:
+            await self._session.close()
+            self._session = None
+
+    async def _rate_limit(self) -> None:
+        """Thread-safe rate limiting."""
+        async with self._rate_limit_semaphore:
+            now = time.time()
+            time_since_last = now - self.last_request_time
+            if time_since_last < Config.RATE_LIMIT_DELAY:
+                await asyncio.sleep(Config.RATE_LIMIT_DELAY - time_since_last)
+            self.last_request_time = time.time()
+
+    def _get_cache_key(self, contract: str, token_id: str) -> str:
+        """Generate cache key for NFT."""
+        return f"{contract.lower()}_{token_id}"
 
+    @retry(
+        stop=stop_after_attempt(Config.MAX_RETRIES),
+        wait=wait_exponential(
+            multiplier=Config.RETRY_BASE_DELAY,
+            min=Config.RETRY_BASE_DELAY,
+            max=Config.RETRY_MAX_DELAY
+        )
+    )
+    async def fetch_nft_metadata(
+        self,
+        contract: str,
+        token_id: str,
+        force_refresh: bool = False
+    ) -> NFTMetadata:
+        """Fetch NFT metadata from OpenSea with caching."""
+        if not self._session:
+            raise NFTError("No active session")
+            
+        cache_key = self._get_cache_key(contract, token_id)
+        
+        # Check cache unless force refresh
+        if not force_refresh and cache_key in self._metadata_cache:
+            return self._metadata_cache[cache_key]
+            
+        await self._rate_limit()
+        
+        url = f"{Config.OPENSEA_BASE_URL}/chain/ethereum/contract/{contract}/nfts/{token_id}"
+        headers = {"X-API-KEY": self.opensea_key}
+        
+        try:
+            async with self._session.get(url, headers=headers) as resp:
+                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_data = data.get("nft", {})
+                
+                metadata = NFTMetadata(
+                    contract=contract,
+                    token_id=token_id,
+                    name=nft_data.get("name", f"NFT #{token_id}"),
+                    collection_name=nft_data.get("collection", {}).get("name", "Unknown"),
+                    image_url=nft_data.get("image_url"),
+                    description=nft_data.get("description"),
+                    traits=nft_data.get("traits", [])
+                )
+                
+                # Cache the result
+                self._metadata_cache[cache_key] = metadata
+                return metadata
+                
+        except asyncio.TimeoutError:
+            raise NFTError(f"Timeout fetching NFT {contract} #{token_id}")
+        except Exception as e:
+            raise NFTError(f"Error fetching NFT metadata: {str(e)}")
+
+    async def fetch_image(self, image_url: str) -> Optional[Image.Image]:
+        """Download and process NFT image with caching."""
+        if not image_url or not self._session:
+            return None
+            
+        # Check image cache
+        if image_url in self._image_cache:
+            return self._image_cache[image_url]
+            
+        try:
+            async with self._session.get(
+                image_url,
+                timeout=aiohttp.ClientTimeout(total=Config.NFT_FETCH_TIMEOUT)
+            ) as resp:
+                if resp.status != 200:
+                    return None
+                data = await resp.read()
+                
+            img = Image.open(BytesIO(data))
+            
+            # Convert to RGB/RGBA if needed
+            if img.mode not in ('RGB', 'RGBA'):
+                img = img.convert('RGBA')
+                
+            # Resize if needed
+            if (img.size[0] > Config.MAX_IMAGE_SIZE[0] or 
+                img.size[1] > Config.MAX_IMAGE_SIZE[1]):
+                img.thumbnail(Config.MAX_IMAGE_SIZE)
+            
+            # Cache the processed image
+            self._image_cache[image_url] = img
+            return img
+            
+        except Exception as e:
+            logger.error(f"Error processing image: {e}")
+            return None
 class ChatInterface:
-    """Handles chat logic with Etherscan (wallet data), OpenSea (NFT images), and OpenAI (chat)."""
+    """Manages chat interactions and context."""
     
     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.context: Dict[str, Any] = {}
+        self.context = ChatContext()
+        self.nft_processor = None
         openai.api_key = openai_key
 
-    @staticmethod
-    def _validate_api_keys(openai_key: str, etherscan_key: str, opensea_key: str) -> Tuple[bool, str]:
-        """
-        Validate all keys. We'll do a minimal check for OpenSea (non-empty).
-        """
-        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
-            )
-
-            # 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 _ensure_nft_processor(self) -> NFTProcessor:
+        """Initialize NFT processor if needed."""
+        if self.nft_processor is None:
+            self.nft_processor = NFTProcessor(self.opensea_key)
+            await self.nft_processor.__aenter__()
+        return self.nft_processor
+
+    def _format_context(self) -> str:
+        """Format wallet data for AI context."""
+        if not self.context.wallet_data:
+            return "No wallet data available."
+
+        lines = ["Current Wallet Analysis:"]
+        for addr, data in self.context.wallet_data.items():
+            lines.append(f"\nWallet {addr[:8]}...{addr[-6:]}:")
+            lines.append(f"ETH Balance: {data['eth_balance']:.4f} ETH")
+
+            if data['tokens']:
+                lines.append("\nToken Holdings:")
+                for token in sorted(
+                    data['tokens'],
+                    key=lambda x: x['balance'],
+                    reverse=True
+                )[:10]:  # Show top 10 tokens
+                    lines.append(f"- {token['symbol']}: {token['balance']:.4f}")
+
+            nft_data = data.get("nft_collections", {})
+            if "collections" in nft_data and nft_data["collections"]:
+                lines.append("\nNFT Collections:")
+                for coll in nft_data["collections"]:
+                    coll_name = coll["collection_name"]
+                    total_items = len(coll["items"])
+                    discussed = sum(1 for item in coll["items"]
+                                  if self.context.is_nft_discussed(
+                                      item["contract"],
+                                      item["token_id"]
+                                  ))
+                    lines.append(f"\n{coll_name}:")
+                    lines.append(f"- Total Items: {total_items}")
+                    lines.append(f"- Discussed: {discussed}/{total_items}")
 
-            # Check OpenSea
-            if not opensea_key.strip():
-                return False, "OpenSea API key is empty!"
-
-            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 wallet data for GPT system 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)")
         return "\n".join(lines)
 
     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
-        """
+        history: Optional[List[Tuple[str, str]]] = None
+    ) -> Tuple[List[Tuple[str, str]], Dict[str, Any], List[Image.Image]]:
+        """Process user message and return chat updates."""
         if history is None:
             history = []
-
+            
         if not message.strip():
-            return history, self.context, []
+            return history, self.context.wallet_data, []
 
-        match = re.search(Config.ETHEREUM_ADDRESS_REGEX, message)
         nft_images: List[Image.Image] = []
-
+        
+        # Check for Ethereum address
+        match = re.search(Config.ETHEREUM_ADDRESS_REGEX, message)
         if match:
             eth_address = match.group(0)
-            # partial
-            partial_text = f"Analyzing {eth_address}..."
-            history.append((message, partial_text))
+            history.append((
+                message,
+                f"Analyzing wallet {eth_address[:8]}...{eth_address[-6:]}..."
+            ))
 
             try:
-                # Etherscan
+                # Fetch wallet data
                 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))
+                    wallet_data = await analyzer.get_wallet_data(eth_address)
+                self.context.wallet_data[eth_address] = wallet_data
+
+                # Process NFTs
+                nft_processor = await self._ensure_nft_processor()
+                nft_info = wallet_data.get("nft_collections", {})
+                
+                if "collections" in nft_info and nft_info["collections"]:
+                    for collection in nft_info["collections"]:
+                        for i in range(0, len(collection["items"]), Config.NFT_BATCH_SIZE):
+                            batch = collection["items"][i:i + Config.NFT_BATCH_SIZE]
+                            metadata_tasks = [
+                                nft_processor.fetch_nft_metadata(
+                                    item["contract"],
+                                    item["token_id"]
+                                )
+                                for item in batch
+                            ]
+                            metadata_results = await asyncio.gather(
+                                *metadata_tasks,
+                                return_exceptions=True
+                            )
+
+                            for meta in metadata_results:
+                                if isinstance(meta, Exception):
+                                    continue
+                                if meta.image_url:
+                                    img = await nft_processor.fetch_image(meta.image_url)
+                                    if img:
+                                        nft_images.append(img)
 
             except Exception as e:
-                err = f"Error analyzing {eth_address}: {str(e)}"
-                logger.error(err)
-                history.append((message, err))
+                error_msg = f"Error analyzing wallet: {str(e)}"
+                logger.error(error_msg)
+                history.append((message, error_msg))
+                return history, self.context.wallet_data, []
 
-        # Now do an OpenAI chat
+        # Process with OpenAI
         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
+            messages = [
+                {"role": "system", "content": Config.SYSTEM_PROMPT},
+                {"role": "system", "content": self._format_context()}
+            ]
+
+            # Add recent history
+            recent_history = history[-Config.HISTORY_LIMIT:]
+            for usr, bot in recent_history:
+                messages.extend([
+                    {"role": "user", "content": usr},
+                    {"role": "assistant", "content": bot}
+                ])
+
+            messages.append({"role": "user", "content": message})
+
             client = openai.OpenAI(api_key=self.openai_key)
-            resp = client.chat.completions.create(
+            response = 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}
-                ],
+                messages=messages,
                 temperature=Config.TEMPERATURE,
                 max_tokens=Config.MAX_TOKENS
             )
-            final_msg = resp.choices[0].message.content
-            history.append((message, final_msg))
 
-            return history, self.context, nft_images
+            ai_message = response.choices[0].message.content
+            history.append((message, ai_message))
+            
+            return history, self.context.wallet_data, 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, []
+            error_msg = f"Chat error: {str(e)}"
+            logger.error(error_msg)
+            history.append((message, error_msg))
+            return history, self.context.wallet_data, []
 
     def clear_context(self) -> Tuple[Dict[str, Any], List[Tuple[str, str]]]:
-        """Clear the wallet context and chat."""
-        self.context = {}
+        """Clear chat context and history."""
+        self.context = ChatContext()
         return {}, []
-#### Part 5: Gradio UI & Launch
-
-class GradioInterface:
-    """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
-        self.demo = self._create_interface()
-
-    def _create_interface(self) -> gr.Blocks:
-        with gr.Blocks(theme=gr.themes.Soft()) as demo:
-            gr.Markdown("""
-            # 🐕 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():
-                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.Row():
-                # Main Chat Column
-                with gr.Column(scale=2):
-                    chatbot = gr.Chatbot(label="Chat History", height=420, value=[])
-                    with gr.Row():
-                        msg_input = gr.Textbox(
-                            label="Message",
-                            placeholder="Enter an ETH address or question..."
-                        )
-                        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
-                    )
-                    wallet_context = gr.JSON(
-                        label="Active Wallet Context",
-                        value={}
-                    )
-                    clear_btn = gr.Button("Clear Context")
-
-            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)
-                    )
-                else:
-                    return (
-                        f"❌ {msg}",
-                        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."""
-                if self.chat_interface:
-                    return self.chat_interface.clear_context()
-                return {}, []
-
-            clear_btn.click(
-                fn=clear_all,
-                inputs=[],
-                outputs=[wallet_context, chatbot]
-            )
-
-            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."""
-                if not self.chat_interface:
-                    return [], {}, []
-                
-                try:
-                    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, []
-
-            # Chat flow
-            msg_input.submit(
-                fn=handle_message,
-                inputs=[msg_input, chatbot, wallet_context],
-                outputs=[chatbot, wallet_context, nft_gallery]
-            ).then(
-                lambda: gr.update(value=""),
-                None,
-                [msg_input]
-            )
-
-            send_btn.click(
-                fn=handle_message,
-                inputs=[msg_input, chatbot, wallet_context],
-                outputs=[chatbot, wallet_context, nft_gallery]
-            ).then(
-                lambda: gr.update(value=""),
-                None,
-                [msg_input]
-            )
-
-        return demo
 
-    def launch(self):
-        self.demo.queue()
-        self.demo.launch()
+    async def cleanup(self):
+        """Cleanup resources."""
+        if self.nft_processor:
+            await self.nft_processor.__aexit__(None, None, None)
+            self.nft_processor = None
 
 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()
+    """Main entry point."""
+    demo = create_gradio_interface()
+    demo.queue()
+    demo.launch()
 
 if __name__ == "__main__":
-    main()
- 
+    main()