# 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, )