Spaces:
Running
Running
| from __future__ import annotations | |
| import json | |
| import logging | |
| import re | |
| from typing import Any | |
| from .budget import BudgetLedger | |
| from .pricing import TokenRates, cost_usd | |
| from .secrets import get_cloud_credentials | |
| logger = logging.getLogger(__name__) | |
| _JSON_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)```", re.S | re.I) | |
| def _extract_json_object(text: str) -> str: | |
| text = (text or "").strip() | |
| if not text: | |
| return text | |
| fence = _JSON_FENCE_RE.search(text) | |
| if fence: | |
| text = fence.group(1).strip() | |
| start = text.find("{") | |
| end = text.rfind("}") | |
| if start >= 0 and end > start: | |
| return text[start : end + 1] | |
| return text | |
| class CloudLLMClient: | |
| def __init__( | |
| self, | |
| *, | |
| budget: BudgetLedger, | |
| rates: TokenRates, | |
| text_model: str = "qwen-plus", | |
| vlm_model: str = "qwen2-omni-flash", | |
| fast_model: str | None = None, | |
| timeout: float = 60.0, | |
| ) -> None: | |
| self.budget = budget | |
| self.rates = rates | |
| self.text_model = text_model | |
| # cheap/low-latency model for scope checks and query rewriting | |
| self.fast_model = fast_model or text_model | |
| self.vlm_model = vlm_model | |
| self.embed_model = "text-embedding-v3" | |
| # Large diagnosis prompts (retrieved context + 3 full causes) routinely take | |
| # >20s to generate, so default to a roomier timeout. | |
| self.timeout = float(timeout) | |
| api_key, base_url = get_cloud_credentials() | |
| try: | |
| from openai import OpenAI | |
| except Exception as e: | |
| raise RuntimeError("pip install openai") from e | |
| self._client = OpenAI(api_key=api_key, base_url=base_url, timeout=self.timeout, max_retries=0) | |
| def _usage_cost( | |
| self, | |
| usage: Any, | |
| *, | |
| assumed_vision_tokens: int = 0, | |
| multimodal_output: bool = True, | |
| ) -> float: | |
| if usage is None: | |
| return 0.0 | |
| prompt = int(getattr(usage, "prompt_tokens", 0) or 0) | |
| completion = int(getattr(usage, "completion_tokens", 0) or 0) | |
| vision = min(prompt, assumed_vision_tokens) if assumed_vision_tokens else 0 | |
| text_in = max(0, prompt - vision) | |
| return cost_usd( | |
| rates=self.rates, | |
| vision_tokens=vision, | |
| text_in_tokens=text_in, | |
| text_out_tokens=completion, | |
| multimodal_output=multimodal_output, | |
| ) | |
| def chat_text_json( | |
| self, | |
| *, | |
| messages: list[dict], | |
| category: str = "diagnosis", | |
| estimated_usd: float = 0.02, | |
| temperature: float = 0.15, | |
| model: str | None = None, | |
| ) -> str: | |
| self.budget.require(estimated_usd, category=category) | |
| use_model = model or self.text_model | |
| try: | |
| resp = self._client.chat.completions.create( | |
| model=use_model, | |
| messages=messages, | |
| temperature=temperature, | |
| response_format={"type": "json_object"}, | |
| timeout=self.timeout, | |
| ) | |
| except Exception: | |
| logger.warning("JSON-mode text chat failed; retrying without response_format", exc_info=True) | |
| resp = self._client.chat.completions.create( | |
| model=use_model, | |
| messages=messages, | |
| temperature=temperature, | |
| timeout=self.timeout, | |
| ) | |
| content = resp.choices[0].message.content or "" | |
| actual = self._usage_cost(resp.usage, multimodal_output=False) | |
| if actual <= 0: | |
| actual = estimated_usd | |
| self.budget.charge(actual, category, note=f"text chat ({use_model})") | |
| self.last_usage = { | |
| "model": use_model, | |
| "prompt_tokens": int(getattr(resp.usage, "prompt_tokens", 0) or 0) if resp.usage else 0, | |
| "completion_tokens": int(getattr(resp.usage, "completion_tokens", 0) or 0) if resp.usage else 0, | |
| "usd": round(actual, 6), | |
| } | |
| return _extract_json_object(content) | |
| def chat_vision_json( | |
| self, | |
| *, | |
| prompt: str, | |
| image_b64_url: str, | |
| category: str = "video_vlm", | |
| estimated_usd: float = 0.01, | |
| vision_tokens: int = 0, | |
| temperature: float = 0.2, | |
| ) -> str: | |
| self.budget.require(estimated_usd, category=category) | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "text", "text": prompt}, | |
| {"type": "image_url", "image_url": {"url": image_b64_url}}, | |
| ], | |
| } | |
| ] | |
| resp = self._client.chat.completions.create( | |
| model=self.vlm_model, | |
| messages=messages, | |
| temperature=temperature, | |
| response_format={"type": "json_object"}, | |
| ) | |
| content = resp.choices[0].message.content or "" | |
| actual = self._usage_cost( | |
| resp.usage, | |
| assumed_vision_tokens=vision_tokens, | |
| multimodal_output=True, | |
| ) | |
| if actual <= 0: | |
| actual = estimated_usd | |
| self.budget.charge(actual, category, note="vision chat") | |
| return content | |
| def embed_texts( | |
| self, | |
| texts: list[str], | |
| *, | |
| model: str | None = None, | |
| category: str = "embedding", | |
| batch_size: int = 96, | |
| ) -> list[list[float]]: | |
| """Batched embeddings (one HTTP call per ``batch_size`` texts).""" | |
| out: list[list[float]] = [] | |
| for i in range(0, len(texts), batch_size): | |
| chunk = [t[:8000] for t in texts[i : i + batch_size]] | |
| self.budget.require(0.0001 * len(chunk), category=category) | |
| resp = self._client.embeddings.create(model=model or self.embed_model, input=chunk) | |
| out.extend(d.embedding for d in resp.data) | |
| usage = getattr(resp, "usage", None) | |
| tokens = int(getattr(usage, "total_tokens", 0) or 0) if usage else sum(max(50, len(t) // 4) for t in chunk) | |
| actual = cost_usd(rates=self.rates, text_in_tokens=tokens) | |
| if actual > 0: | |
| self.budget.charge(actual, category, note=f"embedding x{len(chunk)}") | |
| return out | |
| def embed_text( | |
| self, | |
| text: str, | |
| *, | |
| model: str | None = None, | |
| category: str = "embedding", | |
| estimated_usd: float = 0.0001, | |
| ) -> list[float]: | |
| self.budget.require(estimated_usd, category=category) | |
| resp = self._client.embeddings.create( | |
| model=model or self.embed_model, | |
| input=text[:8000], | |
| ) | |
| vec = resp.data[0].embedding | |
| usage = getattr(resp, "usage", None) | |
| tokens = int(getattr(usage, "total_tokens", 0) or 0) if usage else max(50, len(text) // 4) | |
| actual = cost_usd(rates=self.rates, text_in_tokens=tokens) | |
| if actual <= 0: | |
| actual = estimated_usd | |
| self.budget.charge(actual, category, note="embedding") | |
| return vec | |