llm.py

import os
import json
import re
from typing import Tuple, Dict, List
from concurrent.futures import ThreadPoolExecutor, TimeoutError as ThreadTimeout


# Option 1: Use Hugging Face Inference API (recommended for better quality)
# Option 2: Use larger local model
# Option 3: Use OpenAI/Anthropic API if available

DEBUG_MODE = os.getenv("DEBUG_MODE", "False").lower() == "true"
USE_HF_API = os.getenv("USE_HF_API", "False").lower() == "true"  # Set default to False
HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN", "")

#if HF_TOKEN:
    # huggingface_hub import login
   # login(token=HF_TOKEN)
def log(msg):
    if DEBUG_MODE:
        print(f"[LLM Debug] {msg}")


def get_system_prompt(interviewee_type: str, is_summary: bool = False) -> str:
    """Generate context-aware system prompts"""
    
    base_prompt = """You are an expert medical transcript analyzer specializing in healthcare interviews.

Your task is to extract structured, actionable insights from interview transcripts.

Core Principles:
- Focus on factual, verifiable medical information
- Distinguish between speaker roles accurately
- Filter out pleasantries, disclaimers, and off-topic content
- Extract specific medical terms, dosages, and treatment details
- Identify patterns and clinical reasoning
"""
    
    if is_summary:
        return base_prompt + """
    CROSS-INTERVIEW SYNTHESIS & VALIDATION TASK:
    
    You are analyzing multiple transcripts. Extract verified patterns and flag inconsistencies.
    
    STEP 1 - PATTERN IDENTIFICATION:
    For each theme, count occurrences across transcripts:
    - How many participants mentioned X? (e.g., "7 out of 10 participants")
    - Calculate percentages when relevant
    - What's the range of perspectives?
    
    STEP 2 - CLASSIFY BY CONSENSUS LEVEL:
    - STRONG CONSENSUS (80%+ agreement): Findings most participants agree on
    - MAJORITY VIEW (60-79%): Significant but not universal agreement
    - SPLIT PERSPECTIVES (40-59%): Where views diverge
    - OUTLIERS (<40%): Unique but noteworthy perspectives
    
    STEP 3 - CROSS-VALIDATE:
    - Check for contradictions between transcripts
    - Note where perspectives differ and why
    - Flag quality issues (brief transcripts, vague responses)
    
    STEP 4 - CITE EVIDENCE:
    - Reference specific transcript numbers
    - Include brief supporting quotes/details
    - Distinguish fact from interpretation
    
    OUTPUT FORMAT:
    Start with 2-3 sentence executive overview, then:
    
    **STRONG CONSENSUS FINDINGS:**
    [List with counts and evidence]
    
    **MAJORITY FINDINGS:**
    [List with counts]
    
    **DIVERGENT PERSPECTIVES:**
    [Where participants disagreed and context]
    
    **NOTABLE OUTLIERS:**
    [Unique but important points]
    
    **QUALITY NOTES:**
    [Any gaps or transcript issues]
    
    CRITICAL RULES:
    - NEVER use vague terms like "many," "most," "some" - always use specific numbers
    - ALWAYS cite transcript numbers for claims
    - FLAG weak evidence explicitly
    - Separate facts from interpretations
    - NO JSON output - write in clear narrative prose
    """
    
    if interviewee_type == "HCP":
        return base_prompt + """
Healthcare Professional Analysis Focus:
- Prescribing patterns and medication choices
- Diagnostic reasoning and clinical decision-making
- Treatment protocols and guidelines referenced
- Peer perspectives on efficacy and safety
- Barriers to treatment or adoption
- Off-label uses or emerging practices

Extract and structure:
1. Diagnoses mentioned with context
2. Prescriptions with dosage, frequency, and rationale
3. Treatment strategies and their justifications
4. Clinical guidelines or studies referenced
5. Challenges or barriers discussed
6. Key clinical insights or pearls
"""
    
    elif interviewee_type == "Patient":
        return base_prompt + """
Patient Interview Analysis Focus:
- Symptom descriptions and severity
- Treatment experiences and outcomes
- Side effects and tolerability
- Quality of life impacts
- Adherence challenges and enablers
- Emotional and psychological factors
- Healthcare system interactions

Extract and structure:
1. Primary symptoms with duration and severity
2. Current and past treatments
3. Treatment effectiveness and satisfaction
4. Side effects experienced
5. Concerns and unmet needs
6. Quality of life impacts
7. Support systems and resources
"""
    
    else:
        return base_prompt + """
General Interview Analysis Focus:
- Main themes and topics discussed
- Key insights and observations
- Recommendations or suggestions
- Contextual factors
- Areas of emphasis or concern

Extract and structure relevant information based on interview content.
"""


def build_extraction_template(interviewee_type: str) -> str:
    """Create JSON template for structured data extraction"""
    
    if interviewee_type == "HCP":
        return """{
  "diagnoses": ["condition 1", "condition 2"],
  "prescriptions": ["medication (dose, frequency, indication)"],
  "treatment_rationale": ["reason for treatment choice"],
  "guidelines_mentioned": ["guideline or study name"],
  "clinical_decisions": ["key clinical decision with reasoning"],
  "barriers": ["barrier to treatment"],
  "key_insights": ["notable clinical insight"]
}"""
    
    elif interviewee_type == "Patient":
        return """{
  "symptoms": ["symptom (severity, duration)"],
  "concerns": ["patient concern or question"],
  "treatments_current": ["current treatment"],
  "treatments_past": ["past treatment with outcome"],
  "treatment_response": ["description of how treatment is working"],
  "side_effects": ["side effect experienced"],
  "quality_of_life": ["impact on daily life"],
  "adherence_factors": ["factor affecting medication adherence"]
}"""
    
    else:
        return """{
  "key_insights": ["main insight or finding"],
  "themes": ["recurring theme"],
  "recommendations": ["recommendation or suggestion"],
  "context": ["important contextual information"]
}"""


def parse_structured_response(text: str, interviewee_type: str) -> Dict:
    """Extract structured data from LLM response"""
    
    # Try to find JSON block
    json_match = re.search(r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', text, re.DOTALL)
    
    if json_match:
        try:
            data = json.loads(json_match.group())
            log(f"Successfully extracted JSON: {data}")
            return data
        except json.JSONDecodeError:
            log("Failed to parse JSON from response")
    
    # Fallback: Extract from text using patterns
    data = {}
    
    if interviewee_type == "HCP":
        # Extract diagnoses
        diag_pattern = r'(?:diagnos[ei]s|condition):\s*([^\n]+)'
        data["diagnoses"] = re.findall(diag_pattern, text, re.IGNORECASE)
        
        # Extract prescriptions
        rx_pattern = r'(?:prescri[bp]\w*|medication):\s*([^\n]+)'
        data["prescriptions"] = re.findall(rx_pattern, text, re.IGNORECASE)
        
        # Extract treatment rationale
        treat_pattern = r'(?:treatment|therapy|rationale):\s*([^\n]+)'
        data["treatment_rationale"] = re.findall(treat_pattern, text, re.IGNORECASE)
        
    elif interviewee_type == "Patient":
        # Extract symptoms
        symptom_pattern = r'(?:symptom|complaint|experienc\w*):\s*([^\n]+)'
        data["symptoms"] = re.findall(symptom_pattern, text, re.IGNORECASE)
        
        # Extract concerns
        concern_pattern = r'(?:concern|worry|question|anxious):\s*([^\n]+)'
        data["concerns"] = re.findall(concern_pattern, text, re.IGNORECASE)
        
        # Extract side effects
        se_pattern = r'(?:side effect|adverse|reaction):\s*([^\n]+)'
        data["side_effects"] = re.findall(se_pattern, text, re.IGNORECASE)
    
    # Clean and deduplicate
    for key in data:
        data[key] = list(set([item.strip() for item in data[key] if item.strip()]))
    
    log(f"Extracted data from text: {data}")
    return data


def query_llm_hf_api(prompt: str, max_tokens: int = 500) -> str:
    """Use Hugging Face Inference API for better quality"""
    try:
        from huggingface_hub import InferenceClient
        
        client = InferenceClient(token=HF_TOKEN)
        
        # Use chat completions instead
        messages = [
            {"role": "system", "content": "You are an expert transcript analyzer. Provide detailed, structured analysis."},
            {"role": "user", "content": prompt}
        ]
        
        response = client.chat_completion(
            messages=messages,
            model="microsoft/Phi-3-mini-4k-instruct",
            max_tokens=max_tokens,
            temperature=0.3
        )
        
        return response.choices[0].message.content.strip()
        
    except Exception as e:
        import traceback
        full_error = traceback.format_exc()
        log(f"HF API error: {e}\n{full_error}")
        print(f"[HF API Full Error]\n{full_error}")  # Print to console
        return f"[Error] HF API failed: {e}"


def query_llm_local(prompt: str, max_tokens: int = 500) -> str:
    """Local model optimized for L4 GPU"""
    try:
        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
        import torch
        
        if not hasattr(query_llm_local, 'model'):
            log("Loading local model on L4...")
            query_llm_local.tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-xxl")
            query_llm_local.model = AutoModelForSeq2SeqLM.from_pretrained(
                "google/flan-t5-xxl",
                torch_dtype=torch.float16,
                device_map="auto"
            )
        
        # Tokenize and truncate to 512 tokens
        inputs = query_llm_local.tokenizer(
            prompt, 
            return_tensors="pt", 
            truncation=True, 
            max_length=512
        ).to("cuda")
        
        outputs = query_llm_local.model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            do_sample=False
        )
        
        response = query_llm_local.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response.strip()
        
    except Exception as e:
        log(f"Local model error: {e}")
        return f"[Error] Local model failed: {e}"


def query_llm(
    chunk: str,
    user_context: str,
    interviewee_type: str,
    extract_structured: bool = False,
    is_summary: bool = False,
    timeout: int = 120
) -> Tuple[str, Dict]:
    """
    Main LLM query function with structured extraction
    
    Returns:
        Tuple of (response_text, structured_data_dict)
    """
    
    system_prompt = get_system_prompt(interviewee_type, is_summary)
    extraction_template = build_extraction_template(interviewee_type) if extract_structured else ""
    
    # Build comprehensive prompt
    full_prompt = f"""{system_prompt}

User Instructions:
{user_context}

Transcript Segment to Analyze:
{chunk}

"""
    
    if extract_structured:
        full_prompt += f"""
IMPORTANT: Provide your analysis in two parts:
1. A clear narrative summary (3-5 sentences)
2. Structured data in this exact JSON format:
{extraction_template}

Be specific and include relevant details (dosages, durations, severity levels, etc.)
"""
    
    # Truncate if needed (but increased limit)
    max_prompt_length = 6000  # Increased from 2000
    if len(full_prompt) > max_prompt_length:
        chunk_limit = max_prompt_length - len(system_prompt) - len(user_context) - len(extraction_template) - 500
        chunk = chunk[:chunk_limit]
        full_prompt = f"{system_prompt}\n\nUser Instructions:\n{user_context}\n\nTranscript Segment:\n{chunk}\n\n"
        if extract_structured:
            full_prompt += f"Provide analysis and structured JSON: {extraction_template}"
        log(f"Prompt truncated to {len(full_prompt)} characters")
    
    def generate():
        if os.getenv("USE_LMSTUDIO", "False").lower() == "true":
            return query_llm_lmstudio(full_prompt, max_tokens=600)
        elif USE_HF_API and HF_TOKEN:
            return query_llm_hf_api(full_prompt, max_tokens=600)
        else:
            return query_llm_local(full_prompt, max_tokens=600)
    
    # Execute with timeout
    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(generate)
        try:
            response = future.result(timeout=timeout)
            log(f"LLM response received ({len(response)} chars)")
            
            # Extract structured data if requested
            structured_data = {}
            if extract_structured:
                structured_data = parse_structured_response(response, interviewee_type)
            
            return response, structured_data
            
        except ThreadTimeout:
            log("LLM generation timed out")
            return "[Error] LLM generation timed out.", {}
        except Exception as e:
            log(f"LLM generation failed: {e}")
            return f"[Error] LLM generation failed: {e}", {}
    

def extract_structured_data(text: str, interviewee_type: str) -> Dict:
    """
    Standalone function to extract structured data from existing text
    Useful for post-processing
    """
    return parse_structured_response(text, interviewee_type)