llm_service.py

"""
LLM Service for intelligent ranking and scoring of clinical trials.
Supports Hugging Face models including DeepSeek-V3.2.
"""
import os
import logging
from typing import List, Dict, Optional, Set
import re

# Conditional imports for local model mode (only needed if not using API)
torch = None
AutoTokenizer = None
AutoModelForCausalLM = None
pipeline = None

logger = logging.getLogger(__name__)

# Router statuses that should trigger trying the next endpoint format.
_ROUTER_FALLBACK_STATUSES = (400, 404, 405, 415, 422)
_DEFAULT_FALLBACK_MODELS = (
    "Qwen/Qwen2.5-7B-Instruct",
    "mistralai/Mistral-7B-Instruct-v0.3",
    "google/gemma-2-9b-it",
    "meta-llama/Llama-3.1-8B-Instruct",
)
_BIOMEDICAL_SYNONYMS = {
    "kras": ["k-ras", "kras g12c", "kras g12d", "krasi"],
    "egfr": ["erbb1", "epidermal growth factor receptor"],
    "nsclc": ["non-small cell lung cancer", "non small cell lung cancer"],
    "immunotherapy": ["checkpoint inhibitor", "pd-1", "pd-l1", "ctla-4", "io therapy"],
    "daraxonrasib": ["rmc-6236", "ras(onc) inhibitor", "ras on inhibitor"],
    "adagrasib": ["mrtx849"],
    "sotorasib": ["amg 510"],
    "pancreatic cancer": ["pdac", "pancreatic adenocarcinoma"],
}

def _import_local_dependencies():
    """Import torch and transformers only when needed for local model mode"""
    global torch, AutoTokenizer, AutoModelForCausalLM, pipeline
    if torch is None:
        try:
            import torch
            from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
        except ImportError as e:
            raise ImportError(
                "torch and transformers are required for local model mode. "
                "Install them with: pip install torch transformers accelerate sentencepiece\n"
                "Or use API mode by setting USE_HF_API=true"
            ) from e

class LLMService:
    """Service for interacting with Hugging Face LLM models"""
    
    def __init__(self, model_name: Optional[str] = None, use_api: bool = False, api_token: Optional[str] = None):
        """
        Initialize LLM service
        
        Args:
            model_name: Hugging Face model identifier (e.g., 'deepseek-ai/DeepSeek-V3.2')
                       If None, uses DEEPSEEK_MODEL env var or defaults to DeepSeek-V3.2
            use_api: If True, use Hugging Face Inference API instead of local model
            api_token: Hugging Face API token (required if use_api=True)
        """
        self.model_name = model_name or os.environ.get('DEEPSEEK_MODEL', 'deepseek-ai/DeepSeek-V3.2')
        self.use_api = use_api or os.environ.get('USE_HF_API', 'false').lower() == 'true'
        self.api_token = api_token or os.environ.get('HUGGINGFACE_API_TOKEN', '')
        fallback_models_raw = os.environ.get("HF_FALLBACK_MODELS", "")
        parsed_fallbacks = [m.strip() for m in fallback_models_raw.split(",") if m.strip()]
        fallback_candidates = parsed_fallbacks or list(_DEFAULT_FALLBACK_MODELS)
        self.fallback_models = [m for m in fallback_candidates if m != self.model_name]
        
        self.tokenizer = None
        self.model = None
        self.pipeline = None
        
        if not self.use_api:
            self._load_local_model()
        else:
            if not self.api_token:
                logger.warning("Hugging Face API token not provided. Set HUGGINGFACE_API_TOKEN env var.")
    
    def _load_local_model(self):
        """Load model locally using transformers"""
        try:
            # Import dependencies if not already imported
            _import_local_dependencies()
            
            logger.info(f"Loading model: {self.model_name}")
            
            # Check if CUDA is available
            device = "cuda" if torch.cuda.is_available() else "cpu"
            logger.info(f"Using device: {device}")
            
            # Load tokenizer and model
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_name,
                trust_remote_code=True
            )
            
            # Load model with appropriate settings
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                trust_remote_code=True,
                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
                device_map="auto" if device == "cuda" else None,
                low_cpu_mem_usage=True
            )
            
            if device == "cpu":
                self.model = self.model.to(device)
            
            # Create pipeline for easier text generation
            self.pipeline = pipeline(
                "text-generation",
                model=self.model,
                tokenizer=self.tokenizer,
                device=0 if device == "cuda" else -1,
                torch_dtype=torch.float16 if device == "cuda" else torch.float32
            )
            
            logger.info(f"Model {self.model_name} loaded successfully")
            
        except Exception as e:
            logger.error(f"Error loading model: {str(e)}")
            raise
    
    def rank_studies(self, studies: List[Dict], ranking_terms: str) -> List[Dict]:
        """
        Rank studies based on relevance to ranking terms using LLM
        
        Args:
            studies: List of study dictionaries
            ranking_terms: Terms to use for ranking (e.g., "KRAS mutation, immunotherapy")
        
        Returns:
            List of studies sorted by relevance score (highest first), with ranking_reasoning added
        """
        if not ranking_terms or not ranking_terms.strip():
            return studies
        
        if not studies:
            return studies
        
        try:
            # Score each study
            scored_studies = []
            for study in studies:
                score, reasoning, match_data = self._score_study(study, ranking_terms)
                study_with_score = study.copy()
                study_with_score['relevance_score'] = score
                study_with_score['ranking_reasoning'] = reasoning
                study_with_score['ranking_match_terms'] = match_data.get('matched_terms', [])
                study_with_score['ranking_match_fields'] = match_data.get('matched_fields', {})
                study_with_score['ranking_score_breakdown'] = match_data.get('score_breakdown', {})
                scored_studies.append(study_with_score)
            
            # Sort by score (highest first)
            scored_studies.sort(key=lambda x: x.get('relevance_score', 0), reverse=True)
            
            return scored_studies
            
        except Exception as e:
            logger.error(f"Error ranking studies: {str(e)}")
            # Return original studies if ranking fails
            return studies
    
    def _score_study(self, study: Dict, ranking_terms: str) -> tuple:
        """
        Score a single study's relevance to ranking terms and get reasoning
        
        Args:
            study: Study dictionary
            ranking_terms: Terms to match against
        
        Returns:
            Tuple of (relevance score (0.0 to 1.0), reasoning explanation)
        """
        try:
            # Build context from study
            study_text = self._build_study_context(study)
            hybrid_match = self._compute_hybrid_match(study, ranking_terms)
            
            # Create prompt for scoring with reasoning - more flexible and natural
            prompt = f"""You are a helpful medical research assistant. The user is looking for clinical trials that relate to: "{ranking_terms}"

Consider this clinical trial:
{study_text}

Think about how well this trial matches what the user is looking for. Consider:
- Direct matches (exact terms mentioned)
- Related concepts and synonyms
- Contextual relevance (even if exact terms aren't used)
- Overall alignment with the user's intent

Rate the relevance on a scale of 0.0 to 1.0, where:
- 0.9-1.0: Highly relevant, directly matches what the user wants
- 0.7-0.8: Very relevant, strong connection
- 0.5-0.6: Moderately relevant, some connection
- 0.3-0.4: Somewhat relevant, weak connection
- 0.0-0.2: Not very relevant

Be flexible and consider the user's intent, not just exact word matches.

Provide your response in this format:
SCORE: [number between 0.0 and 1.0]
REASONING: [brief, natural explanation of why this score was assigned]"""

            if self.use_api:
                llm_score, reasoning = self._score_with_reasoning_api(prompt)
            else:
                llm_score, reasoning = self._score_with_reasoning_local(prompt)

            # Keep ranking useful even when HF API/model compatibility fails.
            if self.use_api and self._is_api_failure_reason(reasoning):
                score = hybrid_match['hybrid_score']
                reasoning = (
                    "Hybrid keyword/synonym fallback score used because Hugging Face ranking failed. "
                    f"Details: {reasoning}"
                )
                score_breakdown = {
                    "hybrid_score": hybrid_match['hybrid_score'],
                    "llm_score": None,
                    "final_score": score,
                }
            else:
                score = (0.55 * llm_score) + (0.45 * hybrid_match['hybrid_score'])
                score_breakdown = {
                    "hybrid_score": hybrid_match['hybrid_score'],
                    "llm_score": llm_score,
                    "final_score": score,
                }
                if hybrid_match['matched_terms']:
                    reasoning = (
                        f"{reasoning} Hybrid signals matched: "
                        f"{', '.join(hybrid_match['matched_terms'][:8])}."
                    )
            
            # Ensure score is between 0 and 1
            score = max(0.0, min(1.0, float(score)))
            
            return score, reasoning, {
                "matched_terms": hybrid_match['matched_terms'],
                "matched_fields": hybrid_match['matched_fields'],
                "score_breakdown": score_breakdown,
            }
            
        except Exception as e:
            logger.error(f"Error scoring study {study.get('nctId', 'unknown')}: {str(e)}")
            return 0.5, "Scoring fallback used because the model response could not be processed.", {
                "matched_terms": [],
                "matched_fields": {},
                "score_breakdown": {
                    "hybrid_score": None,
                    "llm_score": None,
                    "final_score": 0.5,
                },
            }

    def _is_api_failure_reason(self, reasoning: str) -> bool:
        """Identify known HF API failure messages returned by scoring helpers."""
        if not reasoning:
            return False
        markers = (
            "Hugging Face token missing",
            "Hugging Face auth failed",
            "Hugging Face rate limit reached",
            "Model/endpoint mismatch",
            "Hugging Face API error",
            "API request failed",
        )
        return any(marker in reasoning for marker in markers)

    def _extract_query_terms(self, ranking_terms: str) -> List[str]:
        """Extract normalized query concepts (supports comma-separated concepts)."""
        raw_parts = [p.strip().lower() for p in (ranking_terms or "").split(",") if p.strip()]
        if raw_parts:
            return raw_parts

        token_terms = [
            t for t in re.findall(r"[a-z0-9][a-z0-9\-\+]*", (ranking_terms or "").lower())
            if len(t) > 2
        ]
        phrase = (ranking_terms or "").strip().lower()
        if phrase and " " in phrase:
            return [phrase] + token_terms
        return token_terms

    def _expand_term_variants(self, term: str) -> Set[str]:
        """Expand a term into aliases/synonyms for biomedical matching."""
        variants = {term}
        synonyms = _BIOMEDICAL_SYNONYMS.get(term.lower(), [])
        variants.update(s.lower() for s in synonyms)

        for canonical, alias_list in _BIOMEDICAL_SYNONYMS.items():
            alias_lower = [a.lower() for a in alias_list]
            if term.lower() == canonical or term.lower() in alias_lower:
                variants.add(canonical)
                variants.update(alias_lower)
        return variants

    def _find_matches_in_text(self, text: str, variants: Set[str]) -> Set[str]:
        """Return matched variants found in text with word-boundary checks."""
        if not text:
            return set()
        haystack = text.lower()
        matched = set()
        for variant in variants:
            pattern = r"\b" + re.escape(variant) + r"\b"
            if re.search(pattern, haystack):
                matched.add(variant)
        return matched

    def _compute_hybrid_match(self, study: Dict, ranking_terms: str) -> Dict:
        """Compute biomedical keyword/synonym relevance signals."""
        stop_words = {
            "a", "an", "and", "or", "the", "to", "for", "of", "in", "on", "with", "by"
        }
        query_terms = [t for t in self._extract_query_terms(ranking_terms) if t not in stop_words]
        if not query_terms:
            return {
                "hybrid_score": 0.5,
                "matched_terms": [],
                "matched_fields": {},
            }

        fields = {
            "title": study.get("title", ""),
            "conditions": " ".join(study.get("conditions", []) or []),
            "summary": study.get("briefSummary", ""),
            "inclusionCriteria": " ".join(study.get("inclusionCriteria", []) or []),
        }

        matched_canonical_terms = set()
        matched_variants_for_display = set()
        matched_fields = {}
        total_hits = 0

        for term in query_terms:
            variants = self._expand_term_variants(term)
            term_matched_any = False
            for field_name, field_text in fields.items():
                field_matches = self._find_matches_in_text(field_text, variants)
                if field_matches:
                    matched_fields.setdefault(field_name, set()).update(field_matches)
                    matched_variants_for_display.update(field_matches)
                    total_hits += len(field_matches)
                    term_matched_any = True
            if term_matched_any:
                matched_canonical_terms.add(term)

        concept_ratio = len(matched_canonical_terms) / max(len(query_terms), 1)
        field_coverage = len(matched_fields) / max(len(fields), 1)
        density = min(total_hits / 8.0, 1.0)
        hybrid_score = max(0.0, min(1.0, (0.65 * concept_ratio) + (0.2 * field_coverage) + (0.15 * density)))

        cleaned_fields = {
            field: sorted(values, key=len, reverse=True)
            for field, values in matched_fields.items()
        }

        return {
            "hybrid_score": hybrid_score,
            "matched_terms": sorted(matched_variants_for_display, key=len, reverse=True),
            "matched_fields": cleaned_fields,
        }
    
    def _build_study_context(self, study: Dict) -> str:
        """Build a text context from study data"""
        parts = []
        
        if study.get('title'):
            parts.append(f"Title: {study['title']}")
        
        if study.get('sponsor'):
            parts.append(f"Sponsor: {study['sponsor']}")
        
        if study.get('briefSummary'):
            parts.append(f"Summary: {study['briefSummary'][:500]}")  # Limit summary length
        
        if study.get('conditions'):
            parts.append(f"Conditions: {', '.join(study['conditions'])}")
        
        if study.get('inclusionCriteria'):
            inclusion_text = ' '.join(study['inclusionCriteria'][:3])  # First 3 criteria
            parts.append(f"Inclusion Criteria: {inclusion_text[:300]}")
        
        return "\n".join(parts)

    def _extract_text_from_hf_response(self, result: object) -> str:
        """Extract generated text across HF inference and chat response formats."""
        if isinstance(result, list) and result:
            first = result[0]
            if isinstance(first, dict):
                return first.get("generated_text", "") or str(first)
            return str(first)

        if isinstance(result, dict):
            # Common text-generation response shape
            if result.get("generated_text"):
                return str(result.get("generated_text"))

            # OpenAI-compatible router chat completion shape
            choices = result.get("choices")
            if isinstance(choices, list) and choices:
                first_choice = choices[0] if isinstance(choices[0], dict) else {}
                message = first_choice.get("message", {}) if isinstance(first_choice, dict) else {}
                content = message.get("content") if isinstance(message, dict) else None
                if content:
                    return str(content)
                if first_choice.get("text"):
                    return str(first_choice.get("text"))

            # Surface API-side errors in logs and parsing fallback
            if result.get("error"):
                return str(result.get("error"))

        return str(result)

    def _parse_score_and_reasoning(self, text: str) -> tuple:
        """Parse score and reasoning robustly from model output."""
        response_text = (text or "").strip()

        # Support SCORE: 0.72, score=72%, or score is 0.72
        score_match = re.search(
            r"(?:^|\b)score\s*[:=\-]?\s*(\d+(?:\.\d+)?)\s*%?",
            response_text,
            re.IGNORECASE | re.MULTILINE,
        )
        reasoning_match = re.search(
            r"(?:^|\b)reasoning\s*[:=\-]\s*(.+?)(?=\n\s*(?:score|reasoning)\s*[:=\-]|$)",
            response_text,
            re.IGNORECASE | re.DOTALL,
        )

        score = 0.5
        if score_match:
            score = float(score_match.group(1))
            if score > 1.0:
                score = score / 100.0
            score = max(0.0, min(1.0, score))
        else:
            numbers = re.findall(r"\d+(?:\.\d+)?", response_text)
            if numbers:
                score = float(numbers[0])
                if score > 1.0:
                    score = score / 100.0
                score = max(0.0, min(1.0, score))

        reasoning = "No specific reasoning provided."
        if reasoning_match:
            reasoning = reasoning_match.group(1).strip()
        elif response_text:
            reasoning = response_text[:200]

        # Keep UI text clean when model returns escaped JSON-ish strings.
        reasoning_clean = reasoning.replace("\\", "").replace("/", "")
        return score, reasoning_clean

    def _call_hf_router_api(
        self,
        prompt: str,
        max_tokens: int,
        temperature: float = 0.7,
        top_p: float = 0.9,
        timeout: int = 60,
        allow_chat_fallback: bool = True,
        model_override: Optional[str] = None,
    ):
        """Post to HF router across model and endpoint fallbacks."""
        import requests

        headers = {
            "Authorization": f"Bearer {self.api_token}",
            "Content-Type": "application/json",
        }
        models_to_try = [model_override] if model_override else [self.model_name] + self.fallback_models
        last_response = None
        last_model = self.model_name

        for model_name in models_to_try:
            attempts = [
                (
                    f"https://router.huggingface.co/v1/models/{model_name}/generate",
                    {
                        "inputs": prompt,
                        "parameters": {
                            "max_new_tokens": max_tokens,
                            "temperature": temperature,
                            "top_p": top_p,
                            "return_full_text": False,
                        },
                    },
                ),
                (
                    "https://router.huggingface.co/v1/completions",
                    {
                        "model": model_name,
                        "prompt": prompt,
                        "max_tokens": max_tokens,
                        "temperature": temperature,
                        "top_p": top_p,
                    },
                ),
            ]
            if allow_chat_fallback:
                attempts.append(
                    (
                        "https://router.huggingface.co/v1/chat/completions",
                        {
                            "model": model_name,
                            "messages": [{"role": "user", "content": prompt}],
                            "max_tokens": max_tokens,
                            "temperature": temperature,
                            "top_p": top_p,
                        },
                    )
                )

            for api_url, payload in attempts:
                response = requests.post(api_url, headers=headers, json=payload, timeout=timeout)
                last_response = response
                last_model = model_name
                if response.status_code not in _ROUTER_FALLBACK_STATUSES:
                    return response, model_name

            logger.warning("HF router model incompatible for ranking: %s", model_name)

        return last_response, last_model
    
    def _score_with_local_model(self, prompt: str) -> float:
        """Score using local model (legacy method)"""
        try:
            # Generate response
            outputs = self.pipeline(
                prompt,
                max_new_tokens=10,
                temperature=0.1,
                do_sample=False,
                return_full_text=False
            )
            
            # Extract score from response
            response_text = outputs[0]['generated_text'].strip()
            
            # Try to extract a number from the response
            import re
            numbers = re.findall(r'\d+\.?\d*', response_text)
            if numbers:
                score = float(numbers[0])
                # Normalize if it's > 1 (might be percentage or 0-100 scale)
                if score > 1.0:
                    score = score / 100.0
                return score
            
            return 0.5  # Default score if parsing fails
            
        except Exception as e:
            logger.error(f"Error in local model scoring: {str(e)}")
            return 0.5
    
    def _score_with_reasoning_local(self, prompt: str) -> tuple:
        """Score with reasoning using local model"""
        try:
            # Generate response with more tokens for reasoning - higher temperature for more flexible/creative ranking
            outputs = self.pipeline(
                prompt,
                max_new_tokens=200,
                temperature=0.7,  # Higher temperature for more flexible, less rigid ranking
                do_sample=True,
                top_p=0.9,  # Nucleus sampling for more diverse responses
                return_full_text=False
            )
            
            # Extract response text
            response_text = outputs[0]['generated_text'].strip()
            return self._parse_score_and_reasoning(response_text)
            
        except Exception as e:
            logger.error(f"Error in local model scoring with reasoning: {str(e)}")
            return 0.5, "Local-model fallback used because response generation failed."
    
    def _score_with_api(self, prompt: str) -> float:
        """Score using Hugging Face Inference API (legacy method)"""
        try:
            response, _ = self._call_hf_router_api(
                prompt, max_tokens=10, temperature=0.1, top_p=1.0, timeout=30, allow_chat_fallback=True
            )
            response.raise_for_status()
            generated_text = self._extract_text_from_hf_response(response.json())
            score, _ = self._parse_score_and_reasoning(generated_text)
            return score
        except Exception as e:
            logger.error(f"Error in API scoring: {str(e)}")
            return 0.5
    
    def _score_with_reasoning_api(self, prompt: str) -> tuple:
        """Score with reasoning using Hugging Face Inference API"""
        try:
            if not self.api_token:
                return 0.5, "Hugging Face token missing. Add HUGGINGFACE_API_TOKEN in Space Secrets."

            response, used_model = self._call_hf_router_api(
                prompt, max_tokens=200, timeout=60, allow_chat_fallback=True
            )
            
            if not response.ok:
                error_preview = response.text[:500]
                logger.error(
                    "HF API request failed (status=%s): %s",
                    response.status_code,
                    error_preview,
                )

                if response.status_code in (401, 403):
                    return 0.5, "Hugging Face auth failed (401/403). Check HUGGINGFACE_API_TOKEN permissions."
                if response.status_code == 429:
                    return 0.5, "Hugging Face rate limit reached (429). Please retry shortly."
                if response.status_code in (400, 404, 422):
                    return 0.5, (
                        f"Model/endpoint mismatch ({response.status_code}). "
                        "Set DEEPSEEK_MODEL or HF_FALLBACK_MODELS to router-supported models."
                    )

                return 0.5, f"Hugging Face API error {response.status_code}: {error_preview[:140]}"

            result = response.json()
            generated_text = self._extract_text_from_hf_response(result)
            score, reasoning = self._parse_score_and_reasoning(generated_text)
            if used_model != self.model_name:
                reasoning = f"{reasoning} (scored using fallback model {used_model})"
            return score, reasoning
            
        except Exception as e:
            logger.error(f"Error in API scoring with reasoning: {str(e)}")
            return 0.5, f"API request failed: {str(e)[:140]}"


# Global LLM service instance (lazy loaded)
_llm_service = None

def get_llm_service() -> Optional[LLMService]:
    """Get or create LLM service instance"""
    global _llm_service
    
    if _llm_service is None:
        try:
            use_api = os.environ.get('USE_HF_API', 'false').lower() == 'true'
            api_token = os.environ.get('HUGGINGFACE_API_TOKEN', '')
            
            print(f"Initializing LLM service - USE_HF_API: {use_api}, Token set: {bool(api_token)}")
            
            if use_api and not api_token:
                logger.warning("USE_HF_API is true but HUGGINGFACE_API_TOKEN is not set!")
                print("ERROR: HUGGINGFACE_API_TOKEN environment variable is required for API mode.")
                return None
            
            _llm_service = LLMService(use_api=use_api, api_token=api_token if api_token else None)
            print(f"LLM service initialized successfully. Model: {_llm_service.model_name}")
        except Exception as e:
            logger.error(f"Failed to initialize LLM service: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
    
    return _llm_service