Skip to content

feat: tier-aware build_model_id_name_mapping#312

Open
NBibikov wants to merge 1 commit intoHanaokaYuzu:masterfrom
NBibikov:feat/tier-aware-model-name-mapping
Open

feat: tier-aware build_model_id_name_mapping#312
NBibikov wants to merge 1 commit intoHanaokaYuzu:masterfrom
NBibikov:feat/tier-aware-model-name-mapping

Conversation

@NBibikov
Copy link
Copy Markdown

Problem

AvailableModel.build_model_id_name_mapping() unconditionally rewrites
every resolved model_id to its BASIC_* enum name, even when the
account is on PLUS or ADVANCED tier. The comment on the existing logic
states the intent, but it means:

  • A Plus account (capacity=4, capacity_field=12) sees its registered
    models as gemini-3-pro, gemini-3-flash, gemini-3-flash-thinking
    the Plus variants (*-plus) are never surfaced through list_models().
  • An Advanced account (capacity=2) has the same issue vs. *-advanced.
  • Callers that feed the registry's model_name back into
    generate_content(model=...) end up with a name/tier mismatch: the
    string resolves via _resolve_model_by_name_model_registry, which
    returns the AvailableModel with the account's real capacity header, but
    the name the caller saw suggests basic tier.

The PLUS and ADVANCED enum members also share model_id values with each
other (they differ only in the capacity encoded in the header), so the
right name cannot be inferred from model_id alone — it has to be chosen
against the caller's tier.

Fix

  • build_model_id_name_mapping now takes (capacity, capacity_field) and
    walks Model in a tier-priority order derived from them. Primary-tier
    ids register first; remaining tiers fill any ids the primary doesn't
    cover (e.g. BASIC_* ids that Plus/Advanced don't share).
  • Defaults keep the previous free-tier (capacity=1, capacity_field=12)
    behaviour so any external caller on the zero-arg form is unchanged.
  • _fetch_user_status forwards the capacity/capacity_field it already
    computes a few lines above.

Tests

New tests/test_available_model.py — pure unit tests, no network/cookies:

  • Each tier's primary ids map to that tier's canonical names (basic /
    -plus / -advanced).
  • BASIC_* unique ids still resolve on higher tiers (fallback works).
  • Zero-arg call still returns the legacy BASIC mapping.
$ python3 -m unittest tests.test_available_model -v
...
Ran 5 tests in 0.000s
OK

Backwards compatibility

The function signature gains optional parameters with defaults matching
the free tier, so external callers on the zero-arg form see no change.

PLUS and ADVANCED enum members share model_ids with each other; the name
stripped to BASIC regardless of tier meant Plus/Advanced accounts would
see basic names in the registry and build request headers that did not
match the name the caller expected.

Pass the account's (capacity, capacity_field) into the mapping so the
primary tier is chosen deliberately, with fallbacks covering any ids the
primary does not provide. Default args keep the legacy BASIC-only
behaviour for any external caller still using the zero-arg form.

Update the client to forward the capacity it has already computed.

Add unit tests covering all tier combinations; default-args parity
guards the legacy contract.
@luuquangvu
Copy link
Copy Markdown
Contributor

Don't worry about that. gemini-3-flash or gemini-3-pro are just labels, and the system automatically determines the membership type. This setup solves the issue of clients using this library to manage multiple membership types with a round-robin mechanism, removing the need to specify which is free or pro. It's like visiting a website and only seeing Flash, Thinking, or Pro, regardless of the actual membership type.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants