Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import os | |
| from typing import Any | |
| import httpx | |
| from fastapi import HTTPException | |
| from inference.config import get_inference_manifest, resolve_timeout_seconds | |
| from inference.resolve_credentials import assert_api_key | |
| def chat_completions_url() -> str: | |
| manifest = get_inference_manifest() | |
| provider = (manifest.get("providers") or {}).get("proxy-openai") or {} | |
| env_name = str(provider.get("baseUrlEnv") or "OPENAI_BASE_URL") | |
| base = os.environ.get(env_name, "").strip() or str( | |
| provider.get("defaultBaseUrl") or "https://api.openai.com/v1" | |
| ) | |
| return f"{base.rstrip('/')}/chat/completions" | |
| def proxy_model_alias(model: str) -> str: | |
| manifest = get_inference_manifest() | |
| model_entry = (manifest.get("models") or {}).get(model) or {} | |
| aliases = model_entry.get("aliases") or {} | |
| provider_aliases = aliases.get("proxy-openai") or {} | |
| return str(provider_aliases.get("completion") or model) | |
| def build_payload( | |
| *, | |
| system_prompt: str, | |
| user_prompt: str, | |
| max_tokens: int, | |
| temperature: float, | |
| json_schema: dict[str, Any] | None = None, | |
| json_mode: bool = False, | |
| google_search: bool = True, | |
| model_id: str | None = None, | |
| ) -> dict[str, Any]: | |
| payload: dict[str, Any] = { | |
| "messages": [ | |
| {"role": "system", "content": system_prompt}, | |
| { | |
| "role": "user", | |
| "content": "\n\n".join( | |
| [ | |
| "SYSTEM INSTRUCTIONS:", | |
| system_prompt, | |
| "---", | |
| "USER INPUT:", | |
| user_prompt, | |
| ] | |
| ), | |
| }, | |
| ], | |
| "stream": False, | |
| "temperature": temperature, | |
| "max_tokens": max_tokens, | |
| } | |
| if google_search: | |
| payload["tools"] = [{"type": "web_search"}] | |
| return payload | |
| def text_from_response(data: dict[str, Any]) -> str: | |
| try: | |
| text = str(data["choices"][0]["message"].get("content") or "") | |
| except (KeyError, IndexError, TypeError): | |
| text = "" | |
| if not text: | |
| raise HTTPException(status_code=502, detail="TabularAgent returned no text.") | |
| return text | |
| async def post( | |
| payload: dict[str, Any], | |
| *, | |
| model: str, | |
| feature_id: str = "python-review", | |
| client: httpx.AsyncClient | None = None, | |
| ) -> tuple[str, dict[str, Any] | None, dict[str, Any]]: | |
| url = chat_completions_url() | |
| request_payload = {**payload, "model": proxy_model_alias(model)} | |
| timeout = resolve_timeout_seconds(feature_id) | |
| owns_client = client is None | |
| active_client = client or httpx.AsyncClient(timeout=timeout) | |
| try: | |
| response = await active_client.post( | |
| url, | |
| json=request_payload, | |
| headers={ | |
| "Authorization": f"Bearer {assert_api_key('proxy-openai')}", | |
| "Content-Type": "application/json", | |
| }, | |
| ) | |
| finally: | |
| if owns_client: | |
| await active_client.aclose() | |
| if response.status_code >= 400: | |
| detail = response.text[:2000] | |
| raise HTTPException( | |
| status_code=502, detail=f"TabularAgent request failed: {detail}" | |
| ) | |
| data = response.json() | |
| text = text_from_response(data) | |
| usage = data.get("usage") | |
| usage_dict = usage if isinstance(usage, dict) else None | |
| return text, usage_dict, data | |