src/analysis/llm.py

import subprocess
import re
import html
import os
import logging
import time
import torch
from typing import Optional, List, Dict
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from transformers.utils import logging as tf_logging
from huggingface_hub import InferenceClient
tf_logging.set_verbosity_error()


            
class HFGenerator:
    def __init__(self, model_name: str = "meta-llama/Llama-3.2-3B-Instruct", use_api: bool = False):
        self.use_api = use_api
        self.model_name = model_name
        self.use_cuda = torch.cuda.is_available()
        self.device = torch.device("cuda" if self.use_cuda else "cpu")
        
        # 1. Token Retrieval
        self.hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")

        if self.use_api:
            print(f"[LLM] Mode: API Inference ({model_name})")
            self.client = InferenceClient(model=model_name, token=self.hf_token)
        else:
            print(f"[LLM] Mode: Local Load | Device: {self.device}")

            # 2. 4-Bit Quantization (Only for GPU)
            if self.use_cuda:
                bnb_config = BitsAndBytesConfig(
                    load_in_4bit=True,
                    bnb_4bit_quant_type="nf4",
                    bnb_4bit_compute_dtype=torch.float16,
                    bnb_4bit_use_double_quant=True,
                )
                quantization_config = bnb_config
            else:
                quantization_config = None

            # 3. Tokenizer
            self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=self.hf_token)
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token

            # 4. Model Loading Logic
            model_kwargs = {
                "token": self.hf_token,
                "trust_remote_code": True,
                "device_map": "auto" if self.use_cuda else "cpu", 
            }
            if quantization_config:
                model_kwargs["quantization_config"] = quantization_config

            print(f"[LLM] Loading {model_name} {'in 4-bit ' if quantization_config else ''}on {self.device}...")
            self.model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)

            self.max_pos = getattr(self.model.config, "max_position_embeddings", 2048)

    @torch.no_grad()
    def generate(self, prompt: str, deterministic: bool = False, max_new_tokens: int = 512, temperature: float = 0.7) -> str:
        if self.use_api:
            return self.client.text_generation(
                prompt, max_new_tokens=max_new_tokens, temperature=temperature, stop_sequences=["</s>"]
            )
        
        # Local Inference Path
        inputs = self.tokenizer(
            prompt, 
            return_tensors="pt", 
            truncation=True, 
            max_length=self.max_pos - max_new_tokens
        ).to(self.device)

        # Use CPU autocast (bfloat16) for a 2x speedup on compatible processors
        autocast_dtype = torch.float16 if self.use_cuda else torch.bfloat16
        device_type = "cuda" if self.use_cuda else "cpu"

        with torch.autocast(device_type=device_type, dtype=autocast_dtype):
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=not deterministic,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
            )
        
        decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return decoded[len(prompt):].strip()
    

def generate_with_ollama(model_name: str, prompt: str) -> str:
    try:
        proc = subprocess.run(
            ["ollama", "run", model_name],
            input=prompt,
            capture_output=True,
            text=True,
            check=True,
        )
        return proc.stdout.strip()
    except Exception as e:
        return f"Ollama error: {e}"


_CLIENT_INSTANCE: Optional[InferenceClient] = None

def get_client() -> InferenceClient:
    global _CLIENT_INSTANCE
    if _CLIENT_INSTANCE is None:
        token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
        # In 2026, simply initializing the client will pick up local credentials 
        # but passing it explicitly is safer for Spaces.
        _CLIENT_INSTANCE = InferenceClient(token=token)
    return _CLIENT_INSTANCE


def generate_with_hf_api(
    model_name: str, 
    prompt: str, 
    system_message: str = "You are a thoughtful assistant.",
    max_tokens: int = 1024,
    temperature: float = 0.7,
    provider: Optional[str] = None
) -> str:
    """
    HF API implementation with detailed logging for RAG tracking.
    """
    client = get_client()
    full_model_id = f"{model_name}:{provider}" if provider else model_name
    
    # Log the request details
    logging.info(f"🚀 HF API Request | Model: {full_model_id} | Max Tokens: {max_tokens}")
    
    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": prompt}
    ]

    start_time = time.time()
    
    try:
        response = client.chat.completions.create(
            model=full_model_id,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            stream=False 
        )
        
        duration = time.time() - start_time
        content = response.choices[0].message.content.strip()
        
        # Log successful response metrics
        logging.info(f"✅ HF API Success | Latency: {duration:.2f}s | Response Len: {len(content)} chars")
        
        return content

    except Exception as e:
        duration = time.time() - start_time
        logging.error(f"❌ HF API Error after {duration:.2f}s: {str(e)}")
        
        # Handle specific error types if needed (e.g., rate limits)
        if "429" in str(e):
            return "HF API error: Rate limit exceeded. Please wait a moment."
            
        return f"HF API error: {str(e)}"
    


def generate_answer(
    transcription_text: str,
    backend: str = "hf",
    model_name: str = "meta-llama/Llama-3.2-3B-Instruct",
    use_hf_api: bool = True,
    max_new_tokens: int = 1024,
) -> str:
    """
    RAG Answer Generation with multi-backend support and performance logging.
    """
    start_time = time.time()
    
    # 2. Build the prompt
    meeting_prompt = f"""
        ### ROLE: COGNITIVE EXPERT & PROJECT MANAGER
        Act as an expert Cognitive Editor and Business Analyst. Your goal is to transform messy, error-prone speech-to-text into a polished, high-value document.

        ### TASK 1: CLEANING
        - Remove ASR hallucinations (e.g., repeating words like "you you you" or "thank you" during silence).
        - Correct homophone errors (e.g., "there" vs "their", "cash" vs "cache").
        - Remove filler words (ums, ahs, "you know") while preserving the speaker's original meaning.

        ### TASK 2: CATEGORIZATION & FORMATTING
        Analyze the cleaned text. Determine the "Content Type" and format the output accordingly:

        1. **If it's a MEETING:**
        - Provide "Executive Summary," "Key Decisions," and "Action Items" (Task | Owner | Deadline).
        2. **If it's a BRAINSTORM/IDEA JUGGLING:**
        - Organize disorganized thoughts into a "Logical Framework."
        - Identify the "Core Concept" and provide "Strategic Guidelines" for the next steps.
        3. **If it's a LESSON/LECTURE:**
        - Create a "Conceptual Map."
        - Summarize into "Key Takeaways" and "Definitions" of complex terms used.

        ### TRANSCRIPTION DATA:
        \"\"\"
        {transcription_text}
        \"\"\"

        ### OUTPUT:
        (Start with a 1-sentence "Intent Identification" e.g., "This appears to be a brainstorming session regarding...")
        """
    logging.info(f"Built prompt with transcription: {transcription_text[:500]}...")  # Log first 500 chars

    try:
        # Route to HF API
        if use_hf_api:
            logging.info(f"🌐 Routing to Hugging Face API (Model: {model_name})")
            response = generate_with_hf_api(model_name, meeting_prompt, max_tokens=max_new_tokens)
        
        # Route to Local Backends
        elif backend == "ollama":
            logging.info(f"🦙 Routing to Ollama (Model: {model_name})")
            response = generate_with_ollama(model_name, meeting_prompt)
            
        elif backend == "hf":
            logging.info(f"🏗️ Routing to Local HF Transformers (Model: {model_name})")
            generator = HFGenerator(model_name=model_name)
            response = generator.generate(
                meeting_prompt,
                max_new_tokens=max_new_tokens,
                deterministic=True,
            ) 
        else:
            response = f"Unsupported backend: {backend}"
            logging.warning(f"⚠️ {response}")

    except Exception as e:
        logging.error(f"❌ Generation Failed: {str(e)}")
        response = f"Error during generation: {str(e)}"

    # 3. Log Performance Metrics
    duration = time.time() - start_time
    logging.info(f"⏱️ Generation Complete | Time: {duration:.2f}s | Response Length: {len(response)} chars")

    return response