feat: xAI provider support added #1228
Conversation
📝 WalkthroughSummary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings. WalkthroughThis pull request adds support for a new XAI provider to the Bifrost system, enabling chat completions, text completions, responses, and streaming operations. The implementation follows OpenAI-compatible patterns and includes test configurations, schema enhancements for xAI-specific JSON formats, and provider registration updates. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20–30 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
🧪 Test Suite AvailableThis PR can be tested by a repository admin. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
core/schemas/bifrost.go (1)
34-54: Add XAI toStandardProvidersfor consistency with other built‑in providersYou’ve introduced
XAI ModelProvider = "xai"butStandardProviders(documented as “the list of all built‑in (non‑custom) providers”) doesn’t include it. That will cause any logic enumerating standard providers from this slice to omit XAI even though it’s now first‑class and has a factory case.Recommend appending
XAItoStandardProvidershere. You can defer adding it toSupportedBaseProvidersunless you explicitly want to allow XAI as a base for custom providers.Also applies to: 67-86
core/internal/testutil/account.go (1)
1082-1106: Validate XAI scenario matrix against actual Grok feature setThe new
AllProviderConfigsentry forschemas.XAIenables a fairly rich set of scenarios (chat, streaming, tools, multi‑image, image URL/base64, etc.) while disabling speech, transcription, and embeddings. This looks reasonable, but correctness depends on how the XAI provider implementation and Grok models behave in practice.I’d recommend double‑checking these flags against what
grok-4-0709(and any vision model you intend to use) actually supports so that CI doesn’t start failing due to unsupported tests, and adjust the scenario booleans accordingly if you hit gaps.core/internal/testutil/text_completion.go (1)
24-35: Capping text completion output viaMaxTokensis a good additionAdding
ParamswithMaxTokens: bifrost.Ptr(100)to the text completion request keeps test outputs bounded while following the existing parameter schema and pointer helper patterns. Looks good as‑is; consider extracting the100to a shared constant later if you want reuse across scenarios.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
core/bifrost.gocore/internal/testutil/account.gocore/internal/testutil/text_completion.gocore/providers/xai/xai.gocore/providers/xai/xai_test.gocore/schemas/bifrost.gocore/schemas/chatcompletions.gocore/schemas/responses.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
core/providers/xai/xai_test.gocore/schemas/responses.gocore/schemas/chatcompletions.gocore/bifrost.gocore/schemas/bifrost.gocore/internal/testutil/account.gocore/internal/testutil/text_completion.gocore/providers/xai/xai.go
🧠 Learnings (6)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
core/providers/xai/xai_test.gocore/schemas/responses.gocore/schemas/chatcompletions.gocore/bifrost.gocore/schemas/bifrost.gocore/internal/testutil/account.gocore/internal/testutil/text_completion.gocore/providers/xai/xai.go
📚 Learning: 2025-12-29T11:54:55.836Z
Learnt from: akshaydeo
Repo: maximhq/bifrost PR: 1153
File: framework/configstore/rdb.go:2221-2246
Timestamp: 2025-12-29T11:54:55.836Z
Learning: In Go reviews, do not flag range-over-int patterns like for i := range n as compile-time errors, assuming Go 1.22+ semantics. Only flag actual range-capable values (slices, arrays, maps, channels, strings) and other compile-time issues. This applies to all Go files across the repository.
Applied to files:
core/providers/xai/xai_test.gocore/schemas/responses.gocore/schemas/chatcompletions.gocore/bifrost.gocore/schemas/bifrost.gocore/internal/testutil/account.gocore/internal/testutil/text_completion.gocore/providers/xai/xai.go
📚 Learning: 2025-12-15T10:16:21.909Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/providers/huggingface/huggingface_test.go:12-63
Timestamp: 2025-12-15T10:16:21.909Z
Learning: In provider tests under core/providers/<provider>/*_test.go, do not require or flag the use of defer for Shutdown(); instead call client.Shutdown() at the end of each test function. This pattern appears consistent across all provider tests. Apply this rule only within this path; for other tests or resources, defer may still be appropriate.
Applied to files:
core/providers/xai/xai_test.go
📚 Learning: 2025-12-19T09:26:54.961Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/providers/utils/utils.go:1050-1051
Timestamp: 2025-12-19T09:26:54.961Z
Learning: Update streaming end-marker handling so HuggingFace is treated as a non-[DONE] provider for backends that do not emit a DONE marker (e.g., meta llama on novita). In core/providers/utils/utils.go, adjust ProviderSendsDoneMarker() (or related logic) to detect providers that may not emit DONE and avoid relying on DONE as the sole end signal. Add tests to cover both DONE-emitting and non-DONE backends, with clear documentation in code comments explaining the rationale and any fallback behavior.
Applied to files:
core/providers/xai/xai_test.gocore/providers/xai/xai.go
📚 Learning: 2025-12-17T08:44:08.788Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1114
File: core/schemas/chatcompletions.go:224-228
Timestamp: 2025-12-17T08:44:08.788Z
Learning: In core/schemas/chatcompletions.go, ensure the schema structures mirror OpenAI's API specifications exactly. Use the valid values for fields like ChatAudioParameters.Format and ChatAudioParameters.Voice as defined by OpenAI's documentation, and avoid adding additional inline documentation or constants to maintain direct compatibility with OpenAI's API.
Applied to files:
core/schemas/chatcompletions.go
📚 Learning: 2025-12-19T08:29:20.286Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1095
File: core/internal/testutil/count_tokens.go:30-67
Timestamp: 2025-12-19T08:29:20.286Z
Learning: In core/internal/testutil test files, enforce using GetTestRetryConfigForScenario() to obtain a generic retry config, then construct a typed retry config (e.g., CountTokensRetryConfig, EmbeddingRetryConfig, TranscriptionRetryConfig) with an empty Conditions slice. Copy only MaxAttempts, BaseDelay, MaxDelay, OnRetry, and OnFinalFail from the generic config. This convention should be consistently applied across all test files in this directory.
Applied to files:
core/internal/testutil/account.gocore/internal/testutil/text_completion.go
🧬 Code graph analysis (5)
core/providers/xai/xai_test.go (4)
core/internal/testutil/setup.go (1)
SetupTest(51-60)core/internal/testutil/account.go (2)
ComprehensiveTestConfig(61-83)TestScenarios(22-58)core/schemas/bifrost.go (1)
XAI(53-53)core/internal/testutil/tests.go (1)
RunAllComprehensiveTests(15-79)
core/bifrost.go (2)
core/schemas/bifrost.go (1)
XAI(53-53)core/providers/xai/xai.go (1)
NewXAIProvider(28-55)
core/schemas/bifrost.go (1)
ui/lib/types/config.ts (1)
ModelProvider(206-209)
core/internal/testutil/account.go (3)
core/schemas/bifrost.go (1)
XAI(53-53)core/schemas/provider.go (4)
ProviderConfig(271-280)NetworkConfig(48-56)ConcurrencyAndBufferSize(131-134)Provider(318-367)core/internal/testutil/cross_provider_scenarios.go (1)
ProviderConfig(45-53)
core/internal/testutil/text_completion.go (1)
core/schemas/textcompletions.go (1)
TextCompletionParameters(120-140)
🔇 Additional comments (12)
core/bifrost.go (1)
17-40: XAI provider wiring in core factory looks consistentImporting
core/providers/xaiand adding theschemas.XAIbranch increateBaseProvidermatches the pattern used for other providers that construct and return(Provider, error). No issues spotted with registration or error propagation here.Also applies to: 1931-1935
core/internal/testutil/account.go (1)
96-120: XAI test account wiring matches existing provider patternsThe additions in
GetConfiguredProviders,GetKeysForProvider, andGetConfigForProviderforschemas.XAIare consistent with how other providers are wired:
GetConfiguredProvidersnow includesschemas.XAI, so comprehensive tests will exercise this provider.GetKeysForProviderreadsXAI_API_KEYinto a single key, with emptyModelsandUseForBatchAPI: bifrost.Ptr(true)—matching the pattern used for other providers.GetConfigForProvidersets reasonable timeouts, retry, and concurrency settings in line with similar HTTP providers.No structural issues spotted here; just ensure
XAI_API_KEYis configured in environments where these tests run so they don’t fail due to missing keys.Also applies to: 123-362, 364-646
core/schemas/responses.go (1)
444-447: Remove or contextualize this comment—this is new schema code with intentional designThe
Annotationsfield is part of the new xAI provider introduction and does not represent a breaking change to existing consumers. The removal ofomitemptyappears intentional: the OpenAPI schema treatsannotationsas a standard response field (not optional), whilelogprobsretainsomitemptyto distinguish supplementary data. The codebase already uses defensive nil checks consistently (if t.Annotations != nil) throughout streaming and utility code, indicating this pattern is established practice.core/providers/xai/xai_test.go (1)
13-56: LGTM! Well-structured test implementation.The test follows established patterns for provider tests, properly gates execution behind the XAI_API_KEY environment variable, and calls
Shutdown()at the end withoutdefer(consistent with other provider tests per learnings). The comprehensive test configuration appropriately covers XAI's supported capabilities while correctly marking embedding as unsupported.core/schemas/chatcompletions.go (3)
286-310: LGTM! Handles xAI's JSON string format correctly.The custom unmarshaller properly supports both standard JSON object format and xAI's JSON string format (where parameters are encoded as a string). The fallback logic is sound and avoids recursion using the type alias pattern.
694-738: LGTM! Correctly maps xAI's reasoning_content field.The unmarshaller properly supports both OpenAI's
reasoningand xAI'sreasoning_contentfields, mapping the latter to the standardReasoningfield when present. The auxiliary struct pattern is correctly used, and the existing post-processing logic forReasoningDetailsis preserved.
845-885: LGTM! Consistent xAI support for streaming deltas.The streaming delta unmarshaller correctly mirrors the non-streaming implementation, supporting both OpenAI's
reasoningand xAI'sreasoning_contentfields. The auxiliary struct pattern is consistently applied, and the post-processing logic forReasoningDetailsis preserved.core/providers/xai/xai.go (5)
16-55: LGTM! Well-structured provider initialization.The constructor properly configures the HTTP client with reasonable defaults (timeouts, connection limits), handles proxy configuration, and sets the default base URL to
https://api.x.ai. The configuration follows established patterns from other providers.
133-159: LGTM! Proper authorization header handling for streaming.The method correctly constructs the Authorization header with the Bearer token when a key is provided and delegates to the shared OpenAI-compatible streaming handler. The explicit header construction is appropriate for streaming contexts.
161-198: LGTM! Consistent implementation for Responses API.Both streaming and non-streaming Responses methods follow the same pattern as the chat completion methods, with proper URL construction, authorization handling, and delegation to shared OpenAI-compatible handlers.
200-277: LGTM! Proper unsupported operation handling.All unsupported operations consistently return
NewUnsupportedOperationErrorwith the appropriate request type, clearly communicating which features are not available for the xAI provider.
97-115: The original review comment is based on a misunderstanding of the handler signatures.HandleOpenAITextCompletionStreamingexpects anauthHeaderparameter (map[string]string) as its 5th argument, whileHandleOpenAITextCompletionexpects akeyparameter (schemas.Key). These are different handlers with different signatures, so passingnilfor authHeader in TextCompletionStream is not inconsistent with passingkeyin TextCompletion. This pattern is intentional—providers like xai, ollama, and sgl passnilfor authHeader, while others like openai, openrouter, and azure construct an authHeader from the key and pass it. No issue exists here.Likely an incorrect or invalid review comment.
Summary
Added support for xAI's API (Grok) to Bifrost, enabling access to models like grok-4-0709 and grok-2-vision-1212.
Changes
XAIProviderthat leverages the OpenAI-compatible API structurereasoning_contentvsreasoning)ToolFunctionParametersto handle both JSON object and string formats for compatibilityType of change
Affected areas
How to test
Breaking changes
Related issues
N/A
Security considerations
The implementation follows the same security patterns as other providers, using API keys for authentication.
Checklist