force update

2026-02-21 15:13:51 -05:00
parent 47c581d037
commit d2f112b50f
4 changed files with 214 additions and 104 deletions
--- a/ai_client.py
+++ b/ai_client.py
@@ -1,4 +1,4 @@
-# ai_client.py
+# ai_client.py
 import tomllib
 from pathlib import Path

@@ -22,6 +22,103 @@ def _load_credentials() -> dict:
    with open("credentials.toml", "rb") as f:
        return tomllib.load(f)

+# ------------------------------------------------------------------ provider errors
+
+class ProviderError(Exception):
+    """
+    Raised when the upstream API returns a hard error we want to surface
+    distinctly in the UI (quota, rate-limit, auth, balance, etc.).
+
+    Attributes
+    ----------
+    kind : str
+        One of: "quota", "rate_limit", "auth", "balance", "network", "unknown"
+    provider : str
+        "gemini" or "anthropic"
+    original : Exception
+        The underlying SDK exception.
+    """
+    def __init__(self, kind: str, provider: str, original: Exception):
+        self.kind = kind
+        self.provider = provider
+        self.original = original
+        super().__init__(str(original))
+
+    # Human-readable banner shown in the Response panel
+    def ui_message(self) -> str:
+        labels = {
+            "quota":      "QUOTA EXHAUSTED",
+            "rate_limit": "RATE LIMITED",
+            "auth":       "AUTH / API KEY ERROR",
+            "balance":    "BALANCE / BILLING ERROR",
+            "network":    "NETWORK / CONNECTION ERROR",
+            "unknown":    "API ERROR",
+        }
+        label = labels.get(self.kind, "API ERROR")
+        return f"[{self.provider.upper()} {label}]\n\n{self.original}"
+
+
+def _classify_anthropic_error(exc: Exception) -> ProviderError:
+    """Map an anthropic SDK exception to a ProviderError."""
+    try:
+        import anthropic
+        if isinstance(exc, anthropic.RateLimitError):
+            return ProviderError("rate_limit", "anthropic", exc)
+        if isinstance(exc, anthropic.AuthenticationError):
+            return ProviderError("auth", "anthropic", exc)
+        if isinstance(exc, anthropic.PermissionDeniedError):
+            return ProviderError("auth", "anthropic", exc)
+        if isinstance(exc, anthropic.APIConnectionError):
+            return ProviderError("network", "anthropic", exc)
+        if isinstance(exc, anthropic.APIStatusError):
+            status = getattr(exc, "status_code", 0)
+            body = str(exc).lower()
+            if status == 429:
+                return ProviderError("rate_limit", "anthropic", exc)
+            if status in (401, 403):
+                return ProviderError("auth", "anthropic", exc)
+            if status == 402:
+                return ProviderError("balance", "anthropic", exc)
+            # Anthropic puts credit-balance errors in the body at 400
+            if "credit" in body or "balance" in body or "billing" in body:
+                return ProviderError("balance", "anthropic", exc)
+            if "quota" in body or "limit" in body or "exceeded" in body:
+                return ProviderError("quota", "anthropic", exc)
+    except ImportError:
+        pass
+    return ProviderError("unknown", "anthropic", exc)
+
+
+def _classify_gemini_error(exc: Exception) -> ProviderError:
+    """Map a google-genai SDK exception to a ProviderError."""
+    body = str(exc).lower()
+    # google-genai surfaces HTTP errors as google.api_core exceptions or
+    # google.genai exceptions; inspect the message text as a reliable fallback.
+    try:
+        from google.api_core import exceptions as gac
+        if isinstance(exc, gac.ResourceExhausted):
+            return ProviderError("quota", "gemini", exc)
+        if isinstance(exc, gac.TooManyRequests):
+            return ProviderError("rate_limit", "gemini", exc)
+        if isinstance(exc, (gac.Unauthenticated, gac.PermissionDenied)):
+            return ProviderError("auth", "gemini", exc)
+        if isinstance(exc, gac.ServiceUnavailable):
+            return ProviderError("network", "gemini", exc)
+    except ImportError:
+        pass
+    # Fallback: parse status code / message string
+    if "429" in body or "quota" in body or "resource exhausted" in body:
+        return ProviderError("quota", "gemini", exc)
+    if "rate" in body and "limit" in body:
+        return ProviderError("rate_limit", "gemini", exc)
+    if "401" in body or "403" in body or "api key" in body or "unauthenticated" in body:
+        return ProviderError("auth", "gemini", exc)
+    if "402" in body or "billing" in body or "balance" in body or "payment" in body:
+        return ProviderError("balance", "gemini", exc)
+    if "connection" in body or "timeout" in body or "unreachable" in body:
+        return ProviderError("network", "gemini", exc)
+    return ProviderError("unknown", "gemini", exc)
+
 # ------------------------------------------------------------------ provider setup

 def set_provider(provider: str, model: str):
@@ -49,24 +146,30 @@ def list_models(provider: str) -> list[str]:

 def _list_gemini_models(api_key: str) -> list[str]:
    from google import genai
-    client = genai.Client(api_key=api_key)
-    models = []
-    for m in client.models.list():
-        name = m.name
-        if name.startswith("models/"):
-            name = name[len("models/"):]
-        if "gemini" in name.lower():
-            models.append(name)
-    return sorted(models)
+    try:
+        client = genai.Client(api_key=api_key)
+        models = []
+        for m in client.models.list():
+            name = m.name
+            if name.startswith("models/"):
+                name = name[len("models/"):]
+            if "gemini" in name.lower():
+                models.append(name)
+        return sorted(models)
+    except Exception as exc:
+        raise _classify_gemini_error(exc) from exc

 def _list_anthropic_models() -> list[str]:
    import anthropic
-    creds = _load_credentials()
-    client = anthropic.Anthropic(api_key=creds["anthropic"]["api_key"])
-    models = []
-    for m in client.models.list():
-        models.append(m.id)
-    return sorted(models)
+    try:
+        creds = _load_credentials()
+        client = anthropic.Anthropic(api_key=creds["anthropic"]["api_key"])
+        models = []
+        for m in client.models.list():
+            models.append(m.id)
+        return sorted(models)
+    except Exception as exc:
+        raise _classify_anthropic_error(exc) from exc


 # --------------------------------------------------------- tool definition
@@ -148,59 +251,60 @@ def _send_gemini(md_content: str, user_message: str, base_dir: str) -> str:
    from google import genai
    from google.genai import types

-    _ensure_gemini_client()
+    try:
+        _ensure_gemini_client()

-    # Gemini chats don't support mutating tools after creation,
-    # so we recreate if None (reset_session clears it).
-    if _gemini_chat is None:
-        _gemini_chat = _gemini_client.chats.create(
-            model=_model,
-            config=types.GenerateContentConfig(
-                tools=[_gemini_tool_declaration()]
+        if _gemini_chat is None:
+            _gemini_chat = _gemini_client.chats.create(
+                model=_model,
+                config=types.GenerateContentConfig(
+                    tools=[_gemini_tool_declaration()]
+                )
            )
-        )

-    full_message = f"<context>\n{md_content}\n</context>\n\n{user_message}"
+        full_message = f"<context>\n{md_content}\n</context>\n\n{user_message}"

-    response = _gemini_chat.send_message(full_message)
+        response = _gemini_chat.send_message(full_message)

-    for _ in range(MAX_TOOL_ROUNDS):
-        # Collect all function calls in this response
-        tool_calls = [
-            part.function_call
+        for _ in range(MAX_TOOL_ROUNDS):
+            tool_calls = [
+                part.function_call
+                for candidate in response.candidates
+                for part in candidate.content.parts
+                if part.function_call is not None
+            ]
+            if not tool_calls:
+                break
+
+            function_responses = []
+            for fc in tool_calls:
+                if fc.name == TOOL_NAME:
+                    script = fc.args.get("script", "")
+                    output = _run_script(script, base_dir)
+                    function_responses.append(
+                        types.Part.from_function_response(
+                            name=TOOL_NAME,
+                            response={"output": output}
+                        )
+                    )
+
+            if not function_responses:
+                break
+
+            response = _gemini_chat.send_message(function_responses)
+
+        text_parts = [
+            part.text
            for candidate in response.candidates
            for part in candidate.content.parts
-            if part.function_call is not None
+            if hasattr(part, "text") and part.text
        ]
-        if not tool_calls:
-            break
+        return "\n".join(text_parts)

-        # Execute each tool call and collect results
-        function_responses = []
-        for fc in tool_calls:
-            if fc.name == TOOL_NAME:
-                script = fc.args.get("script", "")
-                output = _run_script(script, base_dir)
-                function_responses.append(
-                    types.Part.from_function_response(
-                        name=TOOL_NAME,
-                        response={"output": output}
-                    )
-                )
-
-        if not function_responses:
-            break
-
-        response = _gemini_chat.send_message(function_responses)
-
-    # Extract text from final response
-    text_parts = [
-        part.text
-        for candidate in response.candidates
-        for part in candidate.content.parts
-        if hasattr(part, "text") and part.text
-    ]
-    return "\n".join(text_parts)
+    except ProviderError:
+        raise
+    except Exception as exc:
+        raise _classify_gemini_error(exc) from exc

 # ------------------------------------------------------------------ anthropic

@@ -215,55 +319,58 @@ def _send_anthropic(md_content: str, user_message: str, base_dir: str) -> str:
    global _anthropic_history
    import anthropic

-    _ensure_anthropic_client()
+    try:
+        _ensure_anthropic_client()

-    full_message = f"<context>\n{md_content}\n</context>\n\n{user_message}"
-    _anthropic_history.append({"role": "user", "content": full_message})
+        full_message = f"<context>\n{md_content}\n</context>\n\n{user_message}"
+        _anthropic_history.append({"role": "user", "content": full_message})

-    for _ in range(MAX_TOOL_ROUNDS):
-        response = _anthropic_client.messages.create(
-            model=_model,
-            max_tokens=8096,
-            tools=_ANTHROPIC_TOOLS,
-            messages=_anthropic_history
-        )
+        for _ in range(MAX_TOOL_ROUNDS):
+            response = _anthropic_client.messages.create(
+                model=_model,
+                max_tokens=8096,
+                tools=_ANTHROPIC_TOOLS,
+                messages=_anthropic_history
+            )

-        # Always record the assistant turn
-        _anthropic_history.append({
-            "role": "assistant",
-            "content": response.content
-        })
+            _anthropic_history.append({
+                "role": "assistant",
+                "content": response.content
+            })

-        if response.stop_reason != "tool_use":
-            break
+            if response.stop_reason != "tool_use":
+                break

-        # Process tool calls
-        tool_results = []
-        for block in response.content:
-            if block.type == "tool_use" and block.name == TOOL_NAME:
-                script = block.input.get("script", "")
-                output = _run_script(script, base_dir)
-                tool_results.append({
-                    "type": "tool_result",
-                    "tool_use_id": block.id,
-                    "content": output
-                })
+            tool_results = []
+            for block in response.content:
+                if block.type == "tool_use" and block.name == TOOL_NAME:
+                    script = block.input.get("script", "")
+                    output = _run_script(script, base_dir)
+                    tool_results.append({
+                        "type": "tool_result",
+                        "tool_use_id": block.id,
+                        "content": output
+                    })

-        if not tool_results:
-            break
+            if not tool_results:
+                break

-        _anthropic_history.append({
-            "role": "user",
-            "content": tool_results
-        })
+            _anthropic_history.append({
+                "role": "user",
+                "content": tool_results
+            })

-    # Extract final text
-    text_parts = [
-        block.text
-        for block in response.content
-        if hasattr(block, "text") and block.text
-    ]
-    return "\n".join(text_parts)
+        text_parts = [
+            block.text
+            for block in response.content
+            if hasattr(block, "text") and block.text
+        ]
+        return "\n".join(text_parts)
+
+    except ProviderError:
+        raise
+    except Exception as exc:
+        raise _classify_anthropic_error(exc) from exc

 # ------------------------------------------------------------------ unified send

@@ -272,4 +379,4 @@ def send(md_content: str, user_message: str, base_dir: str = ".") -> str:
        return _send_gemini(md_content, user_message, base_dir)
    elif _provider == "anthropic":
        return _send_anthropic(md_content, user_message, base_dir)
-    raise ValueError(f"unknown provider: {_provider}")
+    raise ValueError(f"unknown provider: {_provider}")