llm_client.py

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Unified Hugging Face LLM client for Sushruta Patient 360."""

import json
import logging
import requests
import hashlib
from typing import Any, Generator, List, Union
from config import HF_TOKEN, MEDGEMMA_27B_ENDPOINT, MEDGEMMA_4B_ENDPOINT
from cache_manager import intake_cache, radiology_cache

logger = logging.getLogger(__name__)


class HFModelClient:
    """OpenAI-compatible client for Hugging Face Inference API / Endpoints."""

    def __init__(
        self,
        endpoint_url: str,
        hf_token: str,
        model_name: str = "tgi",
        default_router_model: str = None,
        cache_instance: Any = None,
    ):
        """Initialize the client.

        Args:
            endpoint_url: Custom HF Endpoint URL, or empty string to use router.
            hf_token: Hugging Face API token.
            model_name: Model identifier for API calls.
            default_router_model: Model name to use when falling back to HF serverless router.
            cache_instance: diskcache Cache instance for memoization.
        """
        self.endpoint_url = endpoint_url.strip()
        self.hf_token = hf_token.strip()
        self.model_name = model_name
        self.default_router_model = default_router_model
        self.cache = cache_instance

        # Resolve the active API URL
        if self.endpoint_url:
            # If endpoint is custom, ensure it goes to chat/completions
            if "/v1/chat/completions" in self.endpoint_url:
                self.api_url = self.endpoint_url
            else:
                self.api_url = f"{self.endpoint_url.rstrip('/')}/v1/chat/completions"
            logger.info("Using custom endpoint for model %s: %s", model_name, self.api_url)
        else:
            # Fallback to Hugging Face Router API (unified providers)
            model_to_use = self.default_router_model or "google/gemma-3-27b-it"
            self.api_url = "https://router.huggingface.co/v1/chat/completions"
            self.model_name = model_to_use
            logger.info(
                "Using HF Serverless Router for model %s: %s",
                model_to_use,
                self.api_url,
            )

    def chat_completion(
        self,
        messages: List[dict],
        temperature: float = 0.1,
        max_tokens: int = 2048,
        stream: bool = False,
        **kwargs,
    ) -> Union[str, Generator[str, None, None]]:
        """Call the HuggingFace chat completions API with optional cache lookup.

        Supports streaming and non-streaming responses.
        """
        # Build payload
        payload = {
            "model": self.model_name,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "stream": stream,
            **kwargs,
        }

        # Calculate cache key (for non-streaming, or streaming caching)
        # We serialize messages and parameters to construct a unique key
        key_str = json.dumps(payload, sort_keys=True)
        key_hash = hashlib.sha256(key_str.encode("utf-8")).hexdigest()

        if self.cache is not None:
            cached_val = self.cache.get(key_hash)
            if cached_val is not None:
                logger.info("Cache hit for model %s (hash: %s)", self.model_name, key_hash)
                if stream:
                    # Generator yielding the cached value as a single chunk
                    def cached_stream():
                        yield cached_val
                    return cached_stream()
                else:
                    return cached_val

        headers = {
            "Content-Type": "application/json",
        }
        if self.hf_token:
            headers["Authorization"] = f"Bearer {self.hf_token}"

        try:
            logger.info(
                "Calling HF completion API at %s for model %s (stream=%s)",
                self.api_url,
                self.model_name,
                stream,
            )
            response = requests.post(
                self.api_url, headers=headers, json=payload, stream=stream, timeout=90
            )
            response.raise_for_status()
        except Exception as e:
            logger.error("Error connecting to HF Inference API: %s", e)
            # Try to print response body for debugging if it failed
            try:
                if 'response' in locals() and not stream:
                    logger.error("Response details: %s", response.text)
            except Exception:
                pass
            raise e

        if stream:
            # Generator that yields chunks and collects them to cache the final output
            def stream_generator():
                collected_chunks = []
                for line in response.iter_lines():
                    if not line:
                        continue
                    line_str = line.decode("utf-8").strip()
                    if line_str.startswith("data: "):
                        data_content = line_str[6:]
                        if data_content == "[DONE]":
                            break
                        try:
                            chunk_json = json.loads(data_content)
                            choice = chunk_json.get("choices", [{}])[0]
                            delta = choice.get("delta", {})
                            content = delta.get("content", "")
                            if content:
                                collected_chunks.append(content)
                                yield content
                        except Exception as e:
                            logger.debug("Failed parsing SSE chunk: %s (chunk: %s)", e, line_str)
                
                # After completing stream, cache the full result if caching enabled
                full_text = "".join(collected_chunks)
                if self.cache is not None and full_text:
                    try:
                        self.cache.set(key_hash, full_text)
                    except Exception as e:
                        logger.warning("Failed to write stream to cache: %s", e)

            return stream_generator()
        else:
            response_json = response.json()
            full_text = response_json["choices"][0]["message"]["content"]
            
            # Cache the response
            if self.cache is not None and full_text:
                try:
                    self.cache.set(key_hash, full_text)
                except Exception as e:
                    logger.warning("Failed to write to cache: %s", e)
                    
            return full_text


# Create global instances
# MedGemma 27B for Clinical Assistant (intake_cache)
medgemma_27b = HFModelClient(
    endpoint_url=MEDGEMMA_27B_ENDPOINT,
    hf_token=HF_TOKEN,
    model_name="medgemma-27b",
    default_router_model="Qwen/Qwen2.5-72B-Instruct",
    cache_instance=intake_cache,
)

# MedGemma 4B for Radiology Explainer (radiology_cache)
medgemma_4b = HFModelClient(
    endpoint_url=MEDGEMMA_4B_ENDPOINT,
    hf_token=HF_TOKEN,
    model_name="medgemma-4b",
    default_router_model="Qwen/Qwen2.5-7B-Instruct",
    cache_instance=radiology_cache,
)

# Gemma-3-27B-IT for general roleplay (intake_cache)
gemma_roleplay = HFModelClient(
    endpoint_url="",  # Always use the serverless router
    hf_token=HF_TOKEN,
    model_name="google/gemma-3-27b-it",
    default_router_model="google/gemma-3-27b-it",
    cache_instance=intake_cache,
)