1 line
12 KiB
JSON
1 line
12 KiB
JSON
{"repo": "aashravs/fmcg-ai-analytics", "files": {"llm.py": {"6909aceb1c0c872a493a8179f9d901288989b3f9": {"t": 1785674394.0409923, "n": 12055, "h": "77fdfc18896c04ae", "c": "from __future__ import annotations\n\nimport json\nimport logging\nimport os\nimport time\nfrom dataclasses import dataclass, field\nfrom typing import Any\nfrom urllib.error import HTTPError, URLError\nfrom urllib.request import Request, urlopen\n\nlogger = logging.getLogger(__name__)\n\n\ndef _load_dotenv() -> None:\n base_dir = os.path.dirname(os.path.abspath(__file__))\n env_path = os.path.join(base_dir, \".env\")\n if not os.path.exists(env_path):\n env_path = \".env\"\n if os.path.exists(env_path):\n try:\n with open(env_path, \"r\", encoding=\"utf-8\") as f:\n for line in f:\n line = line.strip()\n if not line or line.startswith(\"#\"):\n continue\n if \"=\" not in line:\n continue\n key, val = line.split(\"=\", 1)\n key = key.strip()\n val = val.strip()\n if len(val) >= 2 and (\n (val[0] == '\"' and val[-1] == '\"') or\n (val[0] == \"'\" and val[-1] == \"'\")\n ):\n val = val[1:-1]\n if key and key not in os.environ:\n os.environ[key] = val\n except Exception:\n pass\n\n\n_load_dotenv()\n\n\nclass LLMError(Exception):\n pass\n\n\nclass LLMAuthenticationError(LLMError):\n pass\n\n\nclass LLMRateLimitError(LLMError):\n def __init__(self, message: str, retry_after: float | None = None) -> None:\n super().__init__(message)\n self.retry_after = retry_after\n\n\nclass LLMServerError(LLMError):\n def __init__(self, message: str, status_code: int) -> None:\n super().__init__(message)\n self.status_code = status_code\n\n\nclass LLMTimeoutError(LLMError):\n pass\n\n\nclass LLMResponseError(LLMError):\n pass\n\n\n@dataclass\nclass LLMResponse:\n raw_text: str\n model: str\n usage: dict[str, int] = field(default_factory=dict)\n finish_reason: str = \"stop\"\n\n def as_json(self) -> dict[str, Any] | None:\n try:\n return json.loads(self.raw_text)\n except (json.JSONDecodeError, TypeError):\n return None\n\n\n_RETRYABLE_STATUS_CODES: frozenset[int] = frozenset({429, 500, 502, 503, 504})\n_BASE_DELAY_SECONDS: float = 1.0\n_MAX_DELAY_SECONDS: float = 30.0\n_BACKOFF_MULTIPLIER: float = 2.0\n\n\nclass LLMClient:\n DEFAULT_MODEL: str = \"grok-3-mini\"\n DEFAULT_BASE_URL: str = \"https://api.x.ai/v1\"\n DEFAULT_TIMEOUT: int = 60\n\n def __init__(\n self,\n api_key: str | None = None,\n model: str = DEFAULT_MODEL,\n base_url: str = DEFAULT_BASE_URL,\n max_retries: int = 3,\n timeout_secs: int = DEFAULT_TIMEOUT,\n ) -> None:\n grok_key = os.getenv(\"GROK_API_KEY\", \"\")\n groq_key = os.getenv(\"GROQ_API_KEY\", \"\")\n resolved_key = api_key\n if not resolved_key:\n if groq_key.startswith(\"gsk_\"):\n resolved_key = groq_key\n elif grok_key.startswith(\"gsk_\"):\n resolved_key = grok_key\n else:\n resolved_key = grok_key or groq_key\n self.api_key: str = resolved_key\n self.model: str = model\n self.base_url: str = base_url.rstrip(\"/\")\n self.max_retries: int = max_retries\n self.timeout_secs: int = timeout_secs\n is_groq = self.api_key.startswith(\"gsk_\")\n if is_groq:\n if self.base_url == self.DEFAULT_BASE_URL:\n self.base_url = \"https://api.groq.com/openai/v1\"\n if self.model == self.DEFAULT_MODEL or self.model == \"grok-3-mini\":\n self.model = \"llama-3.3-70b-versatile\"\n logger.info(\"Groq API key detected. Using Groq base URL and llama-3.3-70b-versatile model.\")\n if not self.api_key:\n raise LLMAuthenticationError(\n \"Grok or Groq API key is required. \"\n \"Pass api_key= or set the GROK_API_KEY or GROQ_API_KEY environment variable.\"\n )\n\n def complete(\n self,\n prompt: str,\n system_prompt: str = \"\",\n temperature: float = 0.0,\n max_tokens: int = 2048,\n response_format: str | None = None,\n ) -> LLMResponse:\n messages = self._build_messages(prompt, system_prompt)\n payload = self._build_payload(messages, temperature, max_tokens, response_format)\n response_body = self._send_with_retries(payload)\n return self._parse_response(response_body)\n\n def _build_messages(self, prompt: str, system_prompt: str) -> list[dict[str, str]]:\n messages: list[dict[str, str]] = []\n if system_prompt:\n messages.append({\"role\": \"system\", \"content\": system_prompt})\n messages.append({\"role\": \"user\", \"content\": prompt})\n return messages\n\n def _build_payload(\n self,\n messages: list[dict[str, str]],\n temperature: float,\n max_tokens: int,\n response_format: str | None,\n ) -> dict[str, Any]:\n payload: dict[str, Any] = {\n \"model\": self.model,\n \"messages\": messages,\n \"temperature\": temperature,\n \"max_tokens\": max_tokens,\n }\n if response_format:\n payload[\"response_format\"] = {\"type\": response_format}\n return payload\n\n def _send_with_retries(self, payload: dict[str, Any]) -> dict[str, Any]:\n url = f\"{self.base_url}/chat/completions\"\n headers = {\n \"Content-Type\": \"application/json\",\n \"Authorization\": f\"Bearer {self.api_key}\",\n \"User-Agent\": \"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36\",\n }\n body = json.dumps(payload).encode(\"utf-8\")\n last_exception: Exception | None = None\n for attempt in range(1, self.max_retries + 1):\n try:\n logger.debug(\n \"Attempt %d/%d → POST %s (model=%s)\",\n attempt, self.max_retries, url, self.model,\n )\n return self._send_request(url, headers, body)\n except LLMAuthenticationError:\n raise\n except LLMRateLimitError as exc:\n last_exception = exc\n delay = exc.retry_after or self._backoff_delay(attempt)\n logger.warning(\n \"Rate limited (429). Retrying in %.1fs (attempt %d/%d).\",\n delay, attempt, self.max_retries,\n )\n except LLMServerError as exc:\n last_exception = exc\n delay = self._backoff_delay(attempt)\n logger.warning(\n \"Server error (HTTP %d). Retrying in %.1fs (attempt %d/%d).\",\n exc.status_code, delay, attempt, self.max_retries,\n )\n except LLMTimeoutError as exc:\n last_exception = exc\n delay = self._backoff_delay(attempt)\n logger.warning(\n \"Request timed out. Retrying in %.1fs (attempt %d/%d).\",\n delay, attempt, self.max_retries,\n )\n if attempt < self.max_retries:\n time.sleep(delay)\n raise LLMError(\n f\"Grok API call failed after {self.max_retries} attempts. \"\n f\"Last error: {last_exception}\"\n ) from last_exception\n\n def _send_request(self, url: str, headers: dict[str, str], body: bytes) -> dict[str, Any]:\n request = Request(url, data=body, headers=headers, method=\"POST\")\n try:\n with urlopen(request, timeout=self.timeout_secs) as resp:\n raw = resp.read().decode(\"utf-8\")\n except HTTPError as exc:\n self._handle_http_error(exc)\n except (URLError, OSError, TimeoutError) as exc:\n if _is_timeout(exc):\n raise LLMTimeoutError(\n f\"Request to {url} timed out after {self.timeout_secs}s.\"\n ) from exc\n raise LLMError(f\"Network error while contacting {url}: {exc}\") from exc\n try:\n return json.loads(raw)\n except json.JSONDecodeError as exc:\n raise LLMResponseError(\n f\"Failed to parse API response as JSON. \"\n f\"Raw response (first 500 chars): {raw[:500]}\"\n ) from exc\n\n def _handle_http_error(self, exc: HTTPError) -> None:\n status = exc.code\n raw_body = \"\"\n try:\n raw_body = exc.read().decode(\"utf-8\")\n except Exception:\n pass\n try:\n error_detail = json.loads(raw_body).get(\"error\", {}).get(\"message\", raw_body[:300])\n except Exception:\n error_detail = raw_body[:300] if raw_body else str(exc.reason)\n if status in (401, 403):\n raise LLMAuthenticationError(\n f\"Authentication failed (HTTP {status}): {error_detail}\\n\"\n f\" Requested URL: {exc.url}\\n\"\n f\" Model used: {self.model}\\n\"\n f\" API Key: {self.api_key[:10]}...\"\n ) from exc\n if status == 429:\n retry_after = _parse_retry_after(exc.headers.get(\"Retry-After\"))\n raise LLMRateLimitError(\n f\"Rate limit exceeded (HTTP 429): {error_detail}\",\n retry_after=retry_after,\n ) from exc\n if status >= 500:\n raise LLMServerError(\n f\"Server error (HTTP {status}): {error_detail}\",\n status_code=status,\n ) from exc\n raise LLMError(\n f\"API request failed (HTTP {status}): {error_detail}\\n\"\n f\" Requested URL: {exc.url}\\n\"\n f\" Model used: {self.model}\"\n ) from exc\n\n def _parse_response(self, body: dict[str, Any]) -> LLMResponse:\n try:\n choices = body[\"choices\"]\n if not choices:\n raise LLMResponseError(\n \"API returned an empty 'choices' array. \"\n f\"Full response: {json.dumps(body)[:500]}\"\n )\n first_choice = choices[0]\n raw_text = first_choice[\"message\"][\"content\"]\n finish_reason = first_choice.get(\"finish_reason\", \"stop\")\n model = body.get(\"model\", self.model)\n usage = body.get(\"usage\", {})\n return LLMResponse(\n raw_text=raw_text,\n model=model,\n usage={\n \"prompt_tokens\": usage.get(\"prompt_tokens\", 0),\n \"completion_tokens\": usage.get(\"completion_tokens\", 0),\n \"total_tokens\": usage.get(\"total_tokens\", 0),\n },\n finish_reason=finish_reason,\n )\n except (KeyError, IndexError, TypeError) as exc:\n raise LLMResponseError(\n f\"Malformed API response — could not extract completion. \"\n f\"Missing key: {exc}. \"\n f\"Response body (first 500 chars): {json.dumps(body)[:500]}\"\n ) from exc\n\n @staticmethod\n def _backoff_delay(attempt: int) -> float:\n delay = _BASE_DELAY_SECONDS * (_BACKOFF_MULTIPLIER ** (attempt - 1))\n return min(delay, _MAX_DELAY_SECONDS)\n\n def __repr__(self) -> str:\n return (\n f\"LLMClient(model={self.model!r}, \"\n f\"base_url={self.base_url!r}, \"\n f\"max_retries={self.max_retries}, \"\n f\"timeout_secs={self.timeout_secs})\"\n )\n\n\ndef _is_timeout(exc: Exception) -> bool:\n if isinstance(exc, TimeoutError):\n return True\n if isinstance(exc, URLError) and isinstance(exc.reason, TimeoutError):\n return True\n if isinstance(exc, OSError) and \"timed out\" in str(exc).lower():\n return True\n return False\n\n\ndef _parse_retry_after(header_value: str | None) -> float | None:\n if header_value is None:\n return None\n try:\n return float(header_value)\n except ValueError:\n return None\n"}}}} |