1 line
20 KiB
JSON
1 line
20 KiB
JSON
{"repo": "0xkript/logicbreaker", "files": {"agents/llm_client.py": {"f7f19c8e072e016a5df9f2479e3b275615d228cc": {"t": 1785674548.9802513, "n": 18587, "h": "48d1a79482f50484", "c": "\"\"\"\nMulti-provider LLM client\n=========================\n\nA single interface over the major LLM APIs. The user picks a provider at\nstartup (or runs with no provider at all -- \"fast scan\" mode, fully heuristic).\n\nEach provider entry defines its endpoint, auth header style, and request/\nresponse shape, so adding a new provider is a one-line dict entry.\n\nThe client is used only for *optional* triage/enrichment and suggested fixes.\nThe core detection + dynamic exploitation + verified patching does not depend\non it, so \"fast scan\" mode is fully functional.\n\"\"\"\n\nimport ast\nimport hashlib\nimport json\nimport os\nimport re\nimport tempfile\n\nimport requests\n\n\ndef _parse_json_lenient(raw):\n \"\"\"Parse a model's JSON reply ROBUSTLY (FIX A, determinism).\n\n Models occasionally wrap the JSON in prose/markdown, add a trailing comma, or\n emit Python-style quoting. A brittle json.loads would drop an otherwise-valid\n verdict; because a dropped verdict used to go UNCACHED, the next run (which\n parsed cleanly) would relabel the finding -- the exact run-to-run drift FIX A\n targets. We try progressively more tolerant parses and only give up if none\n yields an object. Pure function of `raw`, so it is itself deterministic.\"\"\"\n if raw is None:\n raise ValueError(\"empty LLM response\")\n cleaned = re.sub(r\"^```(?:json)?|```$\", \"\", raw.strip(), flags=re.MULTILINE).strip()\n candidates = [cleaned]\n m = re.search(r\"\\{.*\\}\", cleaned, flags=re.DOTALL)\n if m and m.group(0) != cleaned:\n candidates.append(m.group(0))\n for cand in candidates:\n # 1) strict JSON\n try:\n return json.loads(cand)\n except Exception:\n pass\n # 2) JSON with trailing commas removed\n try:\n return json.loads(re.sub(r\",(\\s*[}\\]])\", r\"\\1\", cand))\n except Exception:\n pass\n # 3) Python-literal fallback (single quotes / True/False/None)\n try:\n v = ast.literal_eval(cand)\n if isinstance(v, dict):\n return v\n except Exception:\n pass\n raise ValueError(\"LLM response was not parseable as a JSON object\")\n\n\nPROVIDERS = {\n \"claude\": {\n \"label\": \"Anthropic Claude (claude-sonnet-4-5 / opus)\",\n \"env\": \"ANTHROPIC_API_KEY\",\n \"url\": \"https://api.anthropic.com/v1/messages\",\n \"default_model\": \"claude-sonnet-4-5\",\n \"style\": \"anthropic\",\n \"key_hint\": \"sk-ant-...\",\n },\n \"openai\": {\n \"label\": \"OpenAI (gpt-4o / gpt-4.1)\",\n \"env\": \"OPENAI_API_KEY\",\n \"url\": \"https://api.openai.com/v1/chat/completions\",\n \"default_model\": \"gpt-4o\",\n \"style\": \"openai\",\n \"key_hint\": \"sk-...\",\n },\n \"gemini\": {\n \"label\": \"Google Gemini (gemini-2.0-flash / 1.5-pro)\",\n \"env\": \"GEMINI_API_KEY\",\n \"url\": \"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent\",\n \"default_model\": \"gemini-2.0-flash\",\n \"style\": \"gemini\",\n \"key_hint\": \"AIza...\",\n },\n \"groq\": {\n \"label\": \"Groq (llama-3.3-70b, very fast)\",\n \"env\": \"GROQ_API_KEY\",\n \"url\": \"https://api.groq.com/openai/v1/chat/completions\",\n \"default_model\": \"llama-3.3-70b-versatile\",\n \"style\": \"openai\",\n \"key_hint\": \"gsk_...\",\n },\n \"mistral\": {\n \"label\": \"Mistral AI (mistral-large)\",\n \"env\": \"MISTRAL_API_KEY\",\n \"url\": \"https://api.mistral.ai/v1/chat/completions\",\n \"default_model\": \"mistral-large-latest\",\n \"style\": \"openai\",\n \"key_hint\": \"...\",\n },\n \"deepseek\": {\n \"label\": \"DeepSeek (deepseek-chat)\",\n \"env\": \"DEEPSEEK_API_KEY\",\n \"url\": \"https://api.deepseek.com/chat/completions\",\n \"default_model\": \"deepseek-chat\",\n \"style\": \"openai\",\n \"key_hint\": \"sk-...\",\n },\n \"together\": {\n \"label\": \"Together AI (Llama / Qwen / Mixtral)\",\n \"env\": \"TOGETHER_API_KEY\",\n \"url\": \"https://api.together.xyz/v1/chat/completions\",\n \"default_model\": \"meta-llama/Llama-3.3-70B-Instruct-Turbo\",\n \"style\": \"openai\",\n \"key_hint\": \"...\",\n },\n \"openrouter\": {\n \"label\": \"OpenRouter (gateway to 100+ models)\",\n \"env\": \"OPENROUTER_API_KEY\",\n \"url\": \"https://openrouter.ai/api/v1/chat/completions\",\n \"default_model\": \"anthropic/claude-3.5-sonnet\",\n \"style\": \"openai\",\n \"key_hint\": \"sk-or-...\",\n },\n \"xai\": {\n \"label\": \"xAI Grok (grok-2)\",\n \"env\": \"XAI_API_KEY\",\n \"url\": \"https://api.x.ai/v1/chat/completions\",\n \"default_model\": \"grok-2-latest\",\n \"style\": \"openai\",\n \"key_hint\": \"xai-...\",\n },\n \"cohere\": {\n \"label\": \"Cohere (command-r-plus)\",\n \"env\": \"COHERE_API_KEY\",\n \"url\": \"https://api.cohere.com/v2/chat\",\n \"default_model\": \"command-r-plus\",\n \"style\": \"openai\",\n \"key_hint\": \"...\",\n },\n \"perplexity\": {\n \"label\": \"Perplexity (sonar-pro, web-grounded)\",\n \"env\": \"PERPLEXITY_API_KEY\",\n \"url\": \"https://api.perplexity.ai/chat/completions\",\n \"default_model\": \"sonar-pro\",\n \"style\": \"openai\",\n \"key_hint\": \"pplx-...\",\n },\n \"fireworks\": {\n \"label\": \"Fireworks AI (Llama / Qwen / DeepSeek, fast)\",\n \"env\": \"FIREWORKS_API_KEY\",\n \"url\": \"https://api.fireworks.ai/inference/v1/chat/completions\",\n \"default_model\": \"accounts/fireworks/models/llama-v3p3-70b-instruct\",\n \"style\": \"openai\",\n \"key_hint\": \"fw_...\",\n },\n \"cerebras\": {\n \"label\": \"Cerebras (llama-3.3-70b, fastest inference)\",\n \"env\": \"CEREBRAS_API_KEY\",\n \"url\": \"https://api.cerebras.ai/v1/chat/completions\",\n \"default_model\": \"llama-3.3-70b\",\n \"style\": \"openai\",\n \"key_hint\": \"csk-...\",\n },\n \"sambanova\": {\n \"label\": \"SambaNova (Llama 3.3 70B / 405B)\",\n \"env\": \"SAMBANOVA_API_KEY\",\n \"url\": \"https://api.sambanova.ai/v1/chat/completions\",\n \"default_model\": \"Meta-Llama-3.3-70B-Instruct\",\n \"style\": \"openai\",\n \"key_hint\": \"...\",\n },\n \"nebius\": {\n \"label\": \"Nebius AI Studio (Llama / Qwen / DeepSeek)\",\n \"env\": \"NEBIUS_API_KEY\",\n \"url\": \"https://api.studio.nebius.com/v1/chat/completions\",\n \"default_model\": \"meta-llama/Llama-3.3-70B-Instruct\",\n \"style\": \"openai\",\n \"key_hint\": \"...\",\n },\n}\n\nMENU_ORDER = [\"claude\", \"openai\", \"gemini\", \"groq\", \"mistral\",\n \"deepseek\", \"together\", \"openrouter\", \"xai\", \"cohere\",\n \"perplexity\", \"fireworks\", \"cerebras\", \"sambanova\", \"nebius\"]\n\n\nclass LLMClient:\n def __init__(self, provider=None, api_key=None, model=None):\n self.provider = provider\n self.spec = PROVIDERS.get(provider) if provider else None\n self.api_key = api_key or (os.environ.get(self.spec[\"env\"]) if self.spec else None)\n self.model = model or (self.spec[\"default_model\"] if self.spec else None)\n\n @property\n def available(self) -> bool:\n return bool(self.spec and self.api_key)\n\n def validate_key(self, timeout=20.0):\n \"\"\"\n Really verify the API key works with the selected provider's model by\n sending a tiny test request. Returns (ok: bool, message: str).\n\n This catches: empty keys, malformed keys, random text, and keys for the\n wrong provider -- the provider itself rejects them (401/403), so we\n report a clear error instead of silently 'accepting' a bad key.\n\n CRITICAL: this method BYPASSES the cache entirely. If we used the cache,\n a previous valid key's response would make an invalid key appear valid.\n \"\"\"\n if not self.spec:\n return False, \"No provider selected.\"\n if not self.api_key or not self.api_key.strip():\n return False, \"No API key provided.\"\n\n key = self.api_key.strip()\n # cheap format sanity check per provider (catches obvious typos / Enter)\n fmt_ok, fmt_msg = self._key_format_ok(key)\n if not fmt_ok:\n return False, fmt_msg\n\n # live check: one minimal request. BYPASS THE CACHE so we actually\n # contact the provider and verify THIS key (not a cached response from\n # a previous key). Any non-auth error (e.g. rate limit) still means\n # the key itself is valid.\n old_cache_setting = os.environ.get(\"LB_NO_LLM_CACHE\", \"\")\n os.environ[\"LB_NO_LLM_CACHE\"] = \"1\"\n try:\n self.chat(\"You are a test.\", \"Reply with: ok\", max_tokens=5, timeout=timeout)\n return True, f\"API key verified — {self.provider} ({self.model}) is reachable.\"\n except RuntimeError as e:\n msg = str(e).lower()\n # ONLY treat as auth failure if the error is SPECIFICALLY about\n # authentication (not just any error containing \"invalid\")\n if \"401\" in msg or \"403\" in msg or \"unauthor\" in msg \\\n or \"api key\" in msg or \"authentication\" in msg \\\n or \"invalid api key\" in msg or \"invalid x-api-key\" in msg \\\n or \"permission denied\" in msg or \"forbidden\" in msg:\n return False, (f\"The API key was rejected by {self.provider} \"\n f\"(authentication failed). Check the key is correct and is for \"\n f\"{self.provider}.\")\n if \"429\" in msg or \"rate\" in msg or \"quota\" in msg or \"billing\" in msg \\\n or \"insufficient\" in msg:\n # key is valid, just rate-limited / out of credit\n return True, (f\"API key is valid for {self.provider}, but the account is rate-limited \"\n f\"or out of credit. Scanning can proceed; LLM triage may be limited.\")\n if \"404\" in msg or (\"model\" in msg and \"not found\" in msg):\n return False, (f\"The key may be valid but model '{self.model}' was not found for \"\n f\"{self.provider}. Try a different model.\")\n # network / unknown -> report it so the user isn't misled\n return False, f\"Could not verify the key with {self.provider}: {str(e)[:160]}\"\n except Exception as e:\n return False, f\"Could not reach {self.provider}: {str(e)[:160]}\"\n finally:\n # restore the original cache setting\n if old_cache_setting:\n os.environ[\"LB_NO_LLM_CACHE\"] = old_cache_setting\n else:\n os.environ.pop(\"LB_NO_LLM_CACHE\", None)\n\n def _key_format_ok(self, key):\n \"\"\"Per-provider format check to reject obvious garbage before a network\n call. Conservative: only rejects clearly-wrong shapes.\n\n We do NOT reject keys based on prefix alone for providers whose key\n formats vary. The real validation is the live API request in\n validate_key() -- this is just a quick sanity check.\"\"\"\n # universal: must not contain spaces, must be a reasonable length\n if \" \" in key or \"\\t\" in key:\n return False, \"API key contains spaces — that is not a valid key.\"\n if len(key) < 12:\n return False, \"API key is too short to be valid.\"\n # Only check prefix for providers with a VERY strict, well-known format.\n # For providers with varied formats, skip the prefix check and let the\n # live API request decide.\n strict_prefix = {\n \"groq\": (\"gsk_\",),\n \"xai\": (\"xai-\",),\n \"perplexity\": (\"pplx-\",),\n \"fireworks\": (\"fw_\",),\n \"cerebras\": (\"csk-\",),\n }.get(self.provider, ())\n if strict_prefix and not any(key.startswith(p) for p in strict_prefix):\n return False, (f\"This does not look like a {self.provider} API key \"\n f\"(expected it to start with {' or '.join(strict_prefix)}).\")\n return True, \"\"\n\n def chat(self, system_prompt, user_prompt, temperature=0.0, max_tokens=1024, timeout=45.0):\n # DETERMINISM: default temperature is 0.0. Classification, triage, the\n # safety net and the verdict must be reproducible -- the same code must\n # yield the same result on every run. There is no reason to sample here.\n if not self.available:\n raise RuntimeError(\"No LLM provider configured\")\n\n # content-hash cache for deterministic calls. A temperature-0 call\n # on the same (system, user, max_tokens) inputs yields the same reply,\n # so caching is sound and saves cost+time on re-runs of unchanged files.\n # The cache is on disk so it survives across runs. Disable with\n # LB_NO_LLM_CACHE=1.\n cache_disabled = os.environ.get(\"LB_NO_LLM_CACHE\", \"\") in (\"1\", \"true\", \"True\")\n cache_key = None\n if not cache_disabled and temperature == 0.0:\n cache_key = self._cache_key(system_prompt, user_prompt, max_tokens)\n cached = self._cache_get(cache_key)\n if cached is not None:\n return cached\n\n style = self.spec[\"style\"]\n if style == \"anthropic\":\n result = self._anthropic(system_prompt, user_prompt, temperature, max_tokens, timeout)\n elif style == \"gemini\":\n result = self._gemini(system_prompt, user_prompt, temperature, max_tokens, timeout)\n else:\n result = self._openai(system_prompt, user_prompt, temperature, max_tokens, timeout)\n\n if cache_key is not None:\n self._cache_put(cache_key, result)\n return result\n\n def chat_json(self, system_prompt, user_prompt, **kwargs):\n # JSON/analysis calls are ALWAYS deterministic, regardless of any future\n # change to chat()'s default: force temperature 0 unless explicitly set.\n kwargs.setdefault(\"temperature\", 0.0)\n raw = self.chat(system_prompt, user_prompt, **kwargs)\n # Robust, deterministic parse (handles fences / prose / trailing commas /\n # single-quoted output). Raises ValueError if nothing parses.\n return _parse_json_lenient(raw)\n\n # ---- disk cache for deterministic LLM calls ----------------------\n\n def _cache_key(self, system_prompt: str, user_prompt: str, max_tokens: int) -> str:\n \"\"\"Stable content-hash key for a deterministic call.\n\n Includes a hash of the API key so that changing the key invalidates\n the cache (a different key must never return a cached response from\n a previous key).\"\"\"\n key_hash = hashlib.sha256(self.api_key.encode(\"utf-8\", \"replace\")).hexdigest()[:16] if self.api_key else \"no-key\"\n raw = f\"{self.provider}|{self.model}|{max_tokens}|{key_hash}|{system_prompt}|{user_prompt}\"\n return hashlib.sha256(raw.encode(\"utf-8\", \"replace\")).hexdigest()\n\n def _cache_dir(self) -> str:\n d = os.path.expanduser(\n os.environ.get(\"LB_CACHE_DIR\", \"~/.logicbreaker/cache\"))\n try:\n os.makedirs(d, exist_ok=True)\n except OSError:\n d = os.path.join(tempfile.gettempdir(), \"logicbreaker_cache\")\n os.makedirs(d, exist_ok=True)\n return d\n\n def _cache_get(self, key: str):\n path = os.path.join(self._cache_dir(), f\"{key}.txt\")\n try:\n with open(path, \"r\", encoding=\"utf-8\") as f:\n return f.read()\n except OSError:\n return None\n\n def _cache_put(self, key: str, value: str):\n path = os.path.join(self._cache_dir(), f\"{key}.txt\")\n try:\n with open(path, \"w\", encoding=\"utf-8\") as f:\n f.write(value)\n except OSError:\n pass # cache is best-effort; never fail a call because of it\n\n def _redact(self, text):\n \"\"\"Belt-and-suspenders: never let the API key surface in an error string,\n log line, or anything we raise/print. Shows only the last 4 chars.\"\"\"\n text = text or \"\"\n if self.api_key and self.api_key in text:\n text = text.replace(self.api_key, \"****\" + self.api_key[-4:])\n return text\n\n def _openai(self, system, user, temperature, max_tokens, timeout):\n headers = {\"Authorization\": f\"Bearer {self.api_key}\", \"Content-Type\": \"application/json\"}\n body = {\n \"model\": self.model,\n \"messages\": [{\"role\": \"system\", \"content\": system},\n {\"role\": \"user\", \"content\": user}],\n \"temperature\": temperature, \"max_tokens\": max_tokens,\n }\n r = requests.post(self.spec[\"url\"], headers=headers, json=body, timeout=timeout)\n if r.status_code != 200:\n raise RuntimeError(self._redact(f\"{self.provider} error {r.status_code}: {r.text[:300]}\"))\n return r.json()[\"choices\"][0][\"message\"][\"content\"]\n\n def _anthropic(self, system, user, temperature, max_tokens, timeout):\n headers = {\"x-api-key\": self.api_key, \"anthropic-version\": \"2023-06-01\",\n \"Content-Type\": \"application/json\"}\n body = {\n \"model\": self.model, \"max_tokens\": max_tokens, \"temperature\": temperature,\n \"system\": system,\n \"messages\": [{\"role\": \"user\", \"content\": user}],\n }\n r = requests.post(self.spec[\"url\"], headers=headers, json=body, timeout=timeout)\n if r.status_code != 200:\n raise RuntimeError(self._redact(f\"claude error {r.status_code}: {r.text[:300]}\"))\n data = r.json()\n return \"\".join(block.get(\"text\", \"\") for block in data.get(\"content\", []))\n\n def _gemini(self, system, user, temperature, max_tokens, timeout):\n # SECURITY: pass the key in a HEADER, never as a `?key=` URL query param,\n # so it cannot leak into an exception string, a proxy/access log, or\n # terminal output if the request fails or is traced.\n url = self.spec[\"url\"].format(model=self.model)\n headers = {\"Content-Type\": \"application/json\", \"x-goog-api-key\": self.api_key}\n body = {\n \"system_instruction\": {\"parts\": [{\"text\": system}]},\n \"contents\": [{\"parts\": [{\"text\": user}]}],\n \"generationConfig\": {\"temperature\": temperature, \"maxOutputTokens\": max_tokens},\n }\n r = requests.post(url, headers=headers, json=body, timeout=timeout)\n if r.status_code != 200:\n raise RuntimeError(self._redact(f\"gemini error {r.status_code}: {r.text[:300]}\"))\n data = r.json()\n return data[\"candidates\"][0][\"content\"][\"parts\"][0][\"text\"]\n\n\ndef provider_menu():\n return [(k, PROVIDERS[k][\"label\"], PROVIDERS[k][\"key_hint\"]) for k in MENU_ORDER]\n"}}}} |