src/services/groq_service.py

import json
import logging
from groq import Groq
from config.settings import settings
from utils.prompts import get_coding_prompt

logger = logging.getLogger(__name__)

class GroqService:
    def __init__(self):
        """Initialize Groq client with API key from settings"""
        if not settings.groq_api_key:
            logger.error("❌ GROQ_API_KEY is not set!")
            raise ValueError("GROQ_API_KEY environment variable is required")
        
        try:
            # Simple initialization without extra parameters
            self.client = Groq(
                api_key=settings.groq_api_key
            )
            self.model = settings.groq_model
            logger.info(f"✅ Groq client initialized successfully")
            logger.info(f"✅ Using model: {self.model}")
        except Exception as e:
            logger.error(f"❌ Failed to initialize Groq client: {str(e)}")
            raise
    
    def analyze_provider_notes(self, provider_notes: str) -> dict:
        """
        Analyze provider notes and extract ICD-10 and CPT codes
        
        Args:
            provider_notes: Clinical provider notes text
            response_format: Desired response format (default: "json")
            
        Returns:
            dict: Parsed coding response with ICD-10, CPT codes and summary
        """
        try:
            logger.info(f"📝 Analyzing provider notes (length: {len(provider_notes)})")
            
            # Get the prompt
            prompt = get_coding_prompt(provider_notes)
            
            # Call Groq API
            logger.info(f"🚀 Calling Groq API with model: {self.model}")
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {
                        "role": "system",
                        "content": "You are a medical coding expert specializing in ICD-10 and CPT coding. Always respond with valid JSON."
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                temperature=0.3,
                max_tokens=2000,
                response_format={"type": "json_object"}
            )
            
            # Extract response
            raw_response = response.choices[0].message.content
            logger.info(f"📥 Received response from Groq (length: {len(raw_response)})")
            logger.debug(f"Raw response: {raw_response[:500]}...")
            
            # Parse JSON
            parsed_response = json.loads(raw_response)
            
            # Validate and structure response
            result = self._structure_response(parsed_response)
            
            logger.info(f"✅ Successfully analyzed notes: {len(result.get('icd_codes', []))} ICD codes, {len(result.get('cpt_codes', []))} CPT codes")
            
            return result
            
        except json.JSONDecodeError as e:
            logger.error(f"❌ JSON parsing error: {str(e)}")
            logger.error(f"Raw response: {raw_response}")
            raise ValueError(f"Failed to parse Groq response as JSON: {str(e)}")
        
        except Exception as e:
            logger.error(f"❌ Error analyzing provider notes: {str(e)}")
            raise
    
    def _structure_response(self, parsed_response: dict) -> dict:
        """
        Structure and validate the response from Groq
        
        Args:
            parsed_response: Raw parsed JSON from Groq
            response_format: Desired response format (default: "json")
            
        Returns:
            dict: Properly structured response
        """
        # Handle different possible response formats
        icd_codes = []
        cpt_codes = []
        overall_summary = parsed_response.get("overall_summary", "")
        
        # Extract ICD codes
        raw_icd = parsed_response.get("icd_codes", [])
        if isinstance(raw_icd, list):
            for icd in raw_icd:
                if isinstance(icd, dict):
                    icd_codes.append({
                        "code": icd.get("code", ""),
                        "description": icd.get("description", ""),
                        "explanation": icd.get("explanation", "")
                    })
        
        # Extract CPT codes
        raw_cpt = parsed_response.get("cpt_codes", [])
        if isinstance(raw_cpt, list):
            for cpt in raw_cpt:
                if isinstance(cpt, dict):
                    cpt_codes.append({
                        "code": cpt.get("code", ""),
                        "description": cpt.get("description", ""),
                        "explanation": cpt.get("explanation", "")
                    })
        
        return {
            "icd_codes": icd_codes,
            "cpt_codes": cpt_codes,
            "overall_summary": overall_summary
        }

# Global instance
groq_service = GroqService()