[FEATURE] Make version-gated proto imports statically analyzable (`src/xai_sdk/proto/__init__.py`)

[dsfaccini]

This issue was posted by claude-code using claude-opus-4-8 on behalf of David (@dsfaccini), who reviewed the approach.

Confirmation of Request Source: This is a request about the xAI SDK Python library (its proto package typing), not the underlying xAI API.

Summary

src/xai_sdk/proto/__init__.py selects between the .v5 and .v6 proto packages at runtime based on the installed protobuf major version:

from packaging import version
import google.protobuf

if version.parse(google.protobuf.__version__).major == 5:
    from .v5 import (chat_pb2, sample_pb2, usage_pb2, ...)
elif version.parse(google.protobuf.__version__).major == 6:
    from .v6 import (chat_pb2, sample_pb2, usage_pb2, ...)
else:
    raise ValueError(...)

Static type checkers cannot evaluate this gate, so they infer each proto symbol as a v5 | v6 union. Because the SDK's own public API is annotated with these proto types (e.g. xai_sdk.chat.assistant(...) -> chat_pb2.Message), the union propagates through the entire SDK surface and produces spurious type errors in any downstream code that touches it.

Concretely, in pydantic-ai's models/xai.py, which does from xai_sdk.proto import chat_pb2, sample_pb2, usage_pb2, this generates ~45 false invalid-argument-type errors under Astral's ty — e.g.:

error[invalid-argument-type] Argument to `ResponseFormat.__init__` is incorrect:
  Expected `xai_sdk.proto.v6.chat_pb2.FormatType | str | None`,
  found `xai_sdk.proto.v5.chat_pb2.FormatType | xai_sdk.proto.v6.chat_pb2.FormatType`

The union is not fixable from the consumer side: pinning only the consumer's own chat_pb2 import to one arm makes it worse, because the SDK-returned unions then clash with the pinned arm (diagnostics went 49 → 75 when we pinned only the consumer import via TYPE_CHECKING).

Why the checker can't prune the gate

Two independent reasons, both rooted in the gate, not in the checker:

1. packaging.version.parse is opaque — it is annotated -> Version; nothing tells the checker which branch is taken.
2. The operand isn't a literal — google.protobuf.__version__ is typed str by the types-protobuf stub, so even literal-narrowing tricks (.startswith("6.")) don't fire.

So the checker keeps both branches live and unions their exports. (Note: pyright currently reports 0 errors only because it silently resolves the gate to the first arm, .v5, regardless of what's installed — internally consistent but arbitrary.)

Proposed fix

Wrap the gate so type checkers see a single, definite arm while runtime behavior is completely unchanged. The .v5 and .v6 *_pb2 stubs are structurally identical (verified below), so pinning .v6 for typing is correct, not a fudge:

from typing import TYPE_CHECKING

from packaging import version

import google.protobuf

if TYPE_CHECKING:
    # Type checkers cannot evaluate the runtime protobuf-version gate below, so
    # they infer a v5|v6 union across the whole SDK surface. The v5 and v6 pb2
    # modules are structurally identical (same class sets), so pinning one arm
    # for typing is correct; runtime behavior is unchanged by the elif/else.
    from .v6 import (
        auth_pb2,
        ...  # full name list, identical to the runtime branches
    )
elif version.parse(google.protobuf.__version__).major == 5:
    from .v5 import (...)
elif version.parse(google.protobuf.__version__).major == 6:
    from .v6 import (...)
else:
    raise ValueError(f"Unsupported protobuf version: {google.protobuf.__version__}")

The runtime path is the existing elif/else chain, untouched — under both protobuf 5 and protobuf 6 the same modules are imported at runtime as before. Only the checker's static view changes.

Evidence

Environment: xai-sdk 1.14.0, protobuf 6.33.5, types-protobuf 6.32.1; checked against pydantic-ai's real models/xai.py.

Astral ty on models/xai.py:

	diagnostics
before (current __init__.py)	49
after (patched as above)	4

All 45 removed diagnostics are the proto v5|v6 union cluster. The 4 residuals are pre-existing and unrelated to this issue (grpc.RpcError.code/details, an override signature, and an aclose attribute).

pyright on the same file: 0 errors both before and after. The difference is determinism: today pyright silently resolves the gate to the first arm (.v5); with the patch it resolves .v6 explicitly. No regression either way.

Runtime unchanged: with the patch applied on a protobuf-6 install, from xai_sdk.proto import chat_pb2 still resolves to xai_sdk.proto.v6.chat_pb2 at runtime, exactly as before.

v5/v6 structural identity

Verified that pinning .v6 for typing is honest — the top-level class sets are identical across every *_pb2.pyi: all 16 *_pb2.pyi modules have an identical set of top-level class names between .v5 and .v6 (a diff of ^class lines is empty for each). e.g. chat_pb2 = 51 classes in both, sample_pb2 = 4 in both, usage_pb2 = 3 in both.

PR feasibility

The change is mechanical and self-contained — a single file (src/xai_sdk/proto/__init__.py): add from typing import TYPE_CHECKING, and prepend an if TYPE_CHECKING: from .v6 import (...) block (the same name list already present in the runtime branches), converting the current if to elif. No other files change. I'm opening a PR with exactly this diff alongside this issue.