src/llm.py

import base64
import copy
import json
import os
import re
import struct
from functools import lru_cache
from typing import List, Dict, Callable, Any, Union, Optional
import aioboto3
import aiohttp
import numpy as np
import ollama
import torch
from openai import (
    AsyncOpenAI,
    APIConnectionError,
    RateLimitError,
    Timeout,
    AsyncAzureOpenAI,
)
from pydantic import BaseModel, Field
from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type,
)
from transformers import AutoTokenizer, AutoModelForCausalLM

from .utils import (
    wrap_embedding_func_with_attrs,
    locate_json_string_body_from_string,
    safe_unicode_decode,
    logger,
)

import sys

if sys.version_info < (3, 9):
    from typing import AsyncIterator
else:
    from collections.abc import AsyncIterator

import warnings
import logging

# Suppress flash_attn and torch_dtype warnings from Jina v3
warnings.filterwarnings("ignore", message="flash_attn is not installed")
warnings.filterwarnings("ignore", message="`torch_dtype` is deprecated")
logging.getLogger("transformers_modules").setLevel(logging.ERROR)

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# ============================================================================
# GLOBAL API KEY STORAGE (to avoid async context issues)
# ============================================================================
_GLOBAL_OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")


# ============================================================================
# LLM COMPLETION FUNCTIONS
# ============================================================================

@retry(
    stop=stop_after_attempt(10),
    wait=wait_exponential(multiplier=1, min=4, max=100),
    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
)
async def openai_complete_if_cache(
    model,
    prompt,
    system_prompt=None,
    history_messages=[],
    base_url="https://api.openai.com/v1",
    api_key="",
    **kwargs,
) -> str:
    if api_key:
        os.environ["OPENAI_API_KEY"] = api_key
    
    openai_async_client = (
        AsyncOpenAI() if base_url is None else AsyncOpenAI(base_url=base_url)
    )
    
    kwargs.pop("hashing_kv", None)
    kwargs.pop("keyword_extraction", None)
    
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.extend(history_messages)
    messages.append({"role": "user", "content": prompt})

    logger.debug("===== Query Input to LLM =====")
    logger.debug(f"Query: {prompt}")
    logger.debug(f"System prompt: {system_prompt}")
    
    if "response_format" in kwargs:
        # Use JSON mode for OpenAI models, strip for local models
        if any(prefix in model.lower() for prefix in ["gpt-4", "gpt-5"]):
            kwargs["response_format"] = {"type": "json_object"}
        else:
            kwargs.pop("response_format", None)

    if "gpt-5" in model.lower():
        # GPT-5 uses reasoning tokens that count against max_completion_tokens.
        # Need ~2000+ total (reasoning + output) for reliable responses.
        if "max_tokens" in kwargs:
            max_tokens_value = kwargs.pop("max_tokens")
            kwargs["max_completion_tokens"] = max(max_tokens_value, 2000)
            logger.debug(f"Converted max_tokens to max_completion_tokens for GPT-5")
        elif "max_completion_tokens" not in kwargs:
            kwargs["max_completion_tokens"] = 4000
            logger.debug(f"Set default max_completion_tokens=4000 for GPT-5")
    
    response = await openai_async_client.chat.completions.create(
        model=model, messages=messages, **kwargs
    )

    if hasattr(response, "__aiter__"):
        async def inner():
            async for chunk in response:
                content = chunk.choices[0].delta.content
                if content is None:
                    continue
                if r"\u" in content:
                    content = safe_unicode_decode(content.encode("utf-8"))
                yield content
        return inner()
    else:
        content = response.choices[0].message.content
        if r"\u" in content:
            content = safe_unicode_decode(content.encode("utf-8"))
        return content


class GPTKeywordExtractionFormat(BaseModel):
    high_level_keywords: List[str]
    low_level_keywords: List[str]


async def gpt_4o_mini_complete(
    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
) -> str:
    keyword_extraction = keyword_extraction or kwargs.pop("keyword_extraction", False)
    if keyword_extraction:
        kwargs["response_format"] = GPTKeywordExtractionFormat
    return await openai_complete_if_cache(
        "gpt-4o-mini",
        prompt,
        system_prompt=system_prompt,
        history_messages=history_messages,
        **kwargs,
    )


async def gpt_4o_complete(
    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
) -> str:
    keyword_extraction = keyword_extraction or kwargs.pop("keyword_extraction", False)
    if keyword_extraction:
        kwargs["response_format"] = GPTKeywordExtractionFormat
    return await openai_complete_if_cache(
        "gpt-4o",
        prompt,
        system_prompt=system_prompt,
        history_messages=history_messages,
        **kwargs,
    )


async def gpt_5_complete(
    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
) -> str:
    keyword_extraction = keyword_extraction or kwargs.pop("keyword_extraction", False)
    if keyword_extraction:
        kwargs["response_format"] = GPTKeywordExtractionFormat
    return await openai_complete_if_cache(
        "gpt-5",
        prompt,
        system_prompt=system_prompt,
        history_messages=history_messages,
        **kwargs,
    )


async def gpt_5_mini_complete(
    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
) -> str:
    keyword_extraction = keyword_extraction or kwargs.pop("keyword_extraction", False)
    if keyword_extraction:
        kwargs["response_format"] = GPTKeywordExtractionFormat

    if "max_tokens" in kwargs:
        kwargs["max_completion_tokens"] = kwargs.pop("max_tokens")

    return await openai_complete_if_cache(
        "gpt-5-mini",
        prompt,
        system_prompt=system_prompt,
        history_messages=history_messages,
        **kwargs,
    )


async def gpt_5_nano_complete(
    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
) -> str:
    """GPT-5 Nano - smallest, fastest, most affordable GPT-5 model"""
    keyword_extraction = keyword_extraction or kwargs.pop("keyword_extraction", False)
    if keyword_extraction:
        kwargs["response_format"] = GPTKeywordExtractionFormat

    # GPT-5 models use max_completion_tokens instead of max_tokens
    if "max_tokens" in kwargs:
        kwargs["max_completion_tokens"] = kwargs.pop("max_tokens")
    
    return await openai_complete_if_cache(
        "gpt-5-nano",
        prompt,
        system_prompt=system_prompt,
        history_messages=history_messages,
        **kwargs,
    )


# ============================================================================
# LOCAL / CUSTOM MODEL FUNCTIONS
# ============================================================================
# Configure via environment variables:
#   LOCAL_LLM_BASE_URL  - Base URL for your vLLM / OpenAI-compatible server
#   LOCAL_LLM_API_KEY   - API key (use "dummy" for local servers)
#   LOCAL_LLM_MODEL     - Model name served by the endpoint


def _get_local_llm_config():
    """Get local LLM configuration from environment variables."""
    base_url = os.environ.get("LOCAL_LLM_BASE_URL")
    if not base_url:
        raise ValueError(
            "LOCAL_LLM_BASE_URL not set. "
            "Export it, e.g.: export LOCAL_LLM_BASE_URL='http://localhost:8000/v1'"
        )
    return {
        "base_url": base_url,
        "api_key": os.environ.get("LOCAL_LLM_API_KEY", "dummy"),
        "model": os.environ.get("LOCAL_LLM_MODEL", "llm_base_model"),
    }


async def gpt_oss_120b_complete(
    prompt, system_prompt=None, history_messages=[], **kwargs
) -> str:
    """Complete using a local OpenAI-compatible server (configured via env vars)."""
    kwargs.pop("keyword_extraction", None)
    kwargs.pop("response_format", None)
    kwargs.pop("hashing_kv", None)

    cfg = _get_local_llm_config()
    return await openai_complete_if_cache(
        model=cfg["model"],
        prompt=prompt,
        system_prompt=system_prompt,
        history_messages=history_messages,
        base_url=cfg["base_url"],
        api_key=cfg["api_key"],
        **kwargs,
    )


# ============================================================================
# EMBEDDING FUNCTIONS - LOCAL MODELS
# ============================================================================

@wrap_embedding_func_with_attrs(embedding_dim=1024, max_token_size=8192)
async def local_sentence_embedding(
    texts: list[str],
    model: str = "jinaai/jina-embeddings-v3",
    base_url=None,
    api_key="",
) -> np.ndarray:
    """Jina v3 local embedding (1024-dim, 8192 tokens) - standalone"""
    try:
        import sys
        from pathlib import Path
        
        embedding_models_path = Path(__file__).parent / "embedding_models"
        if str(embedding_models_path) not in sys.path:
            sys.path.insert(0, str(embedding_models_path))
        
        from JinaV3 import JinaV3EmbeddingModel
        from config import MinimalConfig
        
        if not hasattr(local_sentence_embedding, '_model'):
            logger.info("Initializing Jina v3 (local)...")
            
            config = MinimalConfig(
                embedding_model_name=model,
                embedding_batch_size=32,
                embedding_model_dtype="auto",
                embedding_return_as_normalized=True,
                embedding_max_seq_len=8192
            )
            
            local_sentence_embedding._model = JinaV3EmbeddingModel(
                global_config=config,
                embedding_model_name=model
            )
            
            logger.info("✅ Jina v3 ready (1024-dim, standalone)")
        
        embeddings = local_sentence_embedding._model.batch_encode(texts)
        return np.array(embeddings) if not isinstance(embeddings, np.ndarray) else embeddings
        
    except Exception as e:
        logger.error(f"Jina v3 failed ({e})")
        raise


@wrap_embedding_func_with_attrs(embedding_dim=4096, max_token_size=32768)
async def nvidia_nv_embed_v2_embedding(
    texts: list[str],
    model: str = "nvidia/NV-Embed-v2",
    **kwargs
) -> np.ndarray:
    """NVIDIA NV-Embed-v2 local embedding (4096-dim, 32K tokens)"""
    try:
        # Import from local embedding_models directory
        import sys
        from pathlib import Path
        
        embedding_models_path = Path(__file__).parent / "embedding_models"
        if str(embedding_models_path) not in sys.path:
            sys.path.insert(0, str(embedding_models_path))
        
        from NVEmbedV2 import NVEmbedV2EmbeddingModel
        from config import MinimalConfig
        
        if not hasattr(nvidia_nv_embed_v2_embedding, '_model'):
            logger.info("Initializing NVIDIA NV-Embed-v2 (local)...")
            
            config = MinimalConfig(
                embedding_model_name=model,
                embedding_batch_size=16,
                embedding_model_dtype="auto",
                embedding_return_as_normalized=True,
                embedding_max_seq_len=32768
            )
            
            nvidia_nv_embed_v2_embedding._model = NVEmbedV2EmbeddingModel(
                global_config=config,
                embedding_model_name=model
            )
            
            logger.info("✅ NVIDIA NV-Embed-v2 ready (4096-dim, standalone)")
        
        embeddings = nvidia_nv_embed_v2_embedding._model.batch_encode(texts)
        return np.array(embeddings) if not isinstance(embeddings, np.ndarray) else embeddings
        
    except Exception as e:
        logger.warning(f"NVIDIA failed ({e}), using Jina v3")
        return await local_sentence_embedding(texts)



@wrap_embedding_func_with_attrs(embedding_dim=4096, max_token_size=8192)
async def gritlm_embedding(
    texts: list[str],
    model: str = "GritLM/GritLM-7B",
    **kwargs
) -> np.ndarray:
    """GritLM local embedding (7168-dim, 8192 tokens) - standalone"""
    try:
        import sys
        from pathlib import Path
        
        embedding_models_path = Path(__file__).parent / "embedding_models"
        if str(embedding_models_path) not in sys.path:
            sys.path.insert(0, str(embedding_models_path))
        
        from GritLM import GritLMEmbeddingModel
        from config import MinimalConfig
        
        if not hasattr(gritlm_embedding, '_model'):
            logger.info("Initializing GritLM (local)...")
            
            config = MinimalConfig(
                embedding_model_name=model,
                embedding_batch_size=16,
                embedding_model_dtype="auto",
                embedding_return_as_normalized=True,
                embedding_max_seq_len=8192
            )
            
            gritlm_embedding._model = GritLMEmbeddingModel(
                global_config=config,
                embedding_model_name=model
            )
            
            logger.info("✅ GritLM ready (7168-dim, standalone)")
        
        embeddings = gritlm_embedding._model.batch_encode(texts, **kwargs)
        return np.array(embeddings) if not isinstance(embeddings, np.ndarray) else embeddings
        
    except Exception as e:
        logger.warning(f"GritLM failed ({e}), using Jina v3")
        return await local_sentence_embedding(texts)


# ============================================================================
# EMBEDDING FUNCTIONS - CLOUD/API MODELS
# ============================================================================

@wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8192)
@retry(
    stop=stop_after_attempt(10),
    wait=wait_exponential(multiplier=1, min=4, max=100),
    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
)
async def openai_small_embedding(
    texts: list[str],
    model: str = "text-embedding-3-small",
    dimensions: int = 1536,
    base_url="https://api.openai.com/v1",
    api_key="",
) -> np.ndarray:
    """OpenAI text-embedding-3-small - standalone (1536-dim)"""
    
    # ✅ FIX: Use global API key captured at module import time
    global _GLOBAL_OPENAI_API_KEY
    if not api_key or api_key == "":
        api_key = _GLOBAL_OPENAI_API_KEY or os.environ.get("OPENAI_API_KEY", "")
    
    try:
        import sys
        from pathlib import Path
        
        embedding_models_path = Path(__file__).parent / "embedding_models"
        if str(embedding_models_path) not in sys.path:
            sys.path.insert(0, str(embedding_models_path))
        
        from OpenAI import OpenAIEmbeddingModel
        from config import MinimalConfig
        
        if not hasattr(openai_small_embedding, '_model') or openai_small_embedding._last_dims != dimensions:
            logger.info(f"Initializing OpenAI Small (standalone): {model} with {dimensions}-dim")
            
            config = MinimalConfig(
                embedding_model_name=model,
                embedding_batch_size=100,
                embedding_dimensions=dimensions,
                # openai_api_key=api_key  # Pass the API key here
            )
            
            openai_small_embedding._model = OpenAIEmbeddingModel(
                global_config=config,
                embedding_model_name=model,
                api_key=api_key
            )
            openai_small_embedding._last_dims = dimensions
            
            logger.info(f"✅ OpenAI Small ready ({dimensions}-dim, standalone)")
        
        embeddings = openai_small_embedding._model.batch_encode(texts)
        return np.array(embeddings) if not isinstance(embeddings, np.ndarray) else embeddings
        
    except Exception as e:
        logger.warning(f"OpenAI Small standalone failed ({e}), using direct API call")
        
        # API key already checked at the start
        if base_url is None:
            openai_async_client = AsyncOpenAI(api_key=api_key)
        else:
            openai_async_client = AsyncOpenAI(base_url=base_url, api_key=api_key)
        
        response = await openai_async_client.embeddings.create(
            model=model, 
            input=texts, 
            encoding_format="float",
            dimensions=dimensions
        )
        return np.array([dp.embedding for dp in response.data])


@wrap_embedding_func_with_attrs(embedding_dim=3072, max_token_size=8192)
@retry(
    stop=stop_after_attempt(10),
    wait=wait_exponential(multiplier=1, min=4, max=100),
    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
)
async def openai_large_embedding(
    texts: list[str],
    model: str = "text-embedding-3-large",
    dimensions: int = 3072,
    base_url="https://api.openai.com/v1",
    api_key="",
) -> np.ndarray:
    """OpenAI text-embedding-3-large - standalone (3072-dim)"""
    
    # ✅ FIX: Use global API key captured at module import time
    global _GLOBAL_OPENAI_API_KEY
    if not api_key or api_key == "":
        api_key = _GLOBAL_OPENAI_API_KEY or os.environ.get("OPENAI_API_KEY", "")
    
    try:
        import sys
        from pathlib import Path
        
        embedding_models_path = Path(__file__).parent / "embedding_models"
        if str(embedding_models_path) not in sys.path:
            sys.path.insert(0, str(embedding_models_path))
        
        from OpenAI import OpenAIEmbeddingModel
        from config import MinimalConfig
        
        if not hasattr(openai_large_embedding, '_model') or openai_large_embedding._last_dims != dimensions:
            logger.info(f"Initializing OpenAI Large (standalone): {model} with {dimensions}-dim")
            
            # ✅ FIX: Pass API key to config
            config = MinimalConfig(
                embedding_model_name=model,
                embedding_batch_size=100,
                embedding_dimensions=dimensions,
                # openai_api_key=api_key  # Pass the API key here
            )
            
            openai_large_embedding._model = OpenAIEmbeddingModel(
                global_config=config,
                embedding_model_name=model,
                api_key=api_key
            )
            openai_large_embedding._last_dims = dimensions
            
            logger.info(f"✅ OpenAI Large ready ({dimensions}-dim, standalone)")
        
        embeddings = openai_large_embedding._model.batch_encode(texts)
        return np.array(embeddings) if not isinstance(embeddings, np.ndarray) else embeddings
        
    except Exception as e:
        logger.warning(f"OpenAI Large standalone failed ({e}), using direct API call")
        
        # API key already checked at the start
        if base_url is None:
            openai_async_client = AsyncOpenAI(api_key=api_key)
        else:
            openai_async_client = AsyncOpenAI(base_url=base_url, api_key=api_key)
        
        response = await openai_async_client.embeddings.create(
            model=model, 
            input=texts, 
            encoding_format="float",
            dimensions=dimensions
        )
        return np.array([dp.embedding for dp in response.data])


# Backward compatibility alias (defaults to large)
async def openai_cloud_embedding(*args, **kwargs):
    """Alias for backward compatibility - routes to openai_large_embedding"""
    return await openai_large_embedding(*args, **kwargs)


@wrap_embedding_func_with_attrs(embedding_dim=1024, max_token_size=8192)
async def openai_embedding(
    texts: list[str],
    model: str = "jinaai/jina-embeddings-v3",
    base_url=None,
    api_key="",
) -> np.ndarray:
    """
    Smart embedding function - ALWAYS 1024 dimensions.
    
    Uses LOCAL by default (Jina v3, 1024-dim, free).
    Set USE_OPENAI_EMBEDDINGS=1 to use OpenAI (1024-dim, costs money).
    """
    use_openai = os.environ.get("USE_OPENAI_EMBEDDINGS", "0") == "1"
    
    if use_openai:
        logger.info("Using OpenAI cloud embeddings (1024-dim)")
        return await openai_large_embedding(texts, model="text-embedding-3-large", dimensions=1024, base_url=base_url, api_key=api_key)
    else:
        return await local_sentence_embedding(texts, model=model, base_url=base_url, api_key=api_key)


# ============================================================================
# EMBEDDING MODEL REGISTRY
# ============================================================================

EMBEDDING_CONFIGS = {
    # ===== LOCAL EMBEDDINGS (FREE) =====
    "jina-v3": {
        "provider": "local",
        "model": "jinaai/jina-embeddings-v3",
        "dimensions": 1024,
        "max_tokens": 8192,
        "description": "Jina Embeddings v3 (1024-dim, 8K context, free, local)",
        "embedding_func": local_sentence_embedding,
        "cost": "free",
    },
    "nvidia-nv-embed-v2": {
        "provider": "local",
        "model": "nvidia/NV-Embed-v2",
        "dimensions": 4096,
        "max_tokens": 32768,
        "description": "NVIDIA NV-Embed-v2 (4096-dim, 32K context, local, high quality)",
        "embedding_func": nvidia_nv_embed_v2_embedding,
        "cost": "free",
    },
    "gritlm": {
        "provider": "local",
        "model": "GritLM/GritLM-7B",
        "dimensions": 4096,
        "max_tokens": 8192,
        "description": "GritLM (4096-dim, 8K context, local, unified embedding+generation)",
        "embedding_func": gritlm_embedding,
        "cost": "free",
    },
    
    # ===== CLOUD EMBEDDINGS (PAID) =====
    "openai-large": {
        "provider": "openai",
        "model": "text-embedding-3-large",
        "dimensions": 3072,
        "max_tokens": 8192,
        "description": "OpenAI text-embedding-3-large (3072-dim, cloud, costs money)",
        "embedding_func": openai_large_embedding,
        "cost": "$0.13/1M tokens",
    },
    "openai-small": {
        "provider": "openai",
        "model": "text-embedding-3-small",
        "dimensions": 1536,
        "max_tokens": 8192,
        "description": "OpenAI text-embedding-3-small (1536-dim, cloud, costs money)",
        "embedding_func": openai_small_embedding,
        "cost": "$0.02/1M tokens",
    },
}


def get_embedding_func_for_model(embedding_key: str):
    """Get embedding function and config for specified embedding model"""
    if embedding_key not in EMBEDDING_CONFIGS:
        logger.warning(f"Unknown embedding model '{embedding_key}', defaulting to 'jina-v3'")
        embedding_key = "jina-v3"
    
    config = EMBEDDING_CONFIGS[embedding_key]
    logger.info(f"[Embedding] {embedding_key} → {config['description']}")
    
    return config["embedding_func"], config["dimensions"], config


def list_available_embeddings():
    """Print all available embedding models"""
    print("\n" + "=" * 80)
    print("📊 Available Embedding Models")
    print("=" * 80)
    
    print("\n🆓 FREE LOCAL MODELS:")
    for key, config in EMBEDDING_CONFIGS.items():
        if config["cost"] == "free":
            print(f"\n  {key:25s} → {config['description']}")
            print(f"  {'':25s}   Dimensions: {config['dimensions']}, Max tokens: {config['max_tokens']}")
    
    print("\n💰 PAID CLOUD MODELS:")
    for key, config in EMBEDDING_CONFIGS.items():
        if config["cost"] != "free":
            print(f"\n  {key:25s} → {config['description']}")
            print(f"  {'':25s}   Dimensions: {config['dimensions']}, Cost: {config['cost']}")
    
    print("\n" + "=" * 80 + "\n")


# ============================================================================
# MODEL TYPE DETECTION HELPERS
# ============================================================================

LOCAL_MODEL_PREFIXES = [
    "gpt-oss", "qwen", "llama", "mistral", "internlm", "ollama", "local"
]
CLOUD_MODEL_PREFIXES = [
    "gpt-4", "gpt-5", "nvidia", "azure", "bedrock", "zhipu", "openai"
]

def is_local_model(model_name: str) -> bool:
    """Return True if model_name corresponds to a locally hosted model."""
    if not model_name:
        return False
    model_name = model_name.lower()
    local = any(model_name.startswith(prefix) for prefix in LOCAL_MODEL_PREFIXES)
    if not local and not any(model_name.startswith(prefix) for prefix in CLOUD_MODEL_PREFIXES):
        logger.warning(f"[Model Detection] Unknown model prefix for '{model_name}'. Defaulting to cloud mode.")
    return local

def is_cloud_model(model_name: str) -> bool:
    """Return True if model_name corresponds to a cloud-hosted model."""
    if not model_name:
        return False
    model_name = model_name.lower()
    return any(model_name.startswith(prefix) for prefix in CLOUD_MODEL_PREFIXES)


# =====================================================================
# STARTUP SUMMARY (auto-prints when QAFD.llm is imported)
# =====================================================================

# Embedding mode (for reference, no startup print)
_embedding_mode_env = os.environ.get("EMBEDDING_MODEL_KEY", "auto")


if __name__ == "__main__":
    import asyncio

    async def main():
        # Test listing
        list_available_embeddings()
        
        # Test embedding function
        result = await gpt_4o_mini_complete("How are you?")
        print(result)

    asyncio.run(main())