| |
| |
| """Small OpenAI-compatible chat helper used by query data scripts.""" |
|
|
| from __future__ import annotations |
|
|
| import http.client |
| import json |
| import os |
| import random |
| import time |
| from typing import Any, Dict, List |
| from urllib.parse import urlparse |
|
|
|
|
| DEFAULT_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1" |
|
|
|
|
| def chat_completion( |
| messages: List[Dict[str, str]], |
| model: str, |
| base_url: str = DEFAULT_BASE_URL, |
| api_key_env: str = "DASHSCOPE_API_KEY", |
| temperature: float = 0.2, |
| enable_thinking: bool = True, |
| stream: bool = True, |
| timeout: int = 180, |
| ) -> str: |
| api_key = os.environ.get(api_key_env) or os.environ.get("OPENAI_API_KEY") |
| if not api_key: |
| raise RuntimeError(f"Missing API key. Set ${api_key_env} or $OPENAI_API_KEY.") |
|
|
| max_retries = max(1, int(os.environ.get("LLM_API_MAX_RETRIES", "4"))) |
| base_sleep = max(0.0, float(os.environ.get("LLM_API_RETRY_BASE_SEC", "2.0"))) |
| stream_fallback = os.environ.get("LLM_API_STREAM_FALLBACK", "1").lower() not in {"0", "false", "no"} |
| force_non_stream = os.environ.get("LLM_API_DISABLE_STREAM", "0").lower() in {"1", "true", "yes"} |
| requested_stream = bool(stream and not force_non_stream) |
|
|
| errors: list[str] = [] |
| for attempt in range(1, max_retries + 1): |
| try: |
| return _chat_completion_once( |
| messages=messages, |
| model=model, |
| base_url=base_url, |
| api_key=api_key, |
| temperature=temperature, |
| enable_thinking=enable_thinking, |
| stream=requested_stream, |
| timeout=timeout, |
| ) |
| except Exception as exc: |
| errors.append(f"attempt {attempt} stream={requested_stream}: {exc!r}") |
| if requested_stream and stream_fallback: |
| try: |
| return _chat_completion_once( |
| messages=messages, |
| model=model, |
| base_url=base_url, |
| api_key=api_key, |
| temperature=temperature, |
| enable_thinking=enable_thinking, |
| stream=False, |
| timeout=timeout, |
| ) |
| except Exception as fallback_exc: |
| errors.append(f"attempt {attempt} stream=False fallback: {fallback_exc!r}") |
| if attempt >= max_retries: |
| break |
| sleep_sec = base_sleep * (2 ** (attempt - 1)) + random.uniform(0, min(1.0, base_sleep)) |
| time.sleep(sleep_sec) |
|
|
| raise RuntimeError("LLM chat completion failed after retries:\n" + "\n".join(errors[-8:])) |
|
|
|
|
| def _chat_completion_once( |
| messages: List[Dict[str, str]], |
| model: str, |
| base_url: str, |
| api_key: str, |
| temperature: float, |
| enable_thinking: bool, |
| stream: bool, |
| timeout: int, |
| ) -> str: |
| force_stdlib = os.environ.get("LLM_API_FORCE_STDLIB", "0").lower() in {"1", "true", "yes"} |
|
|
| try: |
| if force_stdlib: |
| raise ImportError("LLM_API_FORCE_STDLIB is set") |
| from openai import OpenAI |
| except Exception: |
| return _chat_completion_http( |
| messages=messages, |
| model=model, |
| base_url=base_url, |
| api_key=api_key, |
| temperature=temperature, |
| enable_thinking=enable_thinking, |
| stream=stream, |
| timeout=timeout, |
| ) |
|
|
| try: |
| client = OpenAI(api_key=api_key, base_url=base_url, timeout=timeout, max_retries=0) |
| except TypeError: |
| client = OpenAI(api_key=api_key, base_url=base_url) |
| completion = client.chat.completions.create( |
| model=model, |
| messages=messages, |
| temperature=temperature, |
| extra_body={"enable_thinking": enable_thinking}, |
| stream=stream, |
| ) |
| if not stream: |
| msg = completion.choices[0].message |
| content = getattr(msg, "content", None) |
| if content: |
| return content.strip() |
| |
| try: |
| return msg.model_dump_json(ensure_ascii=False) |
| except Exception: |
| return str(msg) |
|
|
| parts: List[str] = [] |
| for chunk in completion: |
| if not chunk.choices: |
| continue |
| delta = chunk.choices[0].delta |
| content = getattr(delta, "content", None) |
| if content: |
| parts.append(content) |
| result = "".join(parts).strip() |
| if not result: |
| raise RuntimeError("Empty streaming content from chat completions.") |
| return result |
|
|
|
|
| def extract_json_object(text: str) -> Dict[str, Any]: |
| cleaned = text.strip() |
| cleaned = cleaned.replace("```json", "").replace("```", "") |
| try: |
| return json.loads(cleaned) |
| except json.JSONDecodeError: |
| start = cleaned.find("{") |
| end = cleaned.rfind("}") |
| if start == -1 or end == -1 or end <= start: |
| raise |
| return json.loads(cleaned[start : end + 1]) |
|
|
|
|
| def _chat_completion_http( |
| messages: List[Dict[str, str]], |
| model: str, |
| base_url: str, |
| api_key: str, |
| temperature: float, |
| enable_thinking: bool, |
| stream: bool, |
| timeout: int, |
| ) -> str: |
| parsed_url = urlparse(base_url.rstrip("/")) |
| base_path = parsed_url.path.rstrip("/") |
| if base_path.endswith("/chat/completions"): |
| path = base_path |
| else: |
| path = base_path + "/chat/completions" |
| body = json.dumps( |
| { |
| "model": model, |
| "messages": messages, |
| "temperature": temperature, |
| "enable_thinking": enable_thinking, |
| "stream": stream, |
| }, |
| ensure_ascii=False, |
| ).encode("utf-8") |
| conn = _make_https_connection(parsed_url.netloc, timeout) |
| try: |
| conn.request( |
| "POST", |
| path, |
| body=body, |
| headers={ |
| "Authorization": f"Bearer {api_key}", |
| "Content-Type": "application/json", |
| "Accept": "text/event-stream" if stream else "application/json", |
| }, |
| ) |
| resp = conn.getresponse() |
| if resp.status >= 400: |
| error_body = resp.read().decode("utf-8", errors="ignore") |
| raise RuntimeError(f"HTTP {resp.status} from chat completions: {error_body[:1000]}") |
| if not stream: |
| raw_body = resp.read().decode("utf-8", errors="ignore") |
| obj = json.loads(raw_body) |
| content = _extract_message_content(obj) |
| if content: |
| return content.strip() |
| raise RuntimeError(f"Empty message.content from chat completions. Raw response: {raw_body[:2000]}") |
|
|
| parts: List[str] = [] |
| pending = b"" |
| while True: |
| data = resp.read(1024) |
| if not data: |
| break |
| pending += data |
| while b"\n" in pending: |
| line, pending = pending.split(b"\n", 1) |
| line_text = line.decode("utf-8", errors="ignore").strip() |
| if not line_text.startswith("data:"): |
| continue |
| payload = line_text[5:].strip() |
| if payload == "[DONE]": |
| return "".join(parts).strip() |
| try: |
| obj = json.loads(payload) |
| except json.JSONDecodeError: |
| continue |
| for choice in obj.get("choices", []): |
| content = (choice.get("delta") or {}).get("content") |
| if content: |
| parts.append(content) |
| result = "".join(parts).strip() |
| if not result: |
| raise RuntimeError("Empty streaming content from fallback HTTP client.") |
| return result |
| finally: |
| conn.close() |
|
|
|
|
| def _extract_message_content(obj: Dict[str, Any]) -> str: |
| choices = obj.get("choices") or [] |
| if not choices: |
| return "" |
| message = choices[0].get("message") or {} |
| content = message.get("content") |
| if isinstance(content, str): |
| return content |
| if isinstance(content, list): |
| parts = [] |
| for part in content: |
| if isinstance(part, dict): |
| text = part.get("text") or part.get("content") |
| if text: |
| parts.append(str(text)) |
| elif isinstance(part, str): |
| parts.append(part) |
| return "\n".join(parts) |
| return "" |
|
|
|
|
| def _make_https_connection(target_netloc: str, timeout: int) -> http.client.HTTPSConnection: |
| proxy = ( |
| os.environ.get("HTTPS_PROXY") |
| or os.environ.get("https_proxy") |
| or os.environ.get("ALL_PROXY") |
| or os.environ.get("all_proxy") |
| ) |
| if not proxy: |
| return http.client.HTTPSConnection(target_netloc, timeout=timeout) |
| proxy_url = proxy if "://" in proxy else f"http://{proxy}" |
| parsed_proxy = urlparse(proxy_url) |
| if parsed_proxy.scheme not in {"http", "https"} or not parsed_proxy.hostname: |
| raise RuntimeError(f"Unsupported proxy URL: {proxy}") |
| proxy_port = parsed_proxy.port or (443 if parsed_proxy.scheme == "https" else 8080) |
| conn = http.client.HTTPSConnection(parsed_proxy.hostname, proxy_port, timeout=timeout) |
| conn.set_tunnel(target_netloc) |
| return conn |
|
|
