#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Elephant Alpha AI 服務模組 負責與 OpenRouter / Elephant Alpha API 互動，提供高效率、長上下文的 Worker AI 功能 """ import os import time import json import logging import requests from typing import Optional, Dict, Any, List from dataclasses import dataclass logger = logging.getLogger(__name__) # Elephant Alpha 設定（NVIDIA NIM API） NVIDIA_API_KEY = os.getenv('NVIDIA_API_KEY', '') ELEPHANT_ALPHA_BASE_URL = os.getenv( 'ELEPHANT_ALPHA_NEMOTRON_NIM_ENDPOINT', 'https://integrate.api.nvidia.com/v1', ).rstrip('/') ELEPHANT_ALPHA_URL = os.getenv( 'ELEPHANT_ALPHA_URL', f"{ELEPHANT_ALPHA_BASE_URL}/chat/completions", ) DEFAULT_ELEPHANT_MODEL = os.getenv( 'ELEPHANT_ALPHA_MODEL', 'nvidia/llama-3.3-nemotron-super-49b-v1.5', ) ELEPHANT_FALLBACK_MODELS = [ model.strip() for model in os.getenv( 'ELEPHANT_ALPHA_FALLBACK_MODELS', 'nvidia/llama-3.3-nemotron-super-49b-v1.5,nvidia/llama-3.1-nemotron-70b-instruct,meta/llama-3.1-8b-instruct', ).split(',') if model.strip() ] ELEPHANT_TIMEOUT = int(os.getenv('ELEPHANT_TIMEOUT', '120')) # 預設 2 分鐘 ELEPHANT_FALLBACK_HTTP_STATUS_CODES = {403, 404, 408, 409, 425, 429, 500, 502, 503, 504} # Elephant Alpha 定價 (USD per 1M tokens) - NVIDIA NIM 定價 ELEPHANT_PRICING = { 'nvidia/llama-3.1-nemotron-ultra-253b-v1': {'input': 0.10, 'output': 0.40}, 'nvidia/llama-3.3-nemotron-super-49b-v1.5': {'input': 0.10, 'output': 0.40}, } @dataclass class ElephantResponse: """Elephant Alpha 回應結構""" success: bool content: str model: str error: Optional[str] = None total_duration: Optional[float] = None input_tokens: int = 0 output_tokens: int = 0 total_tokens: int = 0 input_cost: float = 0.0 output_cost: float = 0.0 total_cost: float = 0.0 class ElephantService: """Elephant Alpha AI 服務 - 100B 效率型 Worker""" def __init__(self, api_key: str = None, model: str = None): """ 初始化 Elephant 服務 """ self.api_key = api_key or NVIDIA_API_KEY self.model = model or DEFAULT_ELEPHANT_MODEL # W3-A: 護欄 2 — 斷線降級 cache (300s TTL，不每次 ping OpenRouter) _connection_cache: Dict[str, Any] = {"ok": False, "checked_at": 0.0} def check_connection(self, cache_seconds: int = 300) -> bool: """ 檢查 API 是否可用（結果快取 300s）。 cache_seconds=300 與 Anthropic prompt cache TTL 對齊，避免每分鐘 EA loop 都打 API。 """ if not self.api_key: return False now = time.time() cache = ElephantService._connection_cache if cache["checked_at"] and (now - cache["checked_at"]) < cache_seconds: return cache["ok"] try: response = self.generate("hi", timeout=10) result = response.success except Exception: result = False ElephantService._connection_cache = {"ok": result, "checked_at": now} return result @staticmethod def calculate_cost(model: str, input_tokens: int, output_tokens: int) -> Dict[str, float]: """計算費用""" pricing = ELEPHANT_PRICING.get(model, ELEPHANT_PRICING['nvidia/llama-3.1-nemotron-ultra-253b-v1']) input_cost = (input_tokens / 1_000_000) * pricing['input'] output_cost = (output_tokens / 1_000_000) * pricing['output'] return { 'input_cost': round(input_cost, 6), 'output_cost': round(output_cost, 6), 'total_cost': round(input_cost + output_cost, 6) } @staticmethod def _model_candidates(primary_model: str) -> List[str]: candidates = [] for model_name in [primary_model, *ELEPHANT_FALLBACK_MODELS]: if model_name and model_name not in candidates: candidates.append(model_name) return candidates @staticmethod def _has_next_model(model_name: str, model_candidates: List[str]) -> bool: return bool(model_candidates) and model_name != model_candidates[-1] def generate(self, prompt: str, model: str = None, system_prompt: str = None, temperature: float = 0.3, json_mode: bool = False, timeout: int = None) -> ElephantResponse: """ 生成文字（主介面） """ primary_model = model or self.model request_timeout = timeout or ELEPHANT_TIMEOUT if not self.api_key: return ElephantResponse(success=False, content='', model=primary_model, error="API Key 未設定") headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json" } messages = [] if system_prompt: messages.append({"role": "system", "content": system_prompt}) messages.append({"role": "user", "content": prompt}) last_error = "" model_candidates = self._model_candidates(primary_model) for model_name in model_candidates: payload = { "model": model_name, "messages": messages, "temperature": temperature, "max_tokens": 8000 } if json_mode: payload["response_format"] = {"type": "json_object"} try: start_time = time.time() response = requests.post( ELEPHANT_ALPHA_URL, json=payload, headers=headers, timeout=request_timeout ) response.raise_for_status() end_time = time.time() data = response.json() message = data["choices"][0]["message"] content = message.get("content") or message.get("reasoning_content") or message.get("reasoning") or "" # Token 用量 usage = data.get("usage", {}) input_tokens = usage.get("prompt_tokens", 0) output_tokens = usage.get("completion_tokens", 0) costs = self.calculate_cost(model_name, input_tokens, output_tokens) return ElephantResponse( success=True, content=content, model=model_name, total_duration=end_time - start_time, input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=input_tokens + output_tokens, input_cost=costs['input_cost'], output_cost=costs['output_cost'], total_cost=costs['total_cost'] ) except requests.HTTPError as e: status_code = e.response.status_code if e.response is not None else None last_error = f"{model_name}: {e}" if status_code in ELEPHANT_FALLBACK_HTTP_STATUS_CODES and self._has_next_model(model_name, model_candidates): logger.warning(f"[Elephant] NIM 模型/API 暫時不可用，改用 fallback: {model_name} ({status_code})") continue logger.error(f"[Elephant] 生成失敗: {e}") return ElephantResponse(success=False, content='', model=model_name, error=last_error) except (requests.Timeout, requests.ConnectionError) as e: last_error = f"{model_name}: {e}" if self._has_next_model(model_name, model_candidates): logger.warning(f"[Elephant] NIM 暫時性連線錯誤，改用 fallback: {model_name} ({e})") continue logger.error(f"[Elephant] 生成失敗: {e}") return ElephantResponse(success=False, content='', model=model_name, error=last_error) except Exception as e: last_error = f"{model_name}: {e}" logger.error(f"[Elephant] 生成失敗: {e}") return ElephantResponse(success=False, content='', model=model_name, error=last_error) failed_model = model_candidates[-1] if model_candidates else primary_model return ElephantResponse(success=False, content='', model=failed_model, error=last_error or "所有 Elephant fallback model 均不可用") # 單例實例 elephant_service = ElephantService()