feat: Refactor message schema and enhance multi-modal support across AI providers#54
Conversation
Summary by CodeRabbit
Summary by CodeRabbit
WalkthroughThis change refactors the chat message schema and updates all provider implementations and related logic to use the new Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant HTTP Transport
participant Provider
participant Model API
Client->>HTTP Transport: Sends chat completion request (with BifrostMessage[])
HTTP Transport->>Provider: Calls ChatCompletion(ctx, model, key, messages: BifrostMessage[], params)
Provider->>Model API: Formats messages (role-based, images, tool calls) and sends request
Model API-->>Provider: Returns structured response (text, tool calls, thinking)
Provider-->>HTTP Transport: Returns BifrostResponse (BifrostMessage with assistant fields)
HTTP Transport-->>Client: Sends response
Suggested reviewers
Poem
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
This stack of pull requests is managed by Graphite. Learn more about stacking. |
danpiths
force-pushed
the
feature/06-03-reconcile-message-type-across-bifrost
branch
from
116433d to
d4f6e71
Compare
danpiths
changed the title
danpiths
force-pushed
the
feature/06-03-reconcile-message-type-across-bifrost
branch
from
d4f6e71 to
dd6f9a7
Compare
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
🔭 Outside diff range comments (4)
core/providers/azure.go (1)
338-346:⚠️ Potential issueSame pooling hazard for
response.Choicesslice
bifrostResponse.Choices = response.Choicescopies the slice header; its backing array
belongs to the pooledAzureChatResponsethat is released right after return.
Clone the slice (and any embedded pointers) before releasing:- bifrostResponse.Choices = response.Choices + bifrostResponse.Choices = append([]schemas.BifrostResponseChoice(nil), response.Choices...)core/providers/openai.go (1)
248-252:⚠️ Potential issue
content: nullwill be rejected by the OpenAI APIFor non-image messages you forward
msg.Contentdirectly. WhenContentisnil(e.g. an assistant message that only containstool_calls) JSON will serialise it asnull, but OpenAI’s spec requires a non-empty string.- "content": msg.Content, + "content": func() interface{} { + if msg.Content == nil { + return "" + } + return *msg.Content + }(),Failing to do this yields
400with “contentmust be a string”.core/providers/anthropic.go (2)
338-439: 🧹 Nitpick (assertive)Consider refactoring complex nested logic for better maintainability.
The current implementation has deeply nested conditions handling different message types and content. This makes the code harder to read and maintain.
Consider extracting the content building logic into separate helper functions:
func buildMessageContent(msg schemas.BifrostMessage) []map[string]interface{} { var content []map[string]interface{} // Handle image content if imageContent := extractImageContent(msg); imageContent != nil { content = append(content, imageContent) } // Handle text content if textContent := extractTextContent(msg); textContent != nil { content = append(content, textContent) } // Handle thinking content if thinkingContent := extractThinkingContent(msg); thinkingContent != nil { content = append(content, thinkingContent) } // Handle tool calls content = append(content, extractToolCallsContent(msg)...) return content }This would make the main function more readable and easier to test.
473-505: 🧹 Nitpick (assertive)Consider making text concatenation behavior configurable.
The current implementation joins multiple text blocks with newlines, which might not be desired in all use cases.
Consider adding a configuration option or parameter to control text concatenation behavior:
-if content.Len() > 0 { - content.WriteString("\n") -} -content.WriteString(c.Text) +if content.Len() > 0 && provider.config.JoinTextWithNewlines { + content.WriteString("\n") +} +content.WriteString(c.Text)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (11)
core/providers/anthropic.go(7 hunks)core/providers/azure.go(2 hunks)core/providers/bedrock.go(9 hunks)core/providers/cohere.go(4 hunks)core/providers/openai.go(3 hunks)core/providers/vertex.go(1 hunks)core/schemas/bifrost.go(4 hunks)core/schemas/provider.go(1 hunks)core/tests/tests.go(5 hunks)transports/bifrost-http/integrations/genai/types.go(2 hunks)transports/bifrost-http/main.go(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (9)
transports/bifrost-http/integrations/genai/types.go (1)
core/schemas/bifrost.go (2)
Function(101-105)FunctionCall(258-261)
core/schemas/provider.go (1)
core/schemas/bifrost.go (4)
BifrostMessage(143-150)ModelParameters(75-89)BifrostResponse(182-192)BifrostError(310-316)
core/providers/vertex.go (1)
core/schemas/bifrost.go (4)
BifrostMessage(143-150)ModelParameters(75-89)BifrostResponse(182-192)BifrostError(310-316)
core/tests/tests.go (3)
core/schemas/bifrost.go (6)
BifrostMessage(143-150)ModelChatMessageRoleUser(24-24)UserMessage(153-155)ImageContent(172-177)Anthropic(36-36)AssistantMessage(164-169)core/utils.go (1)
Ptr(3-5)core/schemas/provider.go (1)
Provider(142-149)
core/providers/azure.go (1)
core/schemas/bifrost.go (5)
BifrostMessage(143-150)ModelChatMessageRoleAssistant(23-23)ModelParameters(75-89)BifrostResponse(182-192)BifrostError(310-316)
core/providers/openai.go (1)
core/schemas/bifrost.go (5)
BifrostMessage(143-150)ModelParameters(75-89)UserMessage(153-155)ImageContent(172-177)ToolMessage(158-161)
core/providers/cohere.go (1)
core/schemas/bifrost.go (7)
BifrostMessage(143-150)ModelParameters(75-89)BifrostResponse(182-192)BifrostError(310-316)ModelChatMessageRoleChatbot(26-26)BifrostResponseChoice(287-293)AssistantMessage(164-169)
core/schemas/bifrost.go (1)
core/schemas/provider.go (1)
Provider(142-149)
core/providers/bedrock.go (1)
core/schemas/bifrost.go (13)
Function(101-105)FunctionCall(258-261)BifrostError(310-316)ErrorField(319-326)BifrostMessage(143-150)ModelChatMessageRoleAssistant(23-23)ModelChatMessageRoleSystem(25-25)UserMessage(153-155)ImageContent(172-177)ToolMessage(158-161)AssistantMessage(164-169)ModelParameters(75-89)BifrostResponse(182-192)
🔇 Additional comments (10)
transports/bifrost-http/main.go (1)
281-285: LGTM! Good defensive programming practice.The nil check for
bifrostCtxprevents potential runtime panics when calling client methods and provides a clear error response. This improves the robustness of the HTTP handler.transports/bifrost-http/integrations/genai/types.go (2)
56-59: LGTM! Appropriate defensive nil check.The defensive nil check prevents potential runtime panics during append operations. While the comment correctly notes that
ChatCompletionInputis initialized above, this check adds an extra layer of safety.
108-117: LGTM! Essential nil check for pointer field.The nil check on
toolCall.Function.Nameis necessary sinceNameis a pointer to string (as shown in theFunctionCallstruct fromcore/schemas/bifrost.golines 258-261) and could be nil. The logic correctly only creates and appends theFunctionCallwhen the name is available.core/providers/vertex.go (1)
92-92: LGTM! Provider implementation updated to match interface.The
ChatCompletionmethod signature correctly implements the updatedProviderinterface, using[]schemas.BifrostMessageinstead of the old[]schemas.Messagetype. This maintains consistency with the schema migration across all providers.core/providers/openai.go (1)
253-255: Nice touch – conditionaltool_call_idpropagationForwarding
tool_call_idonly when present keeps the payload minimal and avoids confusing the API.core/providers/bedrock.go (2)
418-424: TODO deserves tracking issueThere’s a
//TODO handling image inputs hereleft for system messages. System-level images are currently silently ignored which may surprise callers. Please either:
- Open a GitHub issue and reference it here, or
- Remove the TODO and document that images in system messages are unsupported.
Lingering TODOs reduce confidence in the production path.
503-508: 🧹 Nitpick (assertive)Reduce allocation for trivial text messages
Instead of building a slice and appending a single element every time, initialise it inline:
-message.Content = []BedrockMistralContent{ - {Text: *msg.Content}, -} +message.Content = []BedrockMistralContent{{Text: *msg.Content}}Minor, but this hot path benefits from micro-optimisation under high load.
Likely an incorrect or invalid review comment.
core/providers/anthropic.go (3)
257-261: LGTM! Correct migration to BifrostMessage.The text completion response correctly uses the new
schemas.BifrostMessagestructure with the updated role constantschemas.ModelChatMessageRoleAssistant.
280-280: Breaking change correctly implemented.The method signature has been updated to use
[]schemas.BifrostMessageas documented in the PR objectives. Consumers of this interface will need to update their code.
507-535: Excellent implementation of unified message structure.The refactored response parsing correctly aggregates all content types (text, thinking, tool calls) into a single assistant message, which aligns perfectly with the PR objectives for better message structure and multi-modal support.
Pratham-Mishra04
force-pushed
the
feature/06-03-reconcile-message-type-across-bifrost
branch
from
dd6f9a7 to
2ef8e49
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🔭 Outside diff range comments (1)
core/providers/bedrock.go (1)
96-100:⚠️ Potential issueIncorrect JSON tag – breaks Anthropic image payload
BedrockAnthropicImage.Formatis tagged asjson:"string"instead ofjson:"format", producing:{"string":"png", ...}Anthropic expects the key
format, leading to request rejection.-type BedrockAnthropicImage struct { - Format string `json:"string"` // Image format +type BedrockAnthropicImage struct { + Format string `json:"format"` // Image format
♻️ Duplicate comments (2)
core/schemas/provider.go (1)
148-148: LGTM! Core interface signature updated for schema migration.The
ChatCompletionmethod signature correctly migrates from[]Messageto[]BifrostMessage, aligning with the broader schema refactoring objectives. This breaking change enables enhanced multi-modal support and improved type safety.Based on the previous discussion, the remaining references to the old schema in dependent code will be updated after this core version is tagged, following Bifrost's established approach to managing breaking changes.
core/providers/bedrock.go (1)
431-465: Nil content still results in[null]payload (already flagged earlier)If a message has neither
Contentnor image data,contentremainsniland you send:"content": [null]Bedrock rejects this.
This exact issue was raised in a previous review; consider adding the early-continue guard suggested there.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (11)
core/providers/anthropic.go(7 hunks)core/providers/azure.go(2 hunks)core/providers/bedrock.go(9 hunks)core/providers/cohere.go(4 hunks)core/providers/openai.go(3 hunks)core/providers/vertex.go(1 hunks)core/schemas/bifrost.go(4 hunks)core/schemas/provider.go(1 hunks)core/tests/tests.go(5 hunks)transports/bifrost-http/integrations/genai/types.go(2 hunks)transports/bifrost-http/main.go(1 hunks)
🧰 Additional context used
🧠 Learnings (2)
core/providers/bedrock.go (1)
Learnt from: Pratham-Mishra04
PR: maximhq/bifrost#54
File: core/providers/bedrock.go:241-252
Timestamp: 2025-06-04T09:07:20.845Z
Learning: In the Bifrost codebase, when working with AWS Bedrock provider authentication, the preference is to let AWS handle access key validation naturally rather than adding preemptive checks for empty/blank access keys. This allows AWS to provide its own authentication error messages which can be more informative than custom validation errors.
core/providers/anthropic.go (1)
Learnt from: Pratham-Mishra04
PR: maximhq/bifrost#54
File: core/providers/anthropic.go:325-331
Timestamp: 2025-06-04T08:59:29.619Z
Learning: In the Anthropic provider, when handling system messages, prefer to let the API provider validate content (including empty strings) rather than pre-filtering on the client side. Users should receive errors directly from the provider if content is not supported.
🧬 Code Graph Analysis (9)
core/providers/vertex.go (1)
core/schemas/bifrost.go (4)
BifrostMessage(143-151)ModelParameters(75-89)BifrostResponse(183-193)BifrostError(311-317)
core/schemas/provider.go (1)
core/schemas/bifrost.go (4)
BifrostMessage(143-151)ModelParameters(75-89)BifrostResponse(183-193)BifrostError(311-317)
transports/bifrost-http/integrations/genai/types.go (1)
core/schemas/bifrost.go (2)
Function(101-105)FunctionCall(259-262)
core/tests/tests.go (2)
core/schemas/bifrost.go (5)
BifrostMessage(143-151)ModelChatMessageRoleUser(24-24)UserMessage(154-156)ImageContent(173-178)AssistantMessage(165-170)core/utils.go (1)
Ptr(3-5)
core/providers/azure.go (1)
core/schemas/bifrost.go (6)
BifrostResponseChoice(288-294)BifrostMessage(143-151)ModelChatMessageRoleAssistant(23-23)ModelParameters(75-89)BifrostResponse(183-193)BifrostError(311-317)
core/providers/cohere.go (1)
core/schemas/bifrost.go (7)
BifrostMessage(143-151)ModelParameters(75-89)BifrostResponse(183-193)BifrostError(311-317)ModelChatMessageRoleChatbot(26-26)BifrostResponseChoice(288-294)AssistantMessage(165-170)
core/providers/bedrock.go (1)
core/schemas/bifrost.go (11)
Function(101-105)FunctionCall(259-262)BifrostError(311-317)ErrorField(320-327)BifrostMessage(143-151)ModelChatMessageRoleAssistant(23-23)ModelChatMessageRoleSystem(25-25)UserMessage(154-156)ImageContent(173-178)ToolMessage(159-162)AssistantMessage(165-170)
core/providers/openai.go (1)
core/schemas/bifrost.go (7)
BifrostMessage(143-151)ModelParameters(75-89)BifrostResponse(183-193)BifrostError(311-317)UserMessage(154-156)ImageContent(173-178)ToolMessage(159-162)
core/providers/anthropic.go (2)
core/schemas/bifrost.go (12)
BifrostMessage(143-151)ModelChatMessageRoleAssistant(23-23)ModelParameters(75-89)BifrostResponse(183-193)ModelChatMessageRoleSystem(25-25)UserMessage(154-156)ImageContent(173-178)ToolMessage(159-162)AssistantMessage(165-170)Function(101-105)ToolCall(265-269)BifrostResponseChoice(288-294)core/providers/bedrock.go (1)
BedrockAnthropicSystemMessage(67-69)
🔇 Additional comments (23)
transports/bifrost-http/main.go (1)
281-285: LGTM! Defensive nil check prevents potential crashes.The nil check for
bifrostCtxafter conversion is a good defensive programming practice that prevents potential nil pointer dereferences in subsequent client method calls. The error response with HTTP 500 status code appropriately indicates an internal server error.transports/bifrost-http/integrations/genai/types.go (2)
56-59: LGTM! Defensive nil check prevents potential issues.The nil check before appending to
ChatCompletionInputis a good defensive programming practice, even though the slice is initialized above. This ensures robustness against potential future changes in initialization logic.
108-117: LGTM! Essential nil check prevents dereferencing nil pointer.The nil check for
toolCall.Function.Nameis crucial sinceFunctionCall.Nameis defined as*stringin the schema (core/schemas/bifrost.go:259-262). Moving the entireFunctionCallcreation and appending logic inside the conditional block correctly prevents nil pointer dereference whenNameis nil.core/providers/vertex.go (1)
92-92: LGTM! Provider implementation updated to match interface.The
ChatCompletionmethod signature correctly implements the updatedProviderinterface by changing from[]schemas.Messageto[]schemas.BifrostMessage. This aligns the Vertex provider with the broader schema migration and maintains consistency with other provider implementations.core/tests/tests.go (5)
146-151: Schema migration properly implemented in chat completion tests.The change from
[]schemas.Messageto[]schemas.BifrostMessageand the role constant update fromschemas.RoleUsertoschemas.ModelChatMessageRoleUserare consistent with the new unified message schema.
185-196: Image content properly nested in UserMessage structure.The refactoring correctly moves image content into the nested
UserMessage.ImageContentstructure, which aligns with the new modular message design. The use ofbifrost.Ptr()for the type field is appropriate for creating pointer values.
219-231: Base64 image test properly structured.The base64 image message structure follows the same nested pattern as the URL image test, maintaining consistency in the new schema design.
282-287: Tool call test messages updated correctly.The message structure for tool call tests properly uses the new
BifrostMessagetype and role constant.
300-300: Essential null safety check added for AssistantMessage.The additional null check for
AssistantMessagebefore accessingToolCallsprevents potential null pointer dereferences when the embedded message type is nil.core/providers/azure.go (3)
266-273: Dangling pointer issue correctly resolved.The fix creates a local copy of the text content before assigning its pointer to the
Contentfield, preventing the dangling pointer issue that would occur when the pooled response object is released. This addresses the critical issue identified in previous reviews.
299-299: Method signature updated for schema consistency.The
ChatCompletionmethod signature change from[]schemas.Messageto[]schemas.BifrostMessagealigns with the unified message schema across all providers.
305-315: Message formatting enhanced with proper nil handling.The conditional inclusion of the
"content"field only whenmsg.Contentis non-nil prevents potential issues with nil pointer dereferences and ensures clean JSON serialization. The dereferencing of the content pointer is handled safely.core/providers/cohere.go (5)
144-144: Method signature updated for schema consistency.The
ChatCompletionmethod signature properly migrated to use[]schemas.BifrostMessage, maintaining consistency with other providers.
152-162: Chat history formatting enhanced with nil safety.The conditional inclusion of message content only when non-nil prevents potential null pointer dereferences and ensures robust message formatting for the Cohere API.
173-176: Last message content handling improved.The nil check for the last message content before adding it to the request body prevents potential issues and aligns with the defensive programming approach used throughout the method.
306-312: Response message structure properly updated.The response construction now uses the new
BifrostMessagestructure with embeddedAssistantMessagecontaining tool calls, which correctly implements the modular message design.
341-372: Helper function updated for new return type.The
convertChatHistoryfunction correctly returns*[]schemas.BifrostMessageand properly constructs messages with embeddedAssistantMessagefields, maintaining consistency with the new schema.core/providers/openai.go (5)
117-117: Method signature updated for schema consistency.The
ChatCompletionmethod signature change to use[]schemas.BifrostMessagemaintains consistency with the unified message schema across all providers.
208-208: Helper function signature updated appropriately.The
prepareOpenAIChatRequestfunction signature update to accept[]schemas.BifrostMessagealigns with the schema migration.
212-218: Image content extraction adapted for nested structure.The logic correctly checks for image content in both
UserMessageandToolMessageembedded structures and extracts the content appropriately. The conditional logic handles the new nested message design properly.
226-226: Content pointer properly dereferenced.The dereferencing of
msg.Contentwhen adding text content to the formatted message is correct and necessary for the API payload.
253-255: Tool call ID handling enhanced with nil safety.The nested nil checks for
msg.ToolMessageandmsg.ToolMessage.ToolCallIDprevent potential null pointer dereferences and ensure safe access to the tool call ID field in the new message structure.core/schemas/bifrost.go (1)
22-28: Clarify the need for thechatbotrole constant
ModelChatMessageRoleChatbotisn’t referenced elsewhere in the diff and is not part of the widely-used OpenAI / Anthropic / Mistral role vocabularies.
Confirm it is really needed; otherwise it’s dead weight and a future maintenance trap.
Merge activity
|
Pratham-Mishra04
deleted the
feature/06-03-reconcile-message-type-across-bifrost
branch
…AI providers (#54) ## Description This PR introduces a comprehensive refactoring of the message handling system across all AI providers (Anthropic, Azure, and Bedrock) to support a more robust and extensible message schema. ### Key Changes **🔄 Message Schema Standardization** - Migrated from `schemas.Message` to `schemas.BifrostMessage` across all providers - Updated role constants from `schemas.RoleAssistant/RoleSystem` to `schemas.ModelChatMessageRoleAssistant/ModelChatMessageRoleSystem` - Replaced `schemas.BifrostResponseChoiceMessage` with unified `schemas.BifrostMessage` **🎯 Enhanced Message Structure** - Introduced structured message types: `AssistantMessage`, `UserMessage`, `ToolMessage` - Added support for specialized content types within each message role - Improved type safety and clarity in message handling **🖼️ Multi-Modal Content Support** - Enhanced image content handling across all providers - Better separation of concerns for different content types (text, images, tool calls) - Improved support for mixed content messages **🧠 Anthropic Provider Enhancements** - Added support for "thinking" content in Claude responses - Improved tool call parsing and formatting - Enhanced content block structure handling - Better error handling for malformed JSON in tool arguments **🔧 Bug Fixes & Improvements** - Added null check for secret access key in Bedrock provider - Fixed content handling edge cases in Azure provider - Improved error messages and validation ### Breaking Changes - `ChatCompletion` method signature updated to use `[]schemas.BifrostMessage` instead of `[]schemas.Message` - Message role constants have been renamed and moved - Response message structure has been updated ### Migration Impact This change affects all consumers of the provider interfaces and will require corresponding updates to maintain compatibility.
3 participants
Description
This PR introduces a comprehensive refactoring of the message handling system across all AI providers (Anthropic, Azure, and Bedrock) to support a more robust and extensible message schema.
Key Changes
🔄 Message Schema Standardization
schemas.Messagetoschemas.BifrostMessageacross all providersschemas.RoleAssistant/RoleSystemtoschemas.ModelChatMessageRoleAssistant/ModelChatMessageRoleSystemschemas.BifrostResponseChoiceMessagewith unifiedschemas.BifrostMessage🎯 Enhanced Message Structure
AssistantMessage,UserMessage,ToolMessage🖼️ Multi-Modal Content Support
🧠 Anthropic Provider Enhancements
🔧 Bug Fixes & Improvements
Breaking Changes
ChatCompletionmethod signature updated to use[]schemas.BifrostMessageinstead of[]schemas.MessageMigration Impact
This change affects all consumers of the provider interfaces and will require corresponding updates to maintain compatibility.