Update app.py

app.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """
-Healthcare Billing Chatbot - Hugging Face Spaces App
-Main application file for Hugging Face Spaces deployment
+Hybrid AI Assistant - General Purpose + Healthcare Billing Expert
+A ChatGPT-style assistant that can handle any conversation while specializing in healthcare billing codes
 """
 
 import os
@@ -14,48 +14,17 @@ from dataclasses import dataclass, field
 from enum import Enum
 import requests
 import gradio as gr
-from datetime import datetime, timedelta
-from pathlib import Path
+from datetime import datetime
+import random
 
-# Set up environment variables (for Hugging Face Spaces)
+# Set up environment
 os.environ['OPENROUTER_API_KEY'] = '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'
-os.environ['DATA_SOURCE'] = 'local_csv'
-os.environ['DATA_PATH'] = 'data/sample_codes.csv'
 
 # Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# ============= Data Classes and Enums =============
-
-class ExplanationLength(Enum):
-    SHORT = "short"
-    DETAILED = "detailed"
-    EXTRA = "extra"
-
-class CodeType(Enum):
-    CPT = "CPT"
-    HCPCS = "HCPCS"
-    ICD10 = "ICD-10"
-    DRG = "DRG"
-    UNKNOWN = "UNKNOWN"
-
-class DataSource(Enum):
-    LOCAL_CSV = "local_csv"
-    LOCAL_JSON = "local_json"
-    API = "api"
-
-@dataclass
-class ConversationContext:
-    current_code: Optional[str] = None
-    current_code_type: Optional[CodeType] = None
-    current_code_info: Optional['CodeInfo'] = None
-    conversation_history: List[Dict[str, str]] = field(default_factory=list)
-    last_explanation_length: Optional[ExplanationLength] = None
-    turn_count: int = 0
+# ============= Data Classes =============
 
 @dataclass
 class CodeInfo:
@@ -63,222 +32,139 @@ class CodeInfo:
     description: str
     code_type: str
     additional_info: Optional[str] = None
-    effective_date: Optional[str] = None
     category: Optional[str] = None
-    
-    def to_dict(self) -> Dict[str, Any]:
-        return {
-            'code': self.code,
-            'description': self.description,
-            'code_type': self.code_type,
-            'additional_info': self.additional_info,
-            'effective_date': self.effective_date,
-            'category': self.category
-        }
 
-# ============= HCPCS Loader Module =============
+@dataclass
+class ConversationContext:
+    messages: List[Dict[str, str]] = field(default_factory=list)
+    detected_codes: List[str] = field(default_factory=list)
+    last_topic: Optional[str] = None
+
+# ============= Healthcare Billing Database =============
 
-class HCPCSLoader:
-    """Loader for HCPCS/CPT/ICD-10/DRG codes"""
-    
+class BillingCodesDB:
     def __init__(self):
-        self.data_source = DataSource(os.getenv('DATA_SOURCE', 'local_csv'))
-        self.data_path = os.getenv('DATA_PATH', 'data/sample_codes.csv')
-        self.codes_cache: Dict[str, CodeInfo] = {}
-        self.cache_timestamp: Optional[datetime] = None
-    
-    def load_data(self) -> bool:
-        """Load data from configured source"""
-        try:
-            return self._load_sample_data()
-        except Exception as e:
-            logger.error(f"Failed to load data: {e}")
-            return self._load_sample_data()
-    
-    def _load_sample_data(self) -> bool:
-        """Load sample data as fallback"""
-        logger.info("Loading sample data")
-        
-        sample_codes = {
+        self.codes = {
             'A0429': CodeInfo(
                 code='A0429',
                 description='Ambulance service, basic life support, emergency transport (BLS-emergency)',
                 code_type='HCPCS',
-                additional_info='Ground ambulance emergency transport with BLS level care',
+                additional_info='Ground ambulance emergency transport with BLS level care. Used for emergency situations requiring immediate medical transport.',
                 category='Ambulance Services'
             ),
             'A0428': CodeInfo(
                 code='A0428',
-                description='Ambulance service, basic life support, non-emergency transport (BLS)',
+                description='Ambulance service, basic life support, non-emergency transport',
                 code_type='HCPCS',
-                additional_info='Ground ambulance non-emergency transport with BLS level care',
+                additional_info='Scheduled or non-urgent medical transport with basic life support.',
                 category='Ambulance Services'
             ),
             '99213': CodeInfo(
                 code='99213',
-                description='Office or other outpatient visit for evaluation and management, established patient, low complexity',
+                description='Office visit for established patient, low complexity',
                 code_type='CPT',
-                additional_info='Typically 20-29 minutes with patient',
+                additional_info='Typically 20-29 minutes. For straightforward medical issues.',
                 category='E&M Services'
             ),
             '99214': CodeInfo(
                 code='99214',
-                description='Office or other outpatient visit for evaluation and management, established patient, moderate complexity',
+                description='Office visit for established patient, moderate complexity',
                 code_type='CPT',
-                additional_info='Typically 30-39 minutes with patient',
+                additional_info='Typically 30-39 minutes. For moderately complex medical issues.',
+                category='E&M Services'
+            ),
+            '99215': CodeInfo(
+                code='99215',
+                description='Office visit for established patient, high complexity',
+                code_type='CPT',
+                additional_info='Typically 40-54 minutes. For complex medical decision making.',
                 category='E&M Services'
             ),
             '93000': CodeInfo(
                 code='93000',
-                description='Electrocardiogram, routine ECG with at least 12 leads; with interpretation and report',
+                description='Electrocardiogram (ECG/EKG) with interpretation',
                 code_type='CPT',
-                additional_info='Complete ECG service including tracing, interpretation, and report',
+                additional_info='Complete 12-lead ECG including test, interpretation, and report.',
                 category='Cardiovascular'
             ),
             'DRG470': CodeInfo(
                 code='DRG470',
-                description='Major hip and knee joint replacement or reattachment of lower extremity without MCC',
+                description='Major hip and knee joint replacement without complications',
                 code_type='DRG',
-                additional_info='Medicare Severity-Diagnosis Related Group for joint replacement procedures',
+                additional_info='Medicare payment group for joint replacement surgeries.',
                 category='Orthopedic'
             ),
             'Z79.899': CodeInfo(
                 code='Z79.899',
-                description='Other long term (current) drug therapy',
+                description='Other long term drug therapy',
                 code_type='ICD-10',
-                additional_info='Used to indicate patient is on long-term medication therapy',
-                category='Factors influencing health status'
+                additional_info='Indicates patient is on long-term medication.',
+                category='Diagnosis'
             ),
             'E1399': CodeInfo(
                 code='E1399',
                 description='Durable medical equipment, miscellaneous',
                 code_type='HCPCS',
-                additional_info='DME not otherwise classified, requires detailed description',
-                category='DME'
-            ),
-            'G0442': CodeInfo(
-                code='G0442',
-                description='Annual alcohol screening, 5 to 10 minutes',
-                code_type='HCPCS',
-                additional_info='Medicare-covered annual alcohol misuse screening',
-                category='Preventive Services'
-            ),
-            '90837': CodeInfo(
-                code='90837',
-                description='Psychotherapy, 60 minutes with patient',
-                code_type='CPT',
-                additional_info='Individual psychotherapy session, 53-60 minutes',
-                category='Psychiatric Services'
+                additional_info='For DME not elsewhere classified.',
+                category='Equipment'
             ),
             'J3420': CodeInfo(
                 code='J3420',
-                description='Injection, vitamin B-12 cyanocobalamin, up to 1000 mcg',
+                description='Vitamin B-12 injection',
                 code_type='HCPCS',
-                additional_info='Vitamin B12 injection for treatment of deficiency',
-                category='Drug Administration'
+                additional_info='Cyanocobalamin up to 1000 mcg.',
+                category='Injections'
             ),
             '80053': CodeInfo(
                 code='80053',
                 description='Comprehensive metabolic panel',
                 code_type='CPT',
-                additional_info='14 tests including glucose, kidney function, liver function, and electrolytes',
+                additional_info='14 blood tests including glucose, kidney, and liver function.',
                 category='Laboratory'
             ),
-            'A0425': CodeInfo(
-                code='A0425',
-                description='Ground mileage, per statute mile',
-                code_type='HCPCS',
-                additional_info='Ambulance mileage for ground transport',
-                category='Ambulance Services'
-            ),
-            '99215': CodeInfo(
-                code='99215',
-                description='Office visit, established patient, high complexity',
-                code_type='CPT',
-                additional_info='Typically 40-54 minutes with patient',
-                category='E&M Services'
-            ),
             '70450': CodeInfo(
                 code='70450',
-                description='CT scan head or brain without contrast',
+                description='CT head/brain without contrast',
                 code_type='CPT',
-                additional_info='Computed tomography of head/brain without contrast material',
+                additional_info='Computed tomography of head without contrast material.',
                 category='Radiology'
+            ),
+            '90837': CodeInfo(
+                code='90837',
+                description='Psychotherapy, 60 minutes',
+                code_type='CPT',
+                additional_info='Individual psychotherapy session.',
+                category='Mental Health'
+            ),
+            '36415': CodeInfo(
+                code='36415',
+                description='Venipuncture (blood draw)',
+                code_type='CPT',
+                additional_info='Collection of blood by needle.',
+                category='Laboratory'
+            ),
+            '99282': CodeInfo(
+                code='99282',
+                description='Emergency department visit, low-moderate severity',
+                code_type='CPT',
+                additional_info='ED visit for problems of low to moderate severity.',
+                category='Emergency'
             )
         }
-        
-        self.codes_cache = sample_codes
-        self.cache_timestamp = datetime.now()
-        return True
     
-    def lookup_code(self, code: str) -> Optional[CodeInfo]:
-        """Look up a single code"""
+    def lookup(self, code: str) -> Optional[CodeInfo]:
         code = code.strip().upper()
-        
-        if code in self.codes_cache:
-            return self.codes_cache[code]
-        
-        # Try with DRG prefix
-        if code.startswith('DRG') and len(code) == 6:
-            drg_code = code[3:]
-            if drg_code in self.codes_cache:
-                return self.codes_cache[drg_code]
-        
-        # Try adding DRG prefix
-        if len(code) == 3 and code.isdigit():
+        if code in self.codes:
+            return self.codes[code]
+        if code.isdigit() and len(code) == 3:
             drg_code = f"DRG{code}"
-            if drg_code in self.codes_cache:
-                return self.codes_cache[drg_code]
-        
+            if drg_code in self.codes:
+                return self.codes[drg_code]
         return None
     
-    def get_statistics(self) -> Dict[str, Any]:
-        """Get statistics about loaded codes"""
-        stats = {
-            'total_codes': len(self.codes_cache),
-            'source': self.data_source.value,
-            'codes_by_type': {}
-        }
-        
-        for code_info in self.codes_cache.values():
-            code_type = code_info.code_type
-            stats['codes_by_type'][code_type] = stats['codes_by_type'].get(code_type, 0) + 1
-        
-        return stats
-
-# ============= Code Classifier =============
-
-class CodeClassifier:
-    """Classifies and validates healthcare billing codes"""
-    
-    PATTERNS = {
-        CodeType.CPT: re.compile(r'^[0-9]{5}$'),
-        CodeType.HCPCS: re.compile(r'^[A-V][0-9]{4}$', re.IGNORECASE),
-        CodeType.ICD10: re.compile(r'^[A-Z][0-9]{2}\.?[0-9]{0,3}$', re.IGNORECASE),
-        CodeType.DRG: re.compile(r'^(?:DRG)?[0-9]{3}$', re.IGNORECASE),
-    }
-    
-    @classmethod
-    def classify_code(cls, code: str) -> CodeType:
-        """Classify a code based on its format"""
-        code = code.strip().upper()
-        
-        if code.startswith('DRG'):
-            code = code[3:]
-            if cls.PATTERNS[CodeType.DRG].match(f"DRG{code}"):
-                return CodeType.DRG
-        
-        for code_type, pattern in cls.PATTERNS.items():
-            if pattern.match(code):
-                return code_type
-        
-        return CodeType.UNKNOWN
-    
-    @classmethod
-    def extract_codes_from_text(cls, text: str) -> List[Tuple[str, CodeType]]:
-        """Extract potential codes from free-form text"""
-        codes = []
-        
+    def search_codes(self, text: str) -> List[str]:
+        """Extract potential billing codes from text"""
+        found_codes = []
         patterns = [
             r'\b([A-V][0-9]{4})\b',  # HCPCS
             r'\b([0-9]{5})\b',  # CPT
@@ -287,300 +173,466 @@ class CodeClassifier:
         ]
         
         for pattern in patterns:
-            matches = re.findall(pattern, text, re.IGNORECASE)
+            matches = re.findall(pattern, text.upper())
             for match in matches:
-                code_type = cls.classify_code(match)
-                if code_type != CodeType.UNKNOWN:
-                    codes.append((match.upper(), code_type))
+                if self.lookup(match):
+                    found_codes.append(match)
         
-        return codes
+        return found_codes
 
-# ============= OpenRouter Client =============
+# ============= AI Assistant Class =============
 
-class OpenRouterClient:
-    """Client for OpenRouter LLM API"""
-    
+class HybridAIAssistant:
     def __init__(self):
-        self.api_key = os.getenv('OPENROUTER_API_KEY', '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')
-        self.base_url = 'https://openrouter.ai/api/v1'
-        self.model = 'openai/gpt-3.5-turbo'
-        self.temperature = 0.3
+        self.api_key = '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'
+        self.billing_db = BillingCodesDB()
+        self.context = ConversationContext()
         
         self.headers = {
             'Authorization': f'Bearer {self.api_key}',
             'Content-Type': 'application/json',
             'HTTP-Referer': 'https://huggingface.co',
-            'X-Title': 'Healthcare Billing Chatbot'
+            'X-Title': 'Hybrid AI Assistant'
         }
     
-    def generate_explanation(self, 
-                            code: str, 
-                            code_type: CodeType,
-                            code_info: Optional[CodeInfo],
-                            length: ExplanationLength,
-                            context: Optional[ConversationContext] = None) -> str:
-        """Generate explanation using LLM"""
-        
-        if not code_info:
-            return f"I couldn't find code {code} in our database. This might be an invalid code or one that's not in our current dataset."
-        
-        prompt = self._build_prompt(code, code_type, code_info, length, context)
-        
-        max_tokens = {
-            ExplanationLength.SHORT: 150,
-            ExplanationLength.DETAILED: 300,
-            ExplanationLength.EXTRA: 500
-        }.get(length, 200)
-        
-        payload = {
-            'model': self.model,
-            'messages': [
-                {'role': 'system', 'content': self._get_system_prompt()},
-                {'role': 'user', 'content': prompt}
-            ],
-            'temperature': self.temperature,
-            'max_tokens': max_tokens
-        }
+    def detect_intent(self, message: str) -> Dict[str, Any]:
+        """Detect if the message is about billing codes or general conversation"""
+        lower_msg = message.lower()
         
-        try:
-            response = requests.post(
-                f'{self.base_url}/chat/completions',
-                headers=self.headers,
-                json=payload,
-                timeout=30
-            )
-            response.raise_for_status()
-            
-            result = response.json()
-            explanation = result['choices'][0]['message']['content']
-            return explanation
-            
-        except Exception as e:
-            logger.error(f"OpenRouter API error: {e}")
-            return self._get_fallback_explanation(code, code_type, code_info, length)
-    
-    def _get_system_prompt(self) -> str:
-        return """You are a healthcare billing expert assistant. Provide accurate, 
-        fact-based explanations of healthcare billing codes. Only provide information 
-        you are certain about. Be concise but informative. Use plain language."""
-    
-    def _build_prompt(self, code: str, code_type: CodeType, code_info: CodeInfo, 
-                     length: ExplanationLength, context: Optional[ConversationContext]) -> str:
+        # Check for billing codes in the message
+        codes = self.billing_db.search_codes(message)
         
-        base_prompt = f"""Healthcare Billing Code: {code}
-Type: {code_type.value}
-Official Description: {code_info.description}
-Additional Information: {code_info.additional_info or 'N/A'}
-"""
+        # Keywords that suggest billing/medical coding questions
+        billing_keywords = ['code', 'cpt', 'hcpcs', 'icd', 'drg', 'billing', 'medical code', 
+                          'healthcare code', 'diagnosis code', 'procedure code']
         
-        if length == ExplanationLength.SHORT:
-            base_prompt += "\nProvide a SHORT explanation (3-4 lines) in simple terms."
-        elif length == ExplanationLength.DETAILED:
-            base_prompt += "\nProvide a DETAILED explanation including what it covers and when it's typically used."
-        elif length == ExplanationLength.EXTRA:
-            base_prompt += "\nProvide EXTENDED details with sections for: typical use cases, coverage warnings, and important notes."
+        is_billing = any(keyword in lower_msg for keyword in billing_keywords) or len(codes) > 0
         
-        return base_prompt
+        return {
+            'is_billing': is_billing,
+            'codes_found': codes,
+            'message': message
+        }
     
-    def _get_fallback_explanation(self, code: str, code_type: CodeType, 
-                                 code_info: CodeInfo, length: ExplanationLength) -> str:
-        if length == ExplanationLength.SHORT:
-            return f"{code} ({code_type.value}): {code_info.description[:150]}"
+    def handle_billing_query(self, message: str, codes: List[str]) -> str:
+        """Handle healthcare billing specific queries"""
+        responses = []
+        
+        if codes:
+            for code in codes[:3]:  # Limit to first 3 codes
+                info = self.billing_db.lookup(code)
+                if info:
+                    response = f"**{info.code} ({info.code_type})**\n"
+                    response += f"📋 **Description:** {info.description}\n"
+                    if info.additional_info:
+                        response += f"ℹ️ **Details:** {info.additional_info}\n"
+                    if info.category:
+                        response += f"🏷️ **Category:** {info.category}\n"
+                    responses.append(response)
+        
+        if responses:
+            final_response = "I found information about the billing code(s) you mentioned:\n\n"
+            final_response += "\n---\n".join(responses)
+            final_response += "\n\n💡 **Need more details?** Feel free to ask specific questions about these codes!"
+            return final_response
         else:
-            explanation = f"{code} ({code_type.value}): {code_info.description}"
-            if code_info.additional_info:
-                explanation += f"\n\nAdditional Details: {code_info.additional_info}"
-            return explanation
-
-# ============= Main Chatbot Class =============
-
-class HealthcareBillingChatbot:
-    """Main chatbot class"""
+            return self.get_general_response(message, billing_context=True)
     
-    def __init__(self):
-        self.loader = HCPCSLoader()
-        self.classifier = CodeClassifier()
-        self.llm_client = OpenRouterClient()
-        self.context = ConversationContext()
-        self.loader.load_data()
-    
-    def process_input(self, user_input: str) -> str:
-        """Process user input and generate response"""
-        self.context.turn_count += 1
+    def get_general_response(self, message: str, billing_context: bool = False) -> str:
+        """Get response from OpenRouter API for general queries"""
         
-        codes = self.classifier.extract_codes_from_text(user_input)
-        intent = self._determine_intent(user_input, codes)
-        response = self._generate_response(user_input, codes, intent)
+        # Prepare system prompt
+        system_prompt = """You are a helpful, friendly AI assistant with expertise in healthcare billing codes. 
+        You can assist with any topic - from casual conversation to complex questions. 
+        When discussing medical billing codes, you provide accurate, detailed information.
+        Be conversational, helpful, and engaging. Use emojis occasionally to be friendly.
+        Keep responses concise but informative."""
         
-        self.context.conversation_history.append({
-            'user': user_input,
-            'assistant': response
-        })
+        if billing_context:
+            system_prompt += "\nThe user is asking about medical billing. Provide helpful information even if you don't have specific code details."
         
-        if len(self.context.conversation_history) > 10:
-            self.context.conversation_history = self.context.conversation_history[-10:]
+        # Build conversation history for context
+        messages = [{'role': 'system', 'content': system_prompt}]
         
-        return response
-    
-    def _determine_intent(self, user_input: str, codes: List[Tuple[str, CodeType]]) -> str:
-        """Determine user intent"""
-        lower_input = user_input.lower()
-        
-        if self.context.current_code:
-            if any(keyword in lower_input for keyword in ['more', 'detail', 'explain']):
-                return 'more_details'
+        # Add recent conversation history (last 5 exchanges)
+        for msg in self.context.messages[-10:]:
+            messages.append(msg)
         
-        if any(keyword in lower_input for keyword in ['short', 'brief', 'quick']):
-            return 'explain_short'
-        if any(keyword in lower_input for keyword in ['detail', 'comprehensive']):
-            return 'explain_detailed'
+        # Add current message
+        messages.append({'role': 'user', 'content': message})
         
-        if codes:
-            return 'explain_code'
-        
-        return 'general_query'
-    
-    def _generate_response(self, user_input: str, codes: List[Tuple[str, CodeType]], intent: str) -> str:
-        """Generate response based on intent"""
-        
-        if codes and intent in ['explain_code', 'explain_short', 'explain_detailed']:
-            code, code_type = codes[0]
-            code_info = self.loader.lookup_code(code)
-            
-            if not code_info:
-                return f"I couldn't find code {code} in our database. Please check the code and try again."
-            
-            self.context.current_code = code
-            self.context.current_code_type = code_type
-            self.context.current_code_info = code_info
-            
-            length = ExplanationLength.SHORT
-            if intent == 'explain_detailed':
-                length = ExplanationLength.DETAILED
-            
-            explanation = self.llm_client.generate_explanation(
-                code, code_type, code_info, length, self.context
+        try:
+            response = requests.post(
+                'https://openrouter.ai/api/v1/chat/completions',
+                headers=self.headers,
+                json={
+                    'model': 'openai/gpt-3.5-turbo',
+                    'messages': messages,
+                    'temperature': 0.7,
+                    'max_tokens': 500
+                },
+                timeout=30
             )
             
-            if length != ExplanationLength.EXTRA:
-                explanation += "\n\nWould you like more details about this code?"
-            
-            return explanation
-        
-        if self.context.current_code and intent == 'more_details':
-            if not self.context.current_code_info:
-                return "Please specify a code first."
-            
-            return self.llm_client.generate_explanation(
-                self.context.current_code,
-                self.context.current_code_type,
-                self.context.current_code_info,
-                ExplanationLength.EXTRA,
-                self.context
-            )
-        
-        return self._handle_general_query(user_input)
+            if response.status_code == 200:
+                result = response.json()
+                ai_response = result['choices'][0]['message']['content']
+                
+                # Update context
+                self.context.messages.append({'role': 'user', 'content': message})
+                self.context.messages.append({'role': 'assistant', 'content': ai_response})
+                
+                # Keep only last 20 messages in context
+                if len(self.context.messages) > 20:
+                    self.context.messages = self.context.messages[-20:]
+                
+                return ai_response
+            else:
+                logger.error(f"API error: {response.status_code}")
+                return self.get_fallback_response(message)
+                
+        except Exception as e:
+            logger.error(f"Request failed: {e}")
+            return self.get_fallback_response(message)
     
-    def _handle_general_query(self, user_input: str) -> str:
-        """Handle general queries"""
-        lower_input = user_input.lower()
-        
-        if any(keyword in lower_input for keyword in ['help', 'how']):
-            return """I can help you understand healthcare billing codes! You can:
-• Ask about specific codes (e.g., "What is HCPCS A0429?")
-• Request short or detailed explanations
-• Ask follow-up questions about codes we've discussed
-
-Just type a code or ask a question to get started!"""
+    def get_fallback_response(self, message: str) -> str:
+        """Fallback responses when API fails"""
+        fallbacks = [
+            "I'm having trouble connecting right now, but I'm still here to help! Could you rephrase your question?",
+            "Let me think about that differently. What specific aspect would you like to know more about?",
+            "That's an interesting question! While I process that, is there anything specific you'd like to explore?",
+            "I'm here to help! Could you provide a bit more detail about what you're looking for?"
+        ]
+        return random.choice(fallbacks)
+    
+    def process_message(self, message: str) -> str:
+        """Main method to process any message"""
+        if not message.strip():
+            return "Feel free to ask me anything! I can help with general questions or healthcare billing codes. 😊"
         
-        if any(keyword in lower_input for keyword in ['hello', 'hi', 'hey']):
-            return "Hello! I'm here to help you understand healthcare billing codes. Feel free to ask about any CPT, HCPCS, ICD-10, or DRG code."
+        # Detect intent
+        intent = self.detect_intent(message)
         
-        return "I'm designed to help explain healthcare billing codes. Please provide a specific code (like A0429, 99213, or DRG470) and I'll give you a clear explanation."
+        # Route to appropriate handler
+        if intent['is_billing'] and intent['codes_found']:
+            return self.handle_billing_query(message, intent['codes_found'])
+        else:
+            return self.get_general_response(message, billing_context=intent['is_billing'])
+    
+    def reset_context(self):
+        """Reset conversation context"""
+        self.context = ConversationContext()
 
 # ============= Gradio Interface =============
 
-def create_gradio_interface():
-    """Create and return the Gradio interface"""
-    chatbot = HealthcareBillingChatbot()
+def create_interface():
+    assistant = HybridAIAssistant()
     
-    def process_message(message: str, history: List[Tuple[str, str]]) -> Tuple[str, List[Tuple[str, str]]]:
-        """Process a message and update chat history"""
-        if not message:
-            return "", history
-        
-        try:
-            response = chatbot.process_input(message)
-            history.append((message, response))
-            return "", history
-        except Exception as e:
-            logger.error(f"Error: {e}")
-            error_msg = "I apologize, but I encountered an error. Please try again."
-            history.append((message, error_msg))
-            return "", history
+    # ChatGPT-style CSS
+    custom_css = """
+    /* Main container */
+    .gradio-container {
+        font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Helvetica Neue', Arial, sans-serif !important;
+        max-width: 900px !important;
+        margin: auto !important;
+        background: #ffffff !important;
+    }
     
-    def clear_conversation():
-        """Clear the conversation"""
-        chatbot.context = ConversationContext()
-        return None
+    /* Header styling */
+    .header-container {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        padding: 2rem;
+        border-radius: 15px 15px 0 0;
+        margin-bottom: 0;
+        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
+    }
     
-    # Create the interface
-    with gr.Blocks(title="Healthcare Billing Chatbot", theme=gr.themes.Soft()) as interface:
-        gr.Markdown(
-            """
-            # 🏥 Healthcare Billing Chatbot
-            
-            Ask me about healthcare billing codes (CPT, HCPCS, ICD-10, DRG) and I'll provide clear explanations!
-            
-            ### Example Queries:
-            - "What is HCPCS code A0429?"
-            - "Explain CPT 99213 in detail"
-            - "Tell me about DRG 470"
-            """
-        )
+    .header-title {
+        color: white;
+        font-size: 2rem;
+        font-weight: 700;
+        margin: 0;
+        display: flex;
+        align-items: center;
+        justify-content: center;
+        gap: 0.5rem;
+    }
+    
+    .header-subtitle {
+        color: rgba(255,255,255,0.9);
+        font-size: 1rem;
+        margin-top: 0.5rem;
+        text-align: center;
+    }
+    
+    /* Chat container */
+    #chatbot {
+        height: 500px !important;
+        border: 1px solid #e5e7eb !important;
+        border-radius: 12px !important;
+        box-shadow: 0 2px 6px rgba(0,0,0,0.05) !important;
+        background: #ffffff !important;
+    }
+    
+    /* Message styling */
+    .message {
+        padding: 1rem !important;
+        margin: 0.5rem !important;
+        border-radius: 12px !important;
+        font-size: 15px !important;
+        line-height: 1.6 !important;
+    }
+    
+    .user-message {
+        background: #f3f4f6 !important;
+        border: 1px solid #e5e7eb !important;
+        margin-left: 20% !important;
+    }
+    
+    .bot-message {
+        background: #ffffff !important;
+        border: 1px solid #e5e7eb !important;
+        margin-right: 20% !important;
+    }
+    
+    /* Input area */
+    #input-box {
+        border: 2px solid #e5e7eb !important;
+        border-radius: 12px !important;
+        padding: 14px 16px !important;
+        font-size: 15px !important;
+        transition: all 0.3s ease !important;
+        background: #ffffff !important;
+    }
+    
+    #input-box:focus {
+        border-color: #667eea !important;
+        box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
+        outline: none !important;
+    }
+    
+    /* Buttons */
+    .primary-btn {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+        color: white !important;
+        border: none !important;
+        border-radius: 10px !important;
+        padding: 12px 24px !important;
+        font-weight: 600 !important;
+        font-size: 15px !important;
+        cursor: pointer !important;
+        transition: transform 0.2s ease !important;
+    }
+    
+    .primary-btn:hover {
+        transform: translateY(-1px) !important;
+        box-shadow: 0 4px 12px rgba(102, 126, 234, 0.3) !important;
+    }
+    
+    .secondary-btn {
+        background: #f3f4f6 !important;
+        color: #374151 !important;
+        border: 1px solid #e5e7eb !important;
+        border-radius: 10px !important;
+        padding: 10px 20px !important;
+        font-weight: 500 !important;
+        cursor: pointer !important;
+        transition: all 0.2s ease !important;
+    }
+    
+    .secondary-btn:hover {
+        background: #e5e7eb !important;
+        border-color: #d1d5db !important;
+    }
+    
+    /* Example chips */
+    .example-chip {
+        display: inline-block !important;
+        background: #ffffff !important;
+        border: 1px solid #e5e7eb !important;
+        border-radius: 20px !important;
+        padding: 8px 16px !important;
+        margin: 4px !important;
+        font-size: 14px !important;
+        color: #4b5563 !important;
+        cursor: pointer !important;
+        transition: all 0.2s ease !important;
+    }
+    
+    .example-chip:hover {
+        background: #f9fafb !important;
+        border-color: #667eea !important;
+        color: #667eea !important;
+        transform: translateY(-1px) !important;
+    }
+    
+    /* Info cards */
+    .info-card {
+        background: linear-gradient(135deg, #f6f8fb 0%, #f1f5f9 100%);
+        border: 1px solid #e5e7eb;
+        border-radius: 12px;
+        padding: 1rem;
+        margin: 1rem 0;
+    }
+    
+    /* Responsive design */
+    @media (max-width: 768px) {
+        .gradio-container {
+            padding: 0 !important;
+        }
         
-        chatbot_ui = gr.Chatbot(label="Conversation", height=400)
+        .header-title {
+            font-size: 1.5rem;
+        }
+        
+        .user-message, .bot-message {
+            margin-left: 5% !important;
+            margin-right: 5% !important;
+        }
+    }
+    """
+    
+    with gr.Blocks(css=custom_css, theme=gr.themes.Base()) as app:
+        # Header
+        gr.HTML("""
+            <div class="header-container">
+                <h1 class="header-title">
+                    <span>🤖</span>
+                    <span>AI Assistant</span>
+                    <span style="font-size: 0.8em; background: rgba(255,255,255,0.2); padding: 2px 8px; border-radius: 12px;">PLUS</span>
+                </h1>
+                <p class="header-subtitle">Your intelligent companion for any question + Healthcare Billing Expert</p>
+            </div>
+        """)
+        
+        # Main chat interface
+        chatbot_ui = gr.Chatbot(
+            value=[[None, "👋 **Hello! I'm your AI Assistant!**\n\nI can help you with:\n\n🏥 **Healthcare Billing Codes** - I'm an expert in CPT, HCPCS, ICD-10, and DRG codes\n💬 **General Conversation** - Ask me anything!\n📚 **Learning & Education** - Help with various topics\n✍️ **Writing & Creation** - Stories, emails, ideas\n🔧 **Problem Solving** - Let's work through challenges together\n\n**Try asking:**\n• 'What is billing code A0429?'\n• 'Help me write an email'\n• 'Explain quantum physics simply'\n• 'What's the weather like?'\n\nHow can I assist you today? 😊"]],
+            elem_id="chatbot",
+            show_label=False,
+            type="messages",
+            bubble_full_width=False,
+            height=500
+        )
         
+        # Input section
         with gr.Row():
             msg = gr.Textbox(
-                label="Your Message",
-                placeholder="Type a billing code or ask a question...",
-                lines=2,
-                scale=4
+                placeholder="Ask me anything... (e.g., 'Explain code 99213' or 'Help me write a story')",
+                show_label=False,
+                elem_id="input-box",
+                scale=5,
+                lines=1,
+                max_lines=5
             )
-            submit_btn = gr.Button("Send", variant="primary", scale=1)
+            send_btn = gr.Button("Send", elem_classes="primary-btn", scale=1)
+        
+        # Quick examples
+        gr.HTML("<div style='text-align: center; margin: 1rem 0; color: #6b7280; font-size: 14px;'>Quick Examples</div>")
         
         with gr.Row():
-            clear_btn = gr.Button("🔄 Clear Conversation")
-            
-            with gr.Column():
-                gr.Examples(
-                    examples=[
-                        "What is HCPCS A0429?",
-                        "Explain CPT 99213",
-                        "Tell me about DRG 470",
-                        "What does code 93000 mean?",
-                        "Give me more details"
-                    ],
-                    inputs=msg
-                )
+            ex_col1 = gr.Column(scale=1)
+            ex_col2 = gr.Column(scale=1)
+            ex_col3 = gr.Column(scale=1)
+        
+        with ex_col1:
+            gr.HTML("<div style='color: #667eea; font-weight: 600; font-size: 13px; margin-bottom: 8px;'>🏥 Medical Billing</div>")
+            ex1 = gr.Button("What is code A0429?", elem_classes="example-chip", size="sm")
+            ex2 = gr.Button("Explain CPT 99213", elem_classes="example-chip", size="sm")
+            ex3 = gr.Button("DRG 470 details", elem_classes="example-chip", size="sm")
+        
+        with ex_col2:
+            gr.HTML("<div style='color: #667eea; font-weight: 600; font-size: 13px; margin-bottom: 8px;'>💭 General Questions</div>")
+            ex4 = gr.Button("How does AI work?", elem_classes="example-chip", size="sm")
+            ex5 = gr.Button("Recipe for pasta", elem_classes="example-chip", size="sm")
+            ex6 = gr.Button("Python tutorial", elem_classes="example-chip", size="sm")
+        
+        with ex_col3:
+            gr.HTML("<div style='color: #667eea; font-weight: 600; font-size: 13px; margin-bottom: 8px;'>✍️ Creative Help</div>")
+            ex7 = gr.Button("Write a poem", elem_classes="example-chip", size="sm")
+            ex8 = gr.Button("Email template", elem_classes="example-chip", size="sm")
+            ex9 = gr.Button("Story ideas", elem_classes="example-chip", size="sm")
+        
+        # Control buttons
+        with gr.Row():
+            clear_btn = gr.Button("🔄 New Chat", elem_classes="secondary-btn", size="sm")
+            gr.HTML("<div style='flex-grow: 1;'></div>")
+            gr.HTML("""
+                <div style='text-align: right; color: #6b7280; font-size: 12px;'>
+                    Powered by GPT-3.5 • Healthcare Billing Database
+                </div>
+            """)
+        
+        # Footer info
+        gr.HTML("""
+            <div class="info-card" style="margin-top: 2rem;">
+                <div style="display: flex; justify-content: space-around; text-align: center;">
+                    <div>
+                        <div style="color: #667eea; font-size: 24px; font-weight: bold;">15+</div>
+                        <div style="color: #6b7280; font-size: 12px;">Medical Codes</div>
+                    </div>
+                    <div>
+                        <div style="color: #667eea; font-size: 24px; font-weight: bold;">∞</div>
+                        <div style="color: #6b7280; font-size: 12px;">Topics</div>
+                    </div>
+                    <div>
+                        <div style="color: #667eea; font-size: 24px; font-weight: bold;">24/7</div>
+                        <div style="color: #6b7280; font-size: 12px;">Available</div>
+                    </div>
+                    <div>
+                        <div style="color: #667eea; font-size: 24px; font-weight: bold;">Fast</div>
+                        <div style="color: #6b7280; font-size: 12px;">Responses</div>
+                    </div>
+                </div>
+            </div>
+        """)
         
         # Event handlers
-        msg.submit(process_message, [msg, chatbot_ui], [msg, chatbot_ui])
-        submit_btn.click(process_message, [msg, chatbot_ui], [msg, chatbot_ui])
-        clear_btn.click(clear_conversation, outputs=[chatbot_ui])
-        
-        stats = chatbot.loader.get_statistics()
-        gr.Markdown(
-            f"""
-            ---
-            **Database Info:** {stats['total_codes']} codes loaded | 
-            Types: {', '.join(stats['codes_by_type'].keys())}
-            """
-        )
+        def respond(message, chat_history):
+            if not message.strip():
+                return "", chat_history
+            
+            # Process message
+            response = assistant.process_message(message)
+            
+            # Update chat history
+            chat_history.append({"role": "user", "content": message})
+            chat_history.append({"role": "assistant", "content": response})
+            
+            return "", chat_history
+        
+        def clear_chat():
+            assistant.reset_context()
+            welcome = """👋 **Chat cleared! Ready for a new conversation.**
+
+I'm here to help with anything you need - from healthcare billing codes to general questions!
+
+What would you like to know? 😊"""
+            return [[None, welcome]]
+        
+        # Connect events
+        msg.submit(respond, [msg, chatbot_ui], [msg, chatbot_ui])
+        send_btn.click(respond, [msg, chatbot_ui], [msg, chatbot_ui])
+        clear_btn.click(clear_chat, outputs=[chatbot_ui])
+        
+        # Example button handlers
+        ex1.click(lambda: "What is healthcare billing code A0429?", outputs=msg)
+        ex2.click(lambda: "Can you explain CPT code 99213 in detail?", outputs=msg)
+        ex3.click(lambda: "Tell me about DRG 470", outputs=msg)
+        ex4.click(lambda: "How does artificial intelligence work?", outputs=msg)
+        ex5.click(lambda: "Give me a simple pasta recipe", outputs=msg)
+        ex6.click(lambda: "Teach me Python basics", outputs=msg)
+        ex7.click(lambda: "Write a short poem about nature", outputs=msg)
+        ex8.click(lambda: "Help me write a professional email template", outputs=msg)
+        ex9.click(lambda: "Give me creative story ideas", outputs=msg)
     
-    return interface
+    return app
 
-# Launch the app
-interface = create_gradio_interface()
-interface.launch()
+# Launch
+if __name__ == "__main__":
+    app = create_interface()
+    app.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )