feat: replicate provider #1434
feat: replicate provider #1434
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a full Replicate provider: provider implementation, types, request/response converters for chat/text/responses/images/files, utilities and error handling, DB schema and configstore support for Replicate key settings, UI/docs/CI updates, extensive tests, and transport handlers for image URL/file inputs. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant BifrostCore as Bifrost Core
participant Provider as ReplicateProvider
participant ReplicateAPI as Replicate API
participant DB as Database
Client->>BifrostCore: Send request (chat/text/image/file)
BifrostCore->>Provider: Route operation (key, request)
Provider->>Provider: Translate to Replicate format (version/deployment, inputs, extra_params)
Provider->>ReplicateAPI: POST /predictions or GET stream (SSE)
ReplicateAPI-->>Provider: Immediate response or SSE events
alt Async prediction
Provider->>ReplicateAPI: Poll or await webhook
ReplicateAPI-->>Provider: Status updates (processing → succeeded)
end
Provider->>Provider: Convert Replicate response → Bifrost response (usage, status, outputs)
alt Persist metadata
Provider->>DB: Store metrics/metadata, file signatures
end
Provider-->>BifrostCore: Return response/stream
BifrostCore-->>Client: Deliver response/stream
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
transports/bifrost-http/handlers/inference.go (1)
44-62: Preserve legacytextfield for text completions.Switching the known field to
promptmakes older clients sendingtextfail validation. Add a compatibility alias (like themax_tokensmapping in chat) to avoid a breaking change.🔧 Suggested compatibility patch
var textParamsKnownFields = map[string]bool{ "model": true, "prompt": true, + "text": true, // legacy alias "fallbacks": true, "best_of": true, "echo": true, @@ - if req.Prompt == nil || (req.Prompt.PromptStr == nil && req.Prompt.PromptArray == nil) { - SendError(ctx, fasthttp.StatusBadRequest, "prompt is required for text completion") - return - } + if req.Prompt == nil || (req.Prompt.PromptStr == nil && req.Prompt.PromptArray == nil) { + var legacy struct { + Text *schemas.TextCompletionInput `json:"text"` + } + if err := sonic.Unmarshal(ctx.PostBody(), &legacy); err == nil && legacy.Text != nil { + req.Prompt = legacy.Text + } + } + if req.Prompt == nil || (req.Prompt.PromptStr == nil && req.Prompt.PromptArray == nil) { + SendError(ctx, fasthttp.StatusBadRequest, "prompt is required for text completion") + return + }Also applies to: 617-620
🤖 Fix all issues with AI agents
In `@core/providers/openai/openai.go`:
- Around line 2590-2591: Remove the stdout dump that prints the full JSON
request body (the fmt.Println("jsonBody", string(jsonBody)) call) to avoid
leaking sensitive prompts/PII; either delete that line or replace it with a
structured logger call that only logs redacted fields and is gated behind a
debug flag (e.g., use the existing package logger with logger.Debug and redact
or omit prompt/content before logging), ensuring you operate on the jsonBody
variable where it is constructed in openai.go.
In `@core/providers/replicate/models.go`:
- Around line 69-73: The code builds Replicate URLs and names from model.ID
which includes a provider prefix (e.g., "replicate/owner/name"), producing
invalid paths; update the loop that constructs ReplicateModelResponse (iterating
response.Data and using model.ID) to strip the provider prefix before setting
Name and URL: split model.ID on "/" (strings.SplitN with max 3), if it has three
segments use segments[1]+"/"+segments[2] as the path and
"https://replicate.com/"+thatPath for URL, otherwise fall back to model.ID; keep
using the ReplicateModelResponse struct to assign the corrected Name and URL.
In `@core/providers/replicate/replicate.go`:
- Around line 1374-1379: Remove the stdout debug prints that expose SSE
payloads: delete the fmt.Println calls around sendBackRawResponse and Raw
response printing in the block that sets chunk.ExtraFields.RawResponse; keep the
existing behavior that, when sendBackRawResponse is true, marshals currentEvent
with sonic.Marshal and assigns the JSON string to chunk.ExtraFields.RawResponse,
or if you prefer retain a log use the package's structured logger at a
non-production debug level (not fmt.Println) and ensure it never logs
currentEvent.Data or any raw base64 image data.
- Around line 1191-1206: In ImageGenerationStream, guard the streaming
conversion by validating the incoming request and the produced Replicate input
before setting replicateReq.Stream: inside the
providerUtils.CheckContextAndGetRequestBody callback (and/or before calling
ToReplicateImageGenerationInput) check that request is non-nil and that
ToReplicateImageGenerationInput(request) returns a non-nil value, and if either
is nil return a proper error instead of proceeding; ensure you still set
replicateReq.Stream = schemas.Ptr(true) only after the non-nil check so you
avoid any nil dereference when creating the JSON body.
- Around line 1117-1127: The ImageGeneration method may call
ToReplicateImageGenerationInput with a nil request or nil request.Input and
panic; before invoking providerUtils.CheckContextAndGetRequestBody (or inside
the anonymous conversion passed to it), add a guard that verifies request != nil
and request.Input != nil and return an appropriate *schemas.BifrostError when
they are nil; reference the ImageGeneration method on ReplicateProvider and the
ToReplicateImageGenerationInput conversion to locate where to perform the nil
checks and short-circuit with a clear BifrostError instead of allowing a panic.
In `@core/providers/replicate/types.go`:
- Around line 96-157: The UnmarshalJSON method for
ReplicatePredictionRequestInput is missing "video" in the knownFields map,
causing the video field to be treated as ExtraParams; update the knownFields map
inside ReplicatePredictionRequestInput.UnmarshalJSON to include the "video" key
(e.g., "video": true) so the video field is recognized as a first-class field
and not copied into ExtraParams.
- Around line 199-245: UnmarshalJSON for ReplicateOutput must reset both fields
at start to avoid leaving stale values (which can cause MarshalJSON to see both
set); update ReplicateOutput.UnmarshalJSON to zero-out mc.OutputStr and
mc.OutputArray before attempting any decoding, and when you set the string field
use the pointer helper (e.g., schemas.Ptr) instead of &value for consistency;
ensure on successful string decode you assign mc.OutputStr =
schemas.Ptr(stringValue) and mc.OutputArray = nil, and on successful array
decode you assign mc.OutputArray = arrayValue and mc.OutputStr = nil, returning
an error only if neither decode succeeds.
In `@core/schemas/files.go`:
- Around line 58-62: The uploaded file handling currently doesn't set the new
File.ContentType field; update the multipart file upload handler (the code that
reads multipart.File and multipart.FileHeader when building the request) to
extract the content type from the file header (use
fileHeader.Header.Get("Content-Type")), and assign it to the File struct as a
pointer to the string (e.g., set requestFile.ContentType = &contentType if
non-empty); ensure this is done when constructing the File slice/instance
alongside existing fields like File, Filename, and Purpose so downstream
consumers receive the MIME type.
🧹 Nitpick comments (10)
core/providers/replicate/utils.go (1)
24-46: Avoid hardcoding the provider key in error helpers.
Hardcodingschemas.Replicatehere means custom provider aliases won’t surface in error metadata; consider threading the actual provider key into this helper (and similar ones) soExtraFields.Providermatches configuration. Based on learnings, preferprovider.GetProviderKey()(or a plumbed provider name) over provider constants in error metadata.core/providers/replicate/errors.go (1)
8-35: Prefer pointer helpers for StatusCode.
Keeps pointer creation consistent with repo conventions.♻️ Suggested change
- StatusCode: &statusCode, + StatusCode: schemas.Ptr(statusCode), @@ - StatusCode: &statusCode, + StatusCode: schemas.Ptr(statusCode),Based on learnings, prefer pointer helpers (e.g.,
schemas.Ptr) over&valuefor consistency.core/providers/replicate/files.go (1)
39-45: Prefer pointer helpers for ExpiresAt.
Keeps pointer creation consistent and avoids raw address operators.♻️ Suggested change
- resp.ExpiresAt = &expiresAt + resp.ExpiresAt = schemas.Ptr(expiresAt) @@ - resp.ExpiresAt = &expiresAt + resp.ExpiresAt = schemas.Ptr(expiresAt)Based on learnings, prefer pointer helpers (e.g.,
schemas.Ptr) over&valuefor consistency.Also applies to: 76-82
core/internal/testutil/tests.go (1)
60-60: Document why the pagination test is disabled.
If this is a temporary stack-specific change, please add a short comment/TODO or a config guard so the coverage drop is intentional and traceable. As per coding guidelines, please confirm this aligns with the rest of the PR stack.core/internal/testutil/list_models.go (1)
111-232: Avoid leaving the pagination test commented-out long-term.A fully commented block tends to drift and hides intent. If pagination is intentionally disabled for the current stack, consider removing it (git history keeps it) or gating via a scenario flag/build tag with a short rationale.
core/providers/replicate/images.go (1)
13-77: Usebifrost.Ptrfor pointer fields to match repo conventions.This keeps pointer creation consistent across the codebase and test utilities. Based on learnings, please use
bifrost.Ptr(...)instead of&value.♻️ Suggested refactor
-import ( - schemas "github.com/maximhq/bifrost/core/schemas" -) +import ( + bifrost "github.com/maximhq/bifrost/core" + schemas "github.com/maximhq/bifrost/core/schemas" +) @@ - input := &ReplicatePredictionRequestInput{ - Prompt: &bifrostReq.Input.Prompt, - } + input := &ReplicatePredictionRequestInput{ + Prompt: bifrost.Ptr(bifrostReq.Input.Prompt), + } @@ if isVersionID(bifrostReq.Model) { - request.Version = &bifrostReq.Model + request.Version = bifrost.Ptr(bifrostReq.Model) }core/internal/testutil/account.go (1)
1201-1227: Consider adding environment variable overrides for Replicate model defaults.The model IDs (
openai/gpt-4.1-mini,black-forest-labs/flux-dev) are valid and publicly available on Replicate. However, to maintain consistency with other providers in the test suite (AWS, Vertex, Ollama, etc.) and provide flexibility for CI environments, consider exposing these as environment variable overrides using the existinggetEnvWithDefault()pattern:ChatModel: getEnvWithDefault("REPLICATE_CHAT_MODEL", "openai/gpt-4.1-mini"), TextModel: getEnvWithDefault("REPLICATE_TEXT_MODEL", "openai/gpt-4.1-mini"), ImageGenerationModel: getEnvWithDefault("REPLICATE_IMAGE_MODEL", "black-forest-labs/flux-dev"),core/providers/replicate/text.go (1)
103-112: Prefer pointer helper for joined output string.
Keeps pointer creation consistent with repo style.♻️ Suggested change
- } else if response.Output.OutputArray != nil { - // Join array of strings into a single string - joined := strings.Join(response.Output.OutputArray, "") - textOutput = &joined - } + } else if response.Output.OutputArray != nil { + // Join array of strings into a single string + joined := strings.Join(response.Output.OutputArray, "") + textOutput = schemas.Ptr(joined) + }Based on learnings, prefer using the pointer helper (e.g., schemas.Ptr/bifrost.Ptr) instead of taking addresses directly.
core/providers/replicate/chat.go (1)
67-76: Use pointer helper for prompt fields.
Aligns pointer creation with repo conventions.♻️ Suggested change
- if systemPrompt != "" { - input.SystemPrompt = &systemPrompt - } + if systemPrompt != "" { + input.SystemPrompt = schemas.Ptr(systemPrompt) + } @@ - if len(conversationParts) > 0 { - prompt := strings.Join(conversationParts, "\n") - input.Prompt = &prompt - } + if len(conversationParts) > 0 { + prompt := strings.Join(conversationParts, "\n") + input.Prompt = schemas.Ptr(prompt) + }Based on learnings, prefer using the pointer helper for string pointers instead of address-of locals.
core/providers/replicate/replicate.go (1)
993-1004: Use pointer helper for finishReason.
Keeps pointer creation consistent with repo style.♻️ Suggested change
- finishReason := "stop" + finishReason := schemas.Ptr("stop") finalResponse := &schemas.BifrostChatResponse{ @@ - FinishReason: &finishReason, + FinishReason: finishReason,Based on learnings, prefer using the pointer helper for local string pointers; apply similarly in this file.
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
0cc4f3e to
f7cbe32
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@core/providers/replicate/models.go`:
- Around line 20-51: The loop over deploymentsResponse.Results takes the address
of the loop variable fields (e.g., &deployment.Name and &deployment.Owner) and
stores them into schemas.Model so all entries alias the last iteration; fix by
creating heap-allocated copies using schemas.Ptr when assigning Name, Deployment
and OwnedBy in the schemas.Model initializer (replace direct &deployment.Name
and &deployment.Owner with schemas.Ptr(deployment.Name) /
schemas.Ptr(deployment.Owner)); leave ParseReplicateTimestamp and local desc
handling as-is (or optionally use schemas.Ptr for consistency) and ensure the ID
still uses providerKey + "/" + deploymentID.
In `@core/providers/replicate/replicate.go`:
- Around line 405-423: After polling (or immediately if sync) the code currently
converts a terminal-failed/canceled Replicate prediction into a success Bifrost
response; add a terminal-status check after the optional poll to return a
Bifrost error instead. Specifically, after the block that may call
pollPrediction (and before calling prediction.ToBifrostTextCompletionResponse),
call isTerminalStatus(prediction.Status) and if true return an enriched provider
error (use providerUtils.EnrichError(ctx, err, jsonData, nil,
provider.sendBackRawRequest, provider.sendBackRawResponse) or construct an
appropriate error describing the terminal status), mirroring the same
check/failure-return pattern used elsewhere in this file for chat/image paths
(apply same fix around the other occurrences noted).
♻️ Duplicate comments (3)
core/providers/replicate/types.go (2)
113-140: IncludevideoinknownFieldsto prevent misclassification.🐛 Suggested fix
knownFields := map[string]bool{ "prompt": true, "messages": true, "system_prompt": true, "image": true, "images": true, "number_of_images": true, "quality": true, "background": true, "audio": true, + "video": true, "seed": true, "width": true, "height": true,
221-240: ResetReplicateOutputfields at unmarshal start to avoid stale data.🐛 Suggested fix
func (mc *ReplicateOutput) UnmarshalJSON(data []byte) error { + mc.OutputStr = nil + mc.OutputArray = nil if string(data) == "null" || len(data) == 0 { - mc.OutputStr = nil - mc.OutputArray = nil return nil }core/providers/replicate/replicate.go (1)
1374-1377: Remove stray stdout debug print.🧹 Suggested cleanup
- fmt.Println("sendbackrawresponse", sendBackRawResponse)
🧹 Nitpick comments (3)
core/internal/testutil/tests.go (1)
59-61: Consider gating the pagination test instead of removing it.At Line 60,
RunListModelsPaginationTestis commented out, which drops pagination coverage globally. If pagination is still supported in the stack, consider gating it by provider/scenario or replacing it with a targeted test so coverage isn’t silently lost. As per coding guidelines, please ensure this aligns with the rest of the stack.core/internal/testutil/account.go (1)
1201-1227: Consider exercising Replicate text completion + image stream flows in tests.The PR adds these capabilities, so enabling their scenarios would validate the new paths. If those flows are intentionally out of scope, feel free to ignore.
Suggested test coverage tweak
Scenarios: TestScenarios{ - TextCompletion: false, // Not typical + TextCompletion: true, + TextCompletionStream: true, SimpleChat: true, CompletionStream: true, MultiTurnConversation: true, @@ - ImageGenerationStream: false, + ImageGenerationStream: true, },core/providers/replicate/files.go (1)
39-45: Use the pointer helper forExpiresAtto match repo conventions.♻️ Suggested change
- if expiresAt > 0 { - resp.ExpiresAt = &expiresAt - } + if expiresAt > 0 { + resp.ExpiresAt = schemas.Ptr(expiresAt) + }- if expiresAt > 0 { - resp.ExpiresAt = &expiresAt - } + if expiresAt > 0 { + resp.ExpiresAt = schemas.Ptr(expiresAt) + }Based on learnings, prefer the pointer helper over
&valuefor consistency.Also applies to: 76-82
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
2 times, most recently
from
c8c219b to
e72bf1f
Compare
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
ui/lib/constants/config.ts (1)
98-130: Add Replicate to BaseProvider type, or update PROVIDER_SUPPORTED_REQUESTS type annotation.
PROVIDER_SUPPORTED_REQUESTSis typed asRecord<BaseProvider, string[]>, butBaseProvideris defined as a union of only 5 providers:"openai" | "anthropic" | "cohere" | "gemini" | "bedrock". Replicate is aKnownProvider(referenced inisKeyRequiredByProviderandKnownProvidersNames) but not included inBaseProvider.Since the PR adds Replicate support, either:
- Extend
BaseProviderinui/lib/types/config.tsto include"replicate", or- Change
PROVIDER_SUPPORTED_REQUESTSto useKnownProviderinstead ofBaseProviderif supporting all known providers is the intent.Currently, the type constraint prevents adding a Replicate entry to the mapping.
🤖 Fix all issues with AI agents
In `@core/providers/replicate/images.go`:
- Around line 143-157: When populating the ImageUsage from prediction.Metrics
(in images.go), ensure TotalTokens is computed even if only one of
InputTokenCount or OutputTokenCount is set: create usage.InputTokens and
usage.OutputTokens from prediction.Metrics (using zero if nil) and then set
usage.TotalTokens = usage.InputTokens + usage.OutputTokens so TotalTokens is
always the sum of the two (not left as zero when only one side exists);
reference schemas.ImageUsage, prediction.Metrics.InputTokenCount,
prediction.Metrics.OutputTokenCount, and the usage.TotalTokens assignment to
locate and update the code.
- Around line 24-44: The mapping defaults to "input_images" for 64-char version
IDs because getInputImageFieldName can't match model names; fix by checking
ExtraParams for an explicit override before using getInputImageFieldName. In the
block handling params.InputImages, read a string override (e.g.
ExtraParams["input_image_field"]) from params.ExtraParams (or
bifrostReq.ExtraParams if that's the container) and use that value as fieldName
when non-empty and a string; otherwise fall back to
getInputImageFieldName(bifrostReq.Model). Then proceed with the existing switch
that sets input.ImagePrompt, input.InputImage, input.Image, or input.InputImages
so Flux variants passed as version IDs honor the override.
In `@core/providers/replicate/models.go`:
- Around line 20-45: The code only sets model Name/Deployment when
deployment.CurrentRelease is non-nil, dropping metadata otherwise; change the
logic in the loop over deploymentsResponse.Results so modelName defaults to
deployment.Name (e.g., set modelName = schemas.Ptr(deployment.Name) before the
CurrentRelease check) and only override or set Created when
deployment.CurrentRelease exists and ParseReplicateTimestamp returns >0; ensure
schemas.Model{ Name, Deployment } use that modelName and Created remains
conditional via ParseReplicateTimestamp.
- Around line 65-78: The code currently sets ReplicateModelResponse.Name to
modelID which contains "owner/name"; instead split modelID into owner and name
(e.g., parts := strings.SplitN(modelID, "/", 2)) and set replicateModel.Owner to
the owner part and replicateModel.Name to the name part; if model.OwnedBy is
present you may prefer to set replicateModel.Owner = *model.OwnedBy but still
parse modelID to populate replicateModel.Name from the second segment so Owner
and Name are stored separately (update the code around modelID,
replicateModel.Owner, and replicateModel.Name).
In `@core/providers/replicate/types.go`:
- Around line 118-148: The knownFields map in core/providers/replicate/types.go
is missing the first-class keys reasoning_effort and max_completion_tokens,
causing them to be treated as ExtraParams; update the knownFields map (the
map[string]bool literal named knownFields) to include "reasoning_effort": true
and "max_completion_tokens": true so these fields are recognized and not copied
into ExtraParams during (re)marshalling.
In `@ui/lib/constants/config.ts`:
- Line 52: Update the placeholder value for the replicate config entry to show
the correct Replicate model identifier format (owner/model-name) instead of the
incorrect "replicate/..." prefix; specifically edit the constant keyed as
replicate in ui/lib/constants/config.ts to use an example like
"meta/llama3-1-8b-instruct" (consistent with other examples such as
"openai/gpt-4.1-mini" and "black-forest-labs/flux-dev") so users supply the
owner/model string and the system can prepend "replicate/" internally.
♻️ Duplicate comments (2)
core/providers/replicate/types.go (1)
231-248: ResetReplicateOutputfields before decoding.Reused structs can retain prior values, leaving both fields set and triggering marshal errors. Reset both fields up‑front and clear the alternate on successful parse.
🐛 Reset fields before parsing
func (mc *ReplicateOutput) UnmarshalJSON(data []byte) error { + mc.OutputStr = nil + mc.OutputArray = nil if string(data) == "null" || len(data) == 0 { - mc.OutputStr = nil - mc.OutputArray = nil return nil } @@ if err := sonic.Unmarshal(data, &stringContent); err == nil { mc.OutputStr = schemas.Ptr(stringContent) + mc.OutputArray = nil return nil } @@ if err := sonic.Unmarshal(data, &arrayContent); err == nil { + mc.OutputStr = nil mc.OutputArray = arrayContent return nil }core/providers/replicate/replicate.go (1)
1389-1393: Remove stdout debug prints that can leak raw image data.
This was already flagged earlier; thefmt.Printlncalls still dump raw SSE payloads (including base64 image data) to stdout.🧹 Remove the debug prints
- fmt.Println("sendbackrawresponse", sendBackRawResponse) - // Set raw response to the actual SSE event data if enabled (as JSON string) - if sendBackRawResponse { - if eventJSON, err := sonic.Marshal(currentEvent); err == nil { - chunk.ExtraFields.RawResponse = string(eventJSON) - } - } + // Set raw response to the actual SSE event data if enabled (as JSON string) + if sendBackRawResponse { + if eventJSON, err := sonic.Marshal(currentEvent); err == nil { + chunk.ExtraFields.RawResponse = string(eventJSON) + } + }
🧹 Nitpick comments (2)
core/internal/testutil/account.go (1)
1201-1228: Align Replicate test scenarios with supported features.The PR summary indicates Replicate supports text completions and file operations, but those scenarios are currently disabled/omitted here. If they’re supported, enabling them will keep the comprehensive suite aligned with the provider surface.
♻️ Possible update to Replicate scenarios
Scenarios: TestScenarios{ - TextCompletion: false, // Not typical + TextCompletion: true, + TextCompletionStream: true, SimpleChat: true, CompletionStream: true, @@ ImageGeneration: true, ImageGenerationStream: false, + FileUpload: true, + FileList: true, + FileRetrieve: true, + FileDelete: true, + FileContent: true, },core/providers/replicate/text.go (1)
16-22: Add prompt validation to fail fast.
Although the HTTP handler validates prompt presence upstream, validating at the provider level is a defensive programming best practice. IfToReplicateTextRequestis called outside the HTTP path (e.g., in tests or future code paths), an empty prompt would reach the API. Since prompt is required for Replicate text models, add an error check here to catch issues early.Proposed guard
if bifrostReq.Input.PromptStr != nil { input.Prompt = bifrostReq.Input.PromptStr } else if len(bifrostReq.Input.PromptArray) > 0 { prompt := strings.Join(bifrostReq.Input.PromptArray, "\n") input.Prompt = &prompt + } else { + return nil, fmt.Errorf("prompt is required") }
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
e72bf1f to
c109eed
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@core/internal/testutil/account.go`:
- Around line 1201-1228: Add a Fallbacks entry to the Replicate provider config
in AllProviderConfigs (the config object where Provider: schemas.Replicate is
defined) so it mirrors other providers' behavior; populate Fallbacks to include
the OpenAI provider (or the appropriate primary fallback used elsewhere) and any
model overrides needed for ChatModel/TextModel fallbacks; update the same
provider config object to include the Fallbacks field alongside
ChatModel/TextModel so test scenarios that rely on fallback routing will use
OpenAI as expected.
In `@docs/providers/supported-providers/replicate.mdx`:
- Around line 550-552: The documentation Note currently labels the link as
"Replicate Files API documentation" but points to the predictions page; update
the href in that Note to the correct Replicate Files API documentation URL (or
replace the link text to accurately reflect the predictions docs if a Files API
page doesn't exist) so the label and target match; look for the Note block
containing the existing link to
https://replicate.com/docs/topics/predictions/create-a-prediction and change it
to the correct Files API URL or adjust the copy to "Replicate Predictions API
documentation" accordingly.
- Around line 354-358: The parameter mapping erroneously maps the key
"num_inference_steps" to the singular string "num_inference_step"; update the
mapping so "num_inference_steps" maps to "num_inference_steps" (plural) to match
the Go JSON tags in core/providers/replicate/types.go and the documentation
examples—locate the mapping object that includes "seed", "negative_prompt",
"num_inference_steps", "input_images" and replace the incorrect value.
♻️ Duplicate comments (5)
ui/lib/constants/config.ts (1)
52-52: Update Replicate model placeholder to use correctowner/model-nameformat.The placeholder uses
replicate/llama3-1-8b-instructformat, but Replicate models should be specified asowner/model-name(e.g.,meta/llama-2-7b-chat). Thereplicate/prefix is added internally by Bifrost's provider routing. Based on the test configuration incore/internal/testutil/account.go, correct examples would beopenai/gpt-4.1-miniorblack-forest-labs/flux-dev.Suggested fix
- replicate: "e.g. replicate/llama3-1-8b-instruct, replicate/llama3-2-70b-instruct", + replicate: "e.g. meta/llama-2-7b-chat, black-forest-labs/flux-schnell",core/providers/replicate/models.go (1)
65-78: Owner/Name not split as suggested in prior review.A past review comment suggested splitting
modelIDinto separateOwnerandNamefields forReplicateModelResponse, but the current implementation setsNameto the fullmodelID(e.g., "owner/name"). If theReplicateModelResponse.Namefield is expected to contain only the model name (not owner), this should be addressed.core/providers/replicate/images.go (1)
24-45: Version-ID models cannot specify custom input image field.When using a 64-character version ID instead of a model name (e.g., for Flux variants),
getInputImageFieldName()falls back toinput_images, which may be incorrect for models requiringimage_promptorimage. A prior review suggested allowing an override viaExtraParams.core/providers/replicate/types.go (1)
233-251: ResetReplicateOutputfields before decoding.
UnmarshalJSONonly clears fields fornull; if the struct is reused (e.g., pooling or re-decoding), stale values can leave both fields set and makeMarshalJSONfail.🐛 Proposed fix
func (mc *ReplicateOutput) UnmarshalJSON(data []byte) error { - if string(data) == "null" || len(data) == 0 { - mc.OutputStr = nil - mc.OutputArray = nil - return nil - } + mc.OutputStr = nil + mc.OutputArray = nil + if string(data) == "null" || len(data) == 0 { + return nil + } // First, try to unmarshal as a direct string var stringContent string if err := sonic.Unmarshal(data, &stringContent); err == nil { mc.OutputStr = schemas.Ptr(stringContent) + mc.OutputArray = nil return nil } @@ var arrayContent []string if err := sonic.Unmarshal(data, &arrayContent); err == nil { mc.OutputArray = arrayContent + mc.OutputStr = nil return nil }#!/bin/bash # Verify if ReplicateOutput is ever reused or pooled across decodes. # Expectation: no pooling/reuse; if found, the reset is required. rg -n "ReplicateOutput" -C3 --type=go rg -n "sync\.Pool|pool" -C3 --type=go core/providers/replicatecore/providers/replicate/replicate.go (1)
1389-1394: Remove stdout debug print in image streaming.
This can leak raw image data to logs and is noisy in production.🧹 Suggested fix
- fmt.Println("sendbackrawresponse", sendBackRawResponse) // Set raw response to the actual SSE event data if enabled (as JSON string) if sendBackRawResponse { if eventJSON, err := sonic.Marshal(currentEvent); err == nil { chunk.ExtraFields.RawResponse = string(eventJSON) } }
🧹 Nitpick comments (12)
core/internal/testutil/list_models.go (1)
111-232: Use explicitt.Skip()instead of commenting out the entire pagination test functionLarge commented-out code blocks obscure the intent and make re-enablement less visible. The function is intentionally disabled as part of this stack; uncomment it and use
t.Skip()with a clear message so test discovery shows it as explicitly skipped rather than hidden.♻️ Proposed refactor
-// // RunListModelsPaginationTest executes pagination test for list models -// func RunListModelsPaginationTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) { -// if !testConfig.Scenarios.ListModels { -// t.Logf("List models not supported for provider %s", testConfig.Provider) -// return -// } -// -// t.Run("ListModelsPagination", func(t *testing.T) { -// ... -// }) -// } +// RunListModelsPaginationTest executes pagination test for list models +func RunListModelsPaginationTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) { + t.Skip("ListModels pagination tests temporarily disabled; enable when provider pagination is stable") +}Also uncomment the function reference in
tests.goline 60 so the test runner sees it (as skipped) rather than the function being completely hidden.core/providers/replicate/errors.go (1)
9-22: Consider enriching error message with Title field when available.The
ReplicateErrorstruct includesTitleandTypefields that could provide additional context. Currently onlyDetailis used in the error message.Optional enhancement
func parseReplicateError(body []byte, statusCode int) *schemas.BifrostError { var replicateErr ReplicateError if err := sonic.Unmarshal(body, &replicateErr); err == nil && replicateErr.Detail != "" { + message := replicateErr.Detail + if replicateErr.Title != nil && *replicateErr.Title != "" { + message = *replicateErr.Title + ": " + replicateErr.Detail + } return &schemas.BifrostError{ IsBifrostError: false, StatusCode: &statusCode, Error: &schemas.ErrorField{ - Message: replicateErr.Detail, + Message: message, }, ExtraFields: schemas.BifrostErrorExtraFields{ Provider: schemas.Replicate, }, } }core/providers/replicate/utils.go (2)
26-36: Minor inconsistency in error message handling.The error message construction has a subtle issue: the outer message is always "prediction failed" regardless of whether
prediction.Errorcontains a more specific message. The innerfmt.ErrorfuseserrorMsg(which may be the specific error), but the first argument toNewBifrostOperationErroris always the generic string.♻️ Suggested fix
func checkForErrorStatus(prediction *ReplicatePredictionResponse) *schemas.BifrostError { if prediction.Status == ReplicatePredictionStatusFailed { errorMsg := "prediction failed" if prediction.Error != nil && *prediction.Error != "" { errorMsg = *prediction.Error } return providerUtils.NewBifrostOperationError( - "prediction failed", + errorMsg, fmt.Errorf("%s", errorMsg), schemas.Replicate, ) }
144-150: Accept both uppercase and lowercase hex digits in version ID pattern.The regex
^[a-f0-9]{64}$only matches lowercase hex characters. While Replicate's documented examples use lowercase, the pattern should be more defensive and accept uppercase as well—either by using^[0-9a-fA-F]{64}$or the case-insensitive flag—to guard against any API changes or edge cases.core/providers/replicate/models.go (1)
23-37: Remove redundant assignment.Line 28 reassigns
modelNameto the same value it already has from line 23. This assignment inside theif deployment.CurrentRelease != nilblock is unnecessary.♻️ Suggested fix
modelName := schemas.Ptr(deployment.Name) var created *int64 // Extract information from current release if available if deployment.CurrentRelease != nil { - modelName = schemas.Ptr(deployment.Name) - // Parse created timestamp if deployment.CurrentRelease.CreatedAt != "" {core/providers/replicate/images.go (1)
169-214: Consider making model lists configurable or documenting maintenance requirements.The hardcoded lists of models (
imagePromptModels,inputImageModels,imageModels) will require manual updates as new Flux variants are released. Consider adding a comment noting this maintenance requirement, or allowing model-specific overrides via configuration.core/providers/replicate/files.go (2)
22-56: Consider extracting shared logic to reduce duplication.
ToBifrostFileUploadResponseandToBifrostFileRetrieveResponseshare nearly identical logic for mapping fields, handlingExpiresAt, and setting raw request/response. Consider extracting a helper to populate common fields.♻️ Example refactor
// populateCommonFileFields sets fields shared between upload and retrieve responses func (r *ReplicateFileResponse) populateCommonFileFields( providerName schemas.ModelProvider, latency time.Duration, requestType schemas.RequestType, sendBackRawRequest, sendBackRawResponse bool, rawRequest, rawResponse interface{}, ) (id string, object string, bytes int64, createdAt int64, filename string, purpose schemas.FilePurpose, status schemas.FileStatus, storageBackend schemas.FileStorageBackend, expiresAt *int64, extraFields schemas.BifrostResponseExtraFields) { // ... shared logic }Also applies to: 59-93
13-19: Status inference may not capture all Replicate file states.The logic assumes a file with
ID != ""andSize > 0isProcessed. However, Replicate files could potentially be in intermediate states (uploading, processing) that this heuristic doesn't capture. The current approach is reasonable given Replicate's API doesn't expose explicit status, but consider adding a comment documenting this assumption.core/providers/replicate/text.go (1)
53-65: Potential duplicatetop_kin request payload.
top_kis extracted fromExtraParamsand set oninput.TopK(lines 54-56), but then the entireExtraParamsmap is also assigned toinput.ExtraParams(lines 63-65). If the JSON marshaling merges these,top_kcould appear twice. Consider removingtop_kfromExtraParamsafter extraction, or document thatExtraParamsis intended to be passed through as-is.♻️ Suggested fix
// Top K (from ExtraParams) if topK, ok := schemas.SafeExtractIntPointer(params.ExtraParams["top_k"]); ok { input.TopK = topK + // Remove from ExtraParams to avoid duplication + delete(params.ExtraParams, "top_k") }core/providers/replicate/chat.go (3)
218-230: Preferschemas.Ptr()for pointer creation.Based on learnings, the codebase prefers
schemas.Ptr()over creating local variables for pointer assignment. The text.go file already uses this pattern (line 119).♻️ Suggested fix
// Determine finish reason based on status - var finishReason *string switch response.Status { case ReplicatePredictionStatusSucceeded: - reason := "stop" - finishReason = &reason + finishReason = schemas.Ptr("stop") case ReplicatePredictionStatusFailed: - reason := "error" - finishReason = &reason + finishReason = schemas.Ptr("error") case ReplicatePredictionStatusCanceled: - reason := "stop" - finishReason = &reason + finishReason = schemas.Ptr("stop") }Note: You'll need to declare
var finishReason *stringbefore the switch or handle the nil case.Based on learnings, prefer
bifrost.Ptr()(orschemas.Ptr()) for pointer creation.
263-288: Hardcoded model lists require ongoing maintenance.The
unsupportedModelslist and deepseek prefix check will need updates as new models are added. Consider:
- Adding a comment documenting when this list was last updated
- Allowing overrides via request parameters for edge cases
183-193:BifrostChatResponse.Createdfield intentionally usesint; the cast on line 188 is consistent with schema design.The
Createdfield inBifrostChatResponseis defined asintincore/schemas/chatcompletions.go(line 28), notint64. While line 188 does cast fromint64toint, this pattern is consistent across all providers (Gemini, Cohere, Bedrock, Anthropic). This appears to be a deliberate schema design choice rather than an oversight. However, note that usingintfor Unix timestamps introduces a Y2038 problem (32-bit signed integer overflow). If future-proofing is a concern, consider whetherBifrostChatResponse.Createdshould be changed toint64at the schema level, which would benefit all providers.
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
c109eed to
9ff46c4
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@core/providers/replicate/replicate.go`:
- Around line 196-198: The polling context is currently created from
context.Background(), which prevents parent (Bifrost) cancellation from stopping
the poll; change the poll context creation to derive from the parent context
(use the existing parent variable, e.g., ctx) by calling
context.WithTimeout(ctx, time.Duration(timeoutSeconds)*time.Second) so pollCtx
respects parent cancellation, keep the returned cancel and defer cancel() as-is
to release resources; update the code around pollCtx and cancel where
context.WithTimeout is called.
♻️ Duplicate comments (2)
core/providers/replicate/images.go (1)
24-45: Version-ID models still default toinput_imageswithout override mechanism.When the model is a 64-character version ID,
getInputImageFieldName()cannot match it against hardcoded model names and defaults toinput_images. This may cause issues with Flux variants that requireimage_promptorimagefields. The previous review suggested allowing an explicit field override viaExtraParams, but this doesn't appear to be implemented.Consider supporting an override when the model is a version ID:
🔧 Suggested approach
// Map InputImages to the appropriate field based on model if len(params.InputImages) > 0 { fieldName := getInputImageFieldName(bifrostReq.Model) + // Allow ExtraParams override for version IDs where model name matching fails + if isVersionID(bifrostReq.Model) && params.ExtraParams != nil { + if override, ok := params.ExtraParams["input_image_field"].(string); ok && override != "" { + fieldName = override + delete(params.ExtraParams, "input_image_field") // Remove from extra params to avoid sending + } + } switch fieldName {core/providers/replicate/replicate.go (1)
1389-1389: Remove debug print statement.This
fmt.Printlndebug statement leaks internal state to stdout and should be removed before merging.- fmt.Println("sendbackrawresponse", sendBackRawResponse)
🧹 Nitpick comments (4)
core/internal/testutil/list_models.go (1)
111-232: Consider removing commented-out code or adding a TODO comment.Keeping large blocks of commented-out code makes the file harder to maintain. If pagination testing is temporarily disabled pending broader support, consider:
- Adding a TODO comment explaining when/why it should be re-enabled, or
- Removing the code entirely and restoring from version control when needed.
If this is intentional for the Replicate provider rollout (which may not support pagination), a brief comment explaining the rationale would help future maintainers.
core/providers/replicate/images.go (1)
169-214: Consider documenting the model-field mapping.The hardcoded model lists work for known models but will need maintenance as new models are added. Consider adding a comment indicating where to find the canonical source for these mappings (e.g., Replicate API documentation).
core/providers/replicate/chat.go (1)
218-230: Useschemas.Ptr()for consistency.Lines 222-229 use local variables and pointers (
&reason) instead ofschemas.Ptr(), which is inconsistent with the pattern used intext.go(lines 119-124). Based on learnings, prefer usingbifrost.Ptr()orschemas.Ptr()to create pointers for improved consistency.♻️ Suggested refactor
// Determine finish reason based on status var finishReason *string switch response.Status { case ReplicatePredictionStatusSucceeded: - reason := "stop" - finishReason = &reason + finishReason = schemas.Ptr("stop") case ReplicatePredictionStatusFailed: - reason := "error" - finishReason = &reason + finishReason = schemas.Ptr("error") case ReplicatePredictionStatusCanceled: - reason := "stop" - finishReason = &reason + finishReason = schemas.Ptr("stop") }core/providers/replicate/replicate.go (1)
1519-1527: Use consistent SSE parsing pattern.This block uses
HasPrefix+TrimPrefix, but the text/chat streaming (lines 683-693, 1045-1056) usesstrings.CutPrefix. UsingCutPrefixis more efficient (single operation) and improves consistency.♻️ Suggested refactor
- if strings.HasPrefix(line, "event:") { - currentEvent.Event = strings.TrimSpace(strings.TrimPrefix(line, "event:")) - } else if strings.HasPrefix(line, "data:") { - currentEvent.Data = strings.TrimSpace(strings.TrimPrefix(line, "data:")) - } else if strings.HasPrefix(line, "id:") { - currentEvent.ID = strings.TrimSpace(strings.TrimPrefix(line, "id:")) + if after, ok := strings.CutPrefix(line, "event: "); ok { + currentEvent.Event = strings.TrimSpace(after) + } else if after, ok := strings.CutPrefix(line, "data: "); ok { + currentEvent.Data = after + } else if after, ok := strings.CutPrefix(line, "id: "); ok { + currentEvent.ID = strings.TrimSpace(after) }
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
2 times, most recently
from
42d2123 to
bc7fb91
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
core/internal/testutil/tests.go (1)
24-61: Avoid silently disabling pagination coverage.Commenting out
RunListModelsPaginationTestremoves coverage without a reason or a replacement. If it’s intentionally disabled, consider adding a TODO with rationale (and ideally a tracking issue), or gate it behind a scenario flag instead of commenting it out.💡 Suggested tweak
- // RunListModelsPaginationTest, + // TODO: Re-enable pagination test once list-model pagination is stabilized for all providers. + // RunListModelsPaginationTest,core/schemas/bifrost.go (1)
55-86: Add Replicate to SupportedBaseProviders to match UI.The UI exposes "Replicate" as a selectable base provider option in the custom provider dialog, but
SupportedBaseProvidersdoes not include it. This causes a mismatch: users can select Replicate in the UI, but the backend'sIsSupportedBaseProvider()validation will reject custom provider creation attempts with Replicate as the base provider.Suggested fix
var SupportedBaseProviders = []ModelProvider{ Anthropic, Bedrock, Cohere, Gemini, OpenAI, HuggingFace, + Replicate, }
🤖 Fix all issues with AI agents
In `@core/providers/replicate/images.go`:
- Around line 127-141: The code dereferences prediction.Output (checking
prediction.Output.OutputStr and prediction.Output.OutputArray) without guarding
for nil, which can panic when Replicate returns a prediction with no output;
update the block in core/providers/replicate/images.go that builds response.Data
(the code referencing prediction.Output, prediction.Output.OutputStr,
prediction.Output.OutputArray and appending schemas.ImageData) to first check if
prediction.Output != nil and only then inspect OutputStr/OutputArray, otherwise
skip adding ImageData (or return/continue as appropriate) so the function safely
handles nil outputs.
In `@core/providers/replicate/models.go`:
- Around line 14-51: The list-models response never maps Replicate pagination
into the Bifrost NextPageToken, so when deploymentsResponse.Next is set clients
can't continue pagination; update the function that builds bifrostResponse (the
block iterating deploymentsResponse.Results) to set
bifrostResponse.NextPageToken from deploymentsResponse.Next (or a parsed token
if Bifrost expects a different format) before returning; look for symbols
deploymentsResponse, bifrostResponse, and schemas.BifrostListModelsResponse and
ensure NextPageToken is populated appropriately (e.g., copy
deploymentsResponse.Next or transform it) so callers can resume paging.
In `@core/providers/replicate/replicate.go`:
- Around line 1710-1712: The logger.Debug calls are passing printf-style format
strings and separate args which won't be formatted; wrap the format string with
fmt.Sprintf before passing to provider.logger.Debug (e.g., replace
provider.logger.Debug("error from %s provider: %s", providerName,
string(resp.Body())) with provider.logger.Debug(fmt.Sprintf(...))). Update all
similar sites in this file where provider.logger.Debug is used with format
strings (the shown call near parseReplicateError plus the other occurrences
around the other noted blocks), importing fmt if missing, and keep the same
message text and args when building the formatted string.
♻️ Duplicate comments (2)
core/providers/replicate/models.go (1)
65-78: Split owner/name when building Replicate model responses.
ReplicateModelResponse.Nameshould be just the model name, notowner/name, andOwnershould be populated even whenOwnedByis nil.🐛 Proposed fix
for _, model := range response.Data { modelID := strings.TrimPrefix(model.ID, string(schemas.Replicate)+"/") + owner := "" + name := modelID + if parts := strings.SplitN(modelID, "/", 2); len(parts) == 2 { + owner, name = parts[0], parts[1] + } replicateModel := ReplicateModelResponse{ URL: "https://replicate.com/" + modelID, - Name: modelID, + Owner: owner, + Name: name, } @@ - if model.OwnedBy != nil { - replicateModel.Owner = *model.OwnedBy - } + if model.OwnedBy != nil { + replicateModel.Owner = *model.OwnedBy + }core/providers/replicate/images.go (1)
24-44: Allow explicit input-image field override for version IDs.When the model is a version ID,
getInputImageFieldNamecan’t resolve Flux variants and defaults toinput_images. An override avoids sending invalid payloads for Flux Pro/Fill versions.🔧 Suggested approach
if len(params.InputImages) > 0 { fieldName := getInputImageFieldName(bifrostReq.Model) + if isVersionID(bifrostReq.Model) && params.ExtraParams != nil { + if override, ok := params.ExtraParams["input_image_field"].(string); ok && override != "" { + fieldName = override + } + } switch fieldName {
🧹 Nitpick comments (7)
core/providers/replicate/errors.go (1)
3-34: Preferbifrost.Ptrfor pointer creation.Align with the repo’s pointer helper for consistency and readability.
♻️ Proposed refactor
import ( "github.com/bytedance/sonic" + "github.com/maximhq/bifrost/core/bifrost" schemas "github.com/maximhq/bifrost/core/schemas" ) @@ return &schemas.BifrostError{ IsBifrostError: false, - StatusCode: &statusCode, + StatusCode: bifrost.Ptr(statusCode), Error: &schemas.ErrorField{ Message: replicateErr.Detail, }, @@ return &schemas.BifrostError{ IsBifrostError: false, - StatusCode: &statusCode, + StatusCode: bifrost.Ptr(statusCode), Error: &schemas.ErrorField{ Message: string(body), },Based on learnings, prefer
bifrost.Ptr(...)over&value.core/internal/testutil/account.go (1)
390-398: Minor inconsistency:UseForBatchAPIis true but Replicate doesn't support Batch API.According to the documentation, Replicate does not support the Batch API (marked with ❌). Setting
UseForBatchAPI: bifrost.Ptr(true)is inconsistent, though this may not cause runtime issues if batch tests are gated byTestScenarios.BatchCreateetc. (which arefalsein the Replicate config).Consider setting this to
falsefor accuracy:Suggested fix
case schemas.Replicate: return []schemas.Key{ { Value: *schemas.NewEnvVar("env.REPLICATE_API_KEY"), Models: []string{}, Weight: 1.0, - UseForBatchAPI: bifrost.Ptr(true), + UseForBatchAPI: bifrost.Ptr(false), }, }, nilcore/providers/replicate/files.go (1)
39-44: Usebifrost.PtrforExpiresAtpointers.The repository prefers pointer helpers over address-of; here
resp.ExpiresAt = &expiresAtappears in both upload and retrieve responses. Switch tobifrost.Ptr(expiresAt)for consistency.♻️ Proposed refactor
import ( "time" + bifrost "github.com/maximhq/bifrost/core/bifrost" "github.com/maximhq/bifrost/core/schemas" ) @@ - resp.ExpiresAt = &expiresAt + resp.ExpiresAt = bifrost.Ptr(expiresAt) @@ - resp.ExpiresAt = &expiresAt + resp.ExpiresAt = bifrost.Ptr(expiresAt)Based on learnings, apply
bifrost.Ptr(...)instead of&value.Also applies to: 76-81
core/providers/replicate/chat.go (2)
80-99: Preferbifrost.Ptrfor pointer creation in chat converters.Several pointers are created with
&value(e.g., system prompt, prompt, joined output, finish reason, and message content). Usebifrost.Ptr(...)for consistency with repo conventions.♻️ Proposed refactor
import ( "fmt" "slices" "strings" "time" + bifrost "github.com/maximhq/bifrost/core/bifrost" schemas "github.com/maximhq/bifrost/core/schemas" ) @@ - input.SystemPrompt = &systemPrompt + input.SystemPrompt = bifrost.Ptr(systemPrompt) @@ - input.Prompt = &prompt + input.Prompt = bifrost.Ptr(prompt) @@ - joined := strings.Join(response.Output.OutputArray, "") - contentStr = &joined + joined := strings.Join(response.Output.OutputArray, "") + contentStr = bifrost.Ptr(joined) @@ - message := schemas.ChatMessage{ + message := schemas.ChatMessage{ Role: schemas.ChatMessageRoleAssistant, - Content: &messageContent, + Content: bifrost.Ptr(messageContent), } @@ - finishReason = &reason + finishReason = bifrost.Ptr(reason)Based on learnings, use
bifrost.Ptr(...)instead of address-of.Also applies to: 196-205, 208-215, 219-229
20-104: Sending bothmessagesandprompt/system_promptis redundant for OpenAI Replicate models.When the model is OpenAI-prefixed, you set
input.Messages, but later also buildPrompt/SystemPromptfrom the same content. Per Replicate's API documentation, whenmessagesis provided, Replicate simply ignorespromptandsystem_promptfields. This works correctly but sends redundant data. Consider gating prompt/system_prompt construction with a flag to avoid unnecessary processing.🔧 Suggested optimization
- // openai models support messages - if len(bifrostReq.Input) > 0 && strings.HasPrefix(bifrostReq.Model, string(schemas.OpenAI)) { - input.Messages = bifrostReq.Input - } + useMessages := len(bifrostReq.Input) > 0 && strings.HasPrefix(bifrostReq.Model, string(schemas.OpenAI)) + if useMessages { + input.Messages = bifrostReq.Input + } @@ - if systemPrompt != "" { + if !useMessages && systemPrompt != "" { @@ - if len(conversationParts) > 0 { + if !useMessages && len(conversationParts) > 0 { @@ - if input.Prompt == nil && input.SystemPrompt == nil { + if !useMessages && input.Prompt == nil && input.SystemPrompt == nil { return nil, fmt.Errorf("no content found in chat messages - need at least one user or system message") }core/providers/replicate/replicate.go (2)
129-131: Avoid hardcodingschemas.Replicatein error construction.Several errors are created with a hardcoded provider key. For custom provider aliases, these should reflect
provider.GetProviderKey()/providerName. Consider plumbingproviderNameinto helpers and using it consistently.♻️ Suggested pattern
-func createPrediction( +func createPrediction( ctx *schemas.BifrostContext, client *fasthttp.Client, jsonBody []byte, key schemas.Key, url string, extraHeaders map[string]string, + providerName schemas.ModelProvider, logger schemas.Logger, sendBackRawRequest bool, sendBackRawResponse bool, ) (*ReplicatePredictionResponse, interface{}, time.Duration, *schemas.BifrostError) { @@ - return nil, nil, latency, providerUtils.NewBifrostOperationError(schemas.ErrProviderResponseDecode, decodeErr, schemas.Replicate) + return nil, nil, latency, providerUtils.NewBifrostOperationError(schemas.ErrProviderResponseDecode, decodeErr, providerName)- return nil, providerUtils.NewBifrostOperationError( + return nil, providerUtils.NewBifrostOperationError( "failed to parse deployments response", err, - schemas.Replicate, + providerName, )Based on learnings, prefer
provider.GetProviderKey()over hardcoded provider constants for error context.Also applies to: 181-182, 227-231, 306-311
1042-1053: Usebifrost.Ptrfor string pointers in streaming and pagination.
&finishReasonand&finalCursorshould follow the pointer-helper convention.♻️ Proposed refactor
import ( "bufio" "bytes" "context" "fmt" @@ "time" "github.com/bytedance/sonic" + bifrost "github.com/maximhq/bifrost/core/bifrost" providerUtils "github.com/maximhq/bifrost/core/providers/utils" schemas "github.com/maximhq/bifrost/core/schemas" "github.com/valyala/fasthttp" ) @@ - finishReason := "stop" + finishReason := "stop" finalResponse := &schemas.BifrostChatResponse{ @@ - FinishReason: &finishReason, + FinishReason: bifrost.Ptr(finishReason), @@ - bifrostResp.After = &finalCursor + bifrostResp.After = bifrost.Ptr(finalCursor)Based on learnings, use
bifrost.Ptr(...)instead of address-of.Also applies to: 1848-1851
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
bc7fb91 to
b4c6d0b
Compare
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Fix all issues with AI agents
In `@core/internal/testutil/account.go`:
- Around line 1201-1228: Enable Replicate's file operation scenarios by setting
the file-related TestScenarios flags to true in the Replicate provider entry
(the struct with Provider: schemas.Replicate) within AllProviderConfigs:
add/turn on FileUpload, FileList, FileRetrieve and FileDelete (or the existing
file operation boolean fields on TestScenarios) so the Replicate configuration
matches its own replicate_test.go expectations.
In `@core/providers/replicate/replicate.go`:
- Around line 1371-1377: In ImageGenerationStream, the SSE parsing loop doesn't
skip comment lines (those beginning with ":"), so modify the loop that reads
scanner.Text() to ignore lines starting with ":" before treating empty lines as
event delimiters; specifically, update the section around currentEvent handling
(where line := scanner.Text() and the empty-line check uses
currentEvent.Event/currentEvent.Data) to continue when strings.HasPrefix(line,
":") so keep-alive/comments are skipped just like in
TextCompletionStream/ChatCompletionStream.
- Around line 1574-1578: The scanner error handling in the SSE loop should
mirror TextCompletionStream/ChatCompletionStream by checking context
cancellation before dispatching errors: after obtaining err := scanner.Err(),
keep the provider.logger.Warn(...) but add a guard if ctx.Err() != nil { return
} (or simply skip calling providerUtils.ProcessAndSendError) to avoid sending
errors when the request context was cancelled/expired; update the block around
scanner.Err(), providerUtils.ProcessAndSendError(...), and related symbols so
ProcessAndSendError is only called when ctx.Err() is nil.
In `@core/providers/replicate/text.go`:
- Around line 16-22: After mapping bifrostReq.Input.PromptStr / PromptArray into
the ReplicatePredictionRequestInput (the input variable of type
ReplicatePredictionRequestInput), validate that input.Prompt is non-nil and
contains a non-empty string; if it is nil or empty (e.g., PromptArray was
non-nil but empty) return an error immediately instead of letting the upstream
call proceed. Update the logic around input.Prompt assignment (referencing
bifrostReq.Input.PromptStr, bifrostReq.Input.PromptArray, and
ReplicatePredictionRequestInput) to perform this guard and fail fast with a
clear error when the prompt is empty.
- Around line 86-101: Remove the dead timestamp parsing: delete the
ParseReplicateTimestamp(response.CreatedAt) call and the createdAt variable (and
its fallback to time.Now().Unix()) since the parsed value is never used or
assigned to schemas.BifrostTextCompletionResponse; locate the block around
ParseReplicateTimestamp, createdAt, and the subsequent fallback and remove those
three lines so only the BifrostTextCompletionResponse initialization (including
ExtraFields) remains.
♻️ Duplicate comments (4)
core/schemas/files.go (1)
58-61: LGTM! Verify HTTP handler wiresContentTypefrom multipart headers.The
ContentTypefield addition is well-structured. Based on past review discussion, ensure the HTTP handler attransports/bifrost-http/handlers/inference.go(around line 1830) extracts the Content-Type fromfileHeader.Header.Get("Content-Type")and assigns it tobifrostFileReq.ContentType.#!/bin/bash # Verify the handler populates ContentType rg -n "ContentType" transports/bifrost-http/handlers/inference.go -B 2 -A 2core/providers/replicate/models.go (1)
69-84: Splitowner/namefor correctReplicateModelResponsefields.As noted in a past review,
ReplicateModelResponseexpectsOwnerandNameas separate fields. CurrentlyNameis set to the fullmodelID(e.g.,"owner/model-name"), andOwneris only set frommodel.OwnedBy. This can result inNamecontaining the owner prefix.♻️ Suggested fix
for _, model := range response.Data { modelID := strings.TrimPrefix(model.ID, string(schemas.Replicate)+"/") + owner := "" + name := modelID + if parts := strings.SplitN(modelID, "/", 2); len(parts) == 2 { + owner, name = parts[0], parts[1] + } replicateModel := ReplicateModelResponse{ URL: "https://replicate.com/" + modelID, - Name: modelID, + Name: name, } if model.Description != nil { replicateModel.Description = model.Description } - if model.OwnedBy != nil { + // Prefer owner from model.OwnedBy, fallback to parsed owner + if model.OwnedBy != nil && *model.OwnedBy != "" { replicateModel.Owner = *model.OwnedBy + } else if owner != "" { + replicateModel.Owner = owner }core/providers/replicate/images.go (1)
24-45: Version-ID models default toinput_images, which may be incorrect for Flux variants.When a model is passed as a 64-character version ID (e.g.,
"80a09d66baa990429c2f5ae8a4306bf778a1b3775afd01cc2cc8bdbe9033769c"),getInputImageFieldName()cannot match it against the hardcoded model names and defaults to"input_images". However, FLUX 1.1 Pro expectsimage_promptand FLUX Fill expectsimage.Consider allowing an explicit field override via
ExtraParamsfor version-ID usage:♻️ Suggested approach
if len(params.InputImages) > 0 { fieldName := getInputImageFieldName(bifrostReq.Model) + + // Allow override for version IDs where model name can't be matched + if params.ExtraParams != nil { + if override, ok := params.ExtraParams["input_image_field"].(string); ok && override != "" { + fieldName = override + } + } switch fieldName {core/providers/replicate/replicate.go (1)
1720-1722: Usefmt.Sprintfto wrap format strings inlogger.Debugcalls.The
Logger.Debugmethod signature isDebug(msg string, args ...any)for structured logging—the variadicargsare not format arguments. These calls pass printf-style format strings that won't be interpolated.This was flagged in a previous review. The same pattern appears at lines 1807, 1904, 1991, and 2119.
🔧 Example fix
- provider.logger.Debug("error from %s provider: %s", providerName, string(resp.Body())) + provider.logger.Debug(fmt.Sprintf("error from %s provider: %s", providerName, string(resp.Body())))
🧹 Nitpick comments (8)
core/providers/replicate/models.go (2)
23-37: Remove redundantmodelNameassignment insideCurrentReleaseblock.Line 28 reassigns
modelName = schemas.Ptr(deployment.Name)which is identical to line 23. This assignment is unnecessary sincemodelNameis already set todeployment.Namebefore the conditional.♻️ Suggested fix
modelName := schemas.Ptr(deployment.Name) var created *int64 // Extract information from current release if available if deployment.CurrentRelease != nil { - modelName = schemas.Ptr(deployment.Name) - // Parse created timestamp if deployment.CurrentRelease.CreatedAt != "" {
39-45: Consider settingDeploymenttodeploymentIDinstead ofmodelName.Currently both
NameandDeploymentare set tomodelName(the deployment name without owner). For clarity and to match the pattern whereDeploymentrepresents the full deployment identifier, consider:bifrostModel := schemas.Model{ ID: string(providerKey) + "/" + deploymentID, Name: modelName, - Deployment: modelName, + Deployment: schemas.Ptr(deploymentID), OwnedBy: schemas.Ptr(deployment.Owner), Created: created, }This would make
Deploymentcontain the fullowner/namepath whileNamecontains just the name.core/providers/replicate/images.go (1)
172-215: Consider documenting or externalizing the model-to-field mapping.The hardcoded model lists in
getInputImageFieldNameare Replicate-specific knowledge that may need updates as new Flux variants are released. Consider:
- Adding a code comment with a reference to Replicate's documentation
- Or extracting this mapping to a configuration that can be updated without code changes
This is a minor maintainability suggestion—the current implementation is functional.
core/providers/replicate/text.go (1)
20-22: Use pointer helpers for local string pointers.For consistency with the repo’s pointer helper convention, prefer
schemas.Ptr(...)over taking addresses of locals in these spots.Based on learnings, please use the shared pointer helper for consistency.♻️ Suggested tweak
@@ } else if len(bifrostReq.Input.PromptArray) > 0 { prompt := strings.Join(bifrostReq.Input.PromptArray, "\n") - input.Prompt = &prompt + input.Prompt = schemas.Ptr(prompt) } @@ } else if response.Output.OutputArray != nil { // Join array of strings into a single string joined := strings.Join(response.Output.OutputArray, "") - textOutput = &joined + textOutput = schemas.Ptr(joined) } }Also applies to: 109-112
core/providers/replicate/errors.go (1)
12-15: Prefer pointer helper for status code pointers.Use
schemas.Ptr(statusCode)instead of&statusCodefor consistency with the repo’s pointer helper convention.Based on learnings, please use the shared pointer helper for consistency.♻️ Suggested tweak
@@ - StatusCode: &statusCode, + StatusCode: schemas.Ptr(statusCode), @@ - StatusCode: &statusCode, + StatusCode: schemas.Ptr(statusCode),Also applies to: 25-28
core/providers/replicate/replicate.go (3)
559-559: UnusedcontentBuildervariable.The
contentBuilderaccumulates content from output events but is never read or used afterward. Either remove it or use it for total content tracking if needed.♻️ Suggested removal
- var contentBuilder strings.Builder messageID := prediction.IDAnd remove the
contentBuilder.WriteString(currentEvent.Data)call at line 588.
923-923: UnusedcontentBuildervariable (duplicate pattern).Same as in
TextCompletionStream- thecontentBuilderis populated but never used.
1565-1571: Inconsistent SSE field parsing style.This streaming method uses
HasPrefix/TrimPrefixfor SSE parsing whileTextCompletionStreamandChatCompletionStreamusestrings.CutPrefix. Consider aligning for consistency.♻️ Suggested alignment with other streams
- if strings.HasPrefix(line, "event:") { - currentEvent.Event = strings.TrimSpace(strings.TrimPrefix(line, "event:")) - } else if strings.HasPrefix(line, "data:") { - currentEvent.Data = strings.TrimSpace(strings.TrimPrefix(line, "data:")) - } else if strings.HasPrefix(line, "id:") { - currentEvent.ID = strings.TrimSpace(strings.TrimPrefix(line, "id:")) + if after, ok := strings.CutPrefix(line, "event: "); ok { + currentEvent.Event = strings.TrimSpace(after) + } else if after, ok := strings.CutPrefix(line, "data: "); ok { + currentEvent.Data = after + } else if after, ok := strings.CutPrefix(line, "id: "); ok { + currentEvent.ID = strings.TrimSpace(after) }Note: The
data:handling in other streams also appends with newline for multiline data, which may need consideration here.
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
b4c6d0b to
4caf57e
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@core/providers/replicate/responses.go`:
- Around line 155-162: When params.Instructions is present but
supportsSystemPrompt(bifrostReq.Model) is false, we currently drop the
instructions; instead prepend (or otherwise combine) params.Instructions into
the input prompt so intent is preserved. Modify the branch around
supportsSystemPrompt to, when system prompt is unsupported, set input.Prompt (or
the existing user prompt field used by your request struct) =
params.Instructions + "\n\n" + existingPrompt (making sure to handle nil/empty
existingPrompt and avoid duplicating if input.SystemPrompt already set); keep
the existing behavior of assigning input.SystemPrompt when
supportsSystemPrompt(...) is true.
♻️ Duplicate comments (4)
core/internal/testutil/account.go (1)
1201-1228: File operation scenarios are still disabled for Replicate.
Replicate supports upload/list/retrieve/delete; enabling these keeps comprehensive tests aligned.🐛 Proposed change
Scenarios: TestScenarios{ ImageGeneration: true, ImageGenerationStream: false, + FileUpload: true, + FileList: true, + FileRetrieve: true, + FileDelete: true, },core/providers/replicate/models.go (1)
69-83: Split owner/name when building Replicate model responses.
Nameshould be the model name only, andOwnershould be derived from the ID whenOwnedByis absent.🐛 Proposed change
- modelID := strings.TrimPrefix(model.ID, string(schemas.Replicate)+"/") - replicateModel := ReplicateModelResponse{ - URL: "https://replicate.com/" + modelID, - Name: modelID, - } + modelID := strings.TrimPrefix(model.ID, string(schemas.Replicate)+"/") + owner := "" + name := modelID + if parts := strings.SplitN(modelID, "/", 2); len(parts) == 2 { + owner, name = parts[0], parts[1] + } + replicateModel := ReplicateModelResponse{ + URL: "https://replicate.com/" + modelID, + Name: name, + } + if owner != "" { + replicateModel.Owner = owner + } if model.OwnedBy != nil { replicateModel.Owner = *model.OwnedBy }core/providers/replicate/text.go (1)
10-76: Guard against empty prompt after mapping.
If PromptArray is empty,input.Promptcan still be nil; fail fast to avoid invalid requests.🐛 Proposed change
if bifrostReq.Input.PromptStr != nil { input.Prompt = bifrostReq.Input.PromptStr } else if len(bifrostReq.Input.PromptArray) > 0 { prompt := strings.Join(bifrostReq.Input.PromptArray, "\n") input.Prompt = &prompt } + + if input.Prompt == nil || strings.TrimSpace(*input.Prompt) == "" { + return nil, fmt.Errorf("prompt not provided for text completion request") + }core/providers/replicate/images.go (1)
24-44: Version‑ID models can misroute image inputs.When a model is a 64‑char version ID,
getInputImageFieldNamedefaults toinput_images, which doesn’t match Flux variants (e.g.,image_promptorimage). Consider allowing an explicit override via extra params for version‑ID flows.🔧 Suggested adjustment
- if len(params.InputImages) > 0 { - fieldName := getInputImageFieldName(bifrostReq.Model) + if len(params.InputImages) > 0 { + fieldName := getInputImageFieldName(bifrostReq.Model) + if isVersionID(bifrostReq.Model) && params.ExtraParams != nil { + if override, ok := params.ExtraParams["input_image_field"].(string); ok && override != "" { + fieldName = override + } + }Replicate Flux model input field names image_prompt vs image vs input_image (and behavior when using version IDs)
🧹 Nitpick comments (4)
core/internal/testutil/account.go (1)
1201-1227: Enable Replicate text completion scenarios to match implemented support.
The PR implements text completions, but the Replicate test config disables them, reducing coverage.♻️ Proposed change
Scenarios: TestScenarios{ - TextCompletion: false, // Not typical + TextCompletion: true, + TextCompletionStream: true, SimpleChat: true, CompletionStream: true,core/internal/testutil/list_models.go (1)
111-232: Remove the commented pagination test block instead of leaving dead code.
This keeps the file lean and avoids stale logic.♻️ Proposed change
-// // RunListModelsPaginationTest executes pagination test for list models -// func RunListModelsPaginationTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) { -// ... -// }core/internal/testutil/tests.go (1)
58-62: Prefer removing the commented-out scenario entry.
Keeps the test runner list clean once the test is disabled.♻️ Proposed change
- // RunListModelsPaginationTest,core/providers/replicate/responses.go (1)
189-233: Prefer pointer helpers for local address‑of values.Helps keep pointer creation consistent across the codebase.
Based on learnings, prefer pointer helpers over `&`.♻️ Pointer helper cleanup
- completedAt = &completed + completedAt = schemas.Ptr(completed) @@ - contentStr = &joined + contentStr = schemas.Ptr(joined) @@ - outputMsg := schemas.ResponsesMessage{ - Type: &messageType, - Role: &role, + outputMsg := schemas.ResponsesMessage{ + Type: schemas.Ptr(messageType), + Role: schemas.Ptr(role), @@ - bifrostResponse.Status = &status + bifrostResponse.Status = schemas.Ptr(status)Also applies to: 255-256
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
4caf57e to
5cdf2ac
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@core/internal/testutil/tests.go`:
- Line 60: Add a TODO comment next to the commented-out
RunListModelsPaginationTest explaining why the pagination test is disabled
(e.g., "Replicate doesn't support pagination", "test is flaky", or "temporarily
disabled while backend bug is fixed") and follow the same TODO style used in
batch.go; place the TODO in both locations where RunListModelsPaginationTest is
commented out (the tests.go comment and the commented block in list_models.go
around the pagination test) and include a short action or ticket reference if
applicable.
In `@core/providers/replicate/types.go`:
- Around line 455-463: The ParseReplicateTimestamp function currently uses
time.RFC3339 which fails on timestamps with fractional seconds; update
ParseReplicateTimestamp to parse using time.RFC3339Nano instead of time.RFC3339
so variable-length fractional seconds (micro- to nanoseconds) are accepted and
valid timestamps no longer return 0; keep the existing empty-string check and
error-handling behavior, only change the time format constant used in
time.Parse.
♻️ Duplicate comments (2)
core/providers/replicate/models.go (1)
69-82: Split owner/name for Replicate model responses.As noted in a previous review,
ReplicateModelResponseexpectsOwnerandNameas separate fields, butNamecurrently receivesmodelIDwhich containsowner/namecombined. This may cause inconsistencies with Replicate's expected API response shape.core/providers/replicate/images.go (1)
24-45: Version-ID models default toinput_imagesfield, potentially causing invalid payloads for Flux variants.As flagged in a previous review, when a model is passed as a 64-character version ID,
getInputImageFieldName()cannot match it against hardcoded model names and defaults toinput_images. However, FLUX 1.1 Pro requiresimage_promptand FLUX Fill requiresimage. Consider allowing an explicit field override viaExtraParamsfor version ID scenarios.
🧹 Nitpick comments (4)
core/providers/replicate/models.go (1)
27-28: Redundant assignment ofmodelName.Line 28 reassigns
modelNameto the same value already set on line 23. This assignment inside theif deployment.CurrentRelease != nilblock is unnecessary.♻️ Suggested fix
// Extract information from current release if available if deployment.CurrentRelease != nil { - modelName = schemas.Ptr(deployment.Name) - // Parse created timestamp if deployment.CurrentRelease.CreatedAt != "" {core/providers/replicate/utils.go (2)
24-47: Consider using dynamic provider key instead of hardcodedschemas.Replicate.The
checkForErrorStatusfunction hardcodesschemas.Replicatein error construction. Based on learnings, when handling errors across providers, avoid hardcoding provider constants. For custom providers using Replicate as a base format, error messages should reflect the actual provider key.Consider accepting a
providerKey schemas.ModelProviderparameter and passing it toNewBifrostOperationError.
101-104: Hardcodedschemas.Replicatein error handling.Same as above, the error construction on lines 102 and 104 uses
schemas.Replicatedirectly. For consistency with custom provider support, consider passing the provider key as a parameter to this function.core/providers/replicate/types.go (1)
286-290: Prefer pointer helper forReplicateOutputTextpointer creation.Line 289 uses
&objectContentwhile the same function already usesschemas.Ptrfor string values. For consistency and to avoid raw address-of usage, consider using the helper here too.♻️ Suggested refactor
- mc.OutputObject = &objectContent + mc.OutputObject = schemas.Ptr(objectContent)
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
5cdf2ac to
38f3f93
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
framework/configstore/clientconfig.go (1)
395-456: Prefix env vs literal values inGenerateKeyHashto prevent collisions.Right now, env-sourced values and literal values can hash to the same bytes, which can cause false “unchanged” or “changed” detections. Add an
env:/val:prefix (mirrorsGenerateMCPClientHash) so hashes are unambiguous.🔧 Suggested fix
- if key.Value.IsFromEnv() { - hash.Write([]byte(key.Value.EnvVar)) - } else { - hash.Write([]byte(key.Value.Val)) - } + if key.Value.IsFromEnv() { + hash.Write([]byte("env:" + key.Value.EnvVar)) + } else { + hash.Write([]byte("val:" + key.Value.Val)) + }Based on learnings, apply the
env:/val:prefix rule in this file to avoid hash collisions.transports/bifrost-http/lib/config.go (1)
1724-1738: IncludeReplicateKeyConfigin legacy key hash comparisons.When
ConfigHashis missing,mergeProviderKeys/reconcileProviderKeysrebuild aschemas.KeywithoutReplicateKeyConfig. With the new field, this can yield hash mismatches and unintentionally overwrite DB keys during reconciliation.🔧 Suggested fix
// mergeProviderKeys → legacy dbKeyHash dbKeyHash, err := configstore.GenerateKeyHash(schemas.Key{ Name: dbKey.Name, Value: dbKey.Value, Models: dbKey.Models, Weight: dbKey.Weight, AzureKeyConfig: dbKey.AzureKeyConfig, VertexKeyConfig: dbKey.VertexKeyConfig, BedrockKeyConfig: dbKey.BedrockKeyConfig, + ReplicateKeyConfig: dbKey.ReplicateKeyConfig, }) // reconcileProviderKeys → legacy dbKeyHash dbKeyHash, err := configstore.GenerateKeyHash(schemas.Key{ Name: dbKey.Name, Value: dbKey.Value, Models: dbKey.Models, Weight: dbKey.Weight, AzureKeyConfig: dbKey.AzureKeyConfig, VertexKeyConfig: dbKey.VertexKeyConfig, BedrockKeyConfig: dbKey.BedrockKeyConfig, + ReplicateKeyConfig: dbKey.ReplicateKeyConfig, })
🤖 Fix all issues with AI agents
In `@core/internal/testutil/list_models.go`:
- Line 111: The top-of-block comment for RunListModelsPaginationTest has an
accidental double leading comment marker ("// // ..."); change it to a
single-line comment with one "//" so it reads "// RunListModelsPaginationTest
executes pagination test for list models" to match standard comment style and
correctly document the RunListModelsPaginationTest symbol.
🧹 Nitpick comments (6)
ui/app/workspace/providers/fragments/apiKeysFormFragment.tsx (1)
430-432: Consider clearing errors on change for consistency.The Azure deployments field clears form errors on change (line 272), but the Replicate field does not. This is a minor inconsistency that doesn't affect functionality but could improve UX by clearing stale validation errors as the user types.
Optional: Add error clearing for consistency
onChange={(e) => { + form.clearErrors("key.replicate_key_config.deployments"); // Store as string during editing to allow intermediate invalid states field.onChange(e.target.value); }}ui/lib/types/schemas.ts (1)
150-186: Consider validating Replicate deployments strings like other providers.
replicateKeyConfigSchemaaccepts any string today; aligning it with the Azure/Vertex/Bedrock deployments validation avoids invalid JSON/env strings slipping through.♻️ Suggested refinement
-export const replicateKeyConfigSchema = z.object({ - deployments: z.union([z.record(z.string(), z.string()), z.string()]).optional(), -}); +export const replicateKeyConfigSchema = z + .object({ + deployments: z.union([z.record(z.string(), z.string()), z.string()]).optional(), + }) + .refine( + (data) => { + if (!data.deployments) return true; + if (typeof data.deployments === "object") return true; + if (typeof data.deployments === "string") { + const trimmed = data.deployments.trim(); + if (trimmed === "") return true; + if (trimmed.startsWith("env.")) return true; + try { + const parsed = JSON.parse(trimmed); + return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed); + } catch { + return false; + } + } + return false; + }, + { + message: "Deployments must be a valid JSON object or an environment variable reference", + path: ["deployments"], + }, + );core/providers/replicate/files.go (1)
39-44: Use the pointer helper forExpiresAt.
Prefer the repo’s pointer helper over&for consistency.♻️ Proposed fix
- resp.ExpiresAt = &expiresAt + resp.ExpiresAt = schemas.Ptr(expiresAt) - resp.ExpiresAt = &expiresAt + resp.ExpiresAt = schemas.Ptr(expiresAt)Based on learnings, please prefer pointer helpers over
&for consistency.Also applies to: 76-81
core/providers/replicate/text.go (1)
113-117: Consider distinct finish reason for canceled status.Mapping
ReplicatePredictionStatusCanceledto"stop"loses the semantic distinction between successful completion and cancellation. This could be relevant for clients that need to know if the request was interrupted.💡 Suggested improvement
case ReplicatePredictionStatusCanceled: - finishReason = schemas.Ptr("stop") + finishReason = schemas.Ptr("cancelled")core/providers/replicate/responses.go (1)
18-18: Hardcoded model name check is fragile.The check
strings.HasPrefix(bifrostReq.Model, "openai/") && strings.Contains(bifrostReq.Model, "gpt-5-structured")is brittle and may break when new structured models are added. Consider extracting this to a helper function or using a more flexible pattern (e.g., a configurable list or suffix check).💡 Suggested approach
// Add a helper function for model capability detection func isStructuredResponsesModel(model string) bool { return strings.HasPrefix(model, "openai/") && strings.Contains(model, "gpt-5-structured") }This makes it easier to extend when new structured models are added.
core/providers/replicate/utils.go (1)
25-47: Add nil check for prediction parameter.
checkForErrorStatusdereferencesprediction.Statusandprediction.Errorwithout checking ifpredictionis nil. While callers may validate this, a defensive check would prevent potential panics.💡 Proposed fix
func checkForErrorStatus(prediction *ReplicatePredictionResponse) *schemas.BifrostError { + if prediction == nil { + return nil + } if prediction.Status == ReplicatePredictionStatusFailed {
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
3 times, most recently
from
bfe79b8 to
b5a13a3
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Fix all issues with AI agents
In `@core/providers/replicate/chat.go`:
- Around line 203-209: The webhook_events_filter is being read from
req.ExtraParams (which is nil) so it's ignored; update the extraction to read
from bifrostReq.Params.ExtraParams instead. In the block that already checks
bifrostReq.Params != nil && bifrostReq.Params.ExtraParams != nil, call
schemas.SafeExtractStringSlice on
bifrostReq.Params.ExtraParams["webhook_events_filter"] and assign the result to
req.WebhookEventsFilter (same as other extraction for webhook), using the
existing variables bifrostReq, req and the SafeExtractStringSlice helper to
locate the change.
In `@core/providers/replicate/images.go`:
- Around line 120-127: The code incorrectly reads webhook_events_filter from
request.ExtraParams instead of the incoming bifrost request; update the
extraction to read from bifrostReq.Params.ExtraParams (the same source used for
webhook) by changing the call site that invokes schemas.SafeExtractStringSlice
to pass bifrostReq.Params.ExtraParams["webhook_events_filter"] and then assign
the result to request.WebhookEventsFilter; touch the block around
bifrostReq.Params != nil && bifrostReq.Params.ExtraParams != nil where webhook
and webhook_events_filter are handled and ensure you reference
schemas.SafeExtractStringSlice and request.WebhookEventsFilter.
In `@core/providers/replicate/responses.go`:
- Around line 128-190: The ExtraParams from bifrostReq.Params are never
forwarded into the ReplicatePredictionRequest so webhook_events_filter is being
read from the nil req.ExtraParams; fix by propagating
bifrostReq.Params.ExtraParams into req.ExtraParams (or, if type conversion is
needed, copy/convert its map into req.ExtraParams) before reading
webhook_events_filter and webhook; update the code around
ReplicatePredictionRequest creation and the subsequent extraction block
(symbols: bifrostReq.Params.ExtraParams, req.ExtraParams,
ReplicatePredictionRequest, schemas.SafeExtractStringSlice,
schemas.SafeExtractStringPointer) so both webhook and webhook_events_filter are
read from the provided ExtraParams.
In `@core/providers/replicate/text.go`:
- Around line 76-82: The webhook_events_filter is being read from
req.ExtraParams (which is nil) instead of from the incoming bifrost request;
update the logic in the block that checks bifrostReq.Params.ExtraParams so that
you call schemas.SafeExtractStringSlice on
bifrostReq.Params.ExtraParams["webhook_events_filter"] (similar to how webhook
is extracted from bifrostReq.Params.ExtraParams) and assign the result to
req.WebhookEventsFilter; keep the existing SafeExtractStringPointer use for
webhook and ensure you reference bifrostReq.Params.ExtraParams for both
extra-param keys.
🧹 Nitpick comments (4)
core/providers/replicate/files.go (2)
40-45: Preferschemas.Ptr()over address-of operator for pointer creation.For consistency with codebase conventions, use
schemas.Ptr(expiresAt)instead of&expiresAt.♻️ Suggested fix
// Add ExpiresAt if present if r.ExpiresAt != "" { expiresAt := ParseReplicateTimestamp(r.ExpiresAt) if expiresAt > 0 { - resp.ExpiresAt = &expiresAt + resp.ExpiresAt = schemas.Ptr(expiresAt) } }Apply the same change at line 80.
Based on learnings: "In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically."
Also applies to: 76-81
22-56: Consider extracting shared logic between upload and retrieve response converters.Both
ToBifrostFileUploadResponseandToBifrostFileRetrieveResponsehave nearly identical implementations with only theRequestTypediffering. This could be DRYed up with a helper function, though it's acceptable as-is given the type differences.Also applies to: 59-93
core/providers/replicate/utils.go (1)
55-63: Minor simplification possible forparsePreferHeader.The logic can be simplified by returning the prefix check result directly.
♻️ Suggested simplification
func parsePreferHeader(extraHeaders map[string]string) bool { if preferValue, exists := extraHeaders["Prefer"]; exists { - if strings.HasPrefix(preferValue, "wait") { - return true - } - return false + return strings.HasPrefix(preferValue, "wait") } return false }core/providers/replicate/types.go (1)
465-481:ToBifrostPredictionStatusis effectively a no-op.The function returns the same string values as the input constants. Unless there's a planned divergence between Replicate and Bifrost status values, this could be simplified to a direct cast.
♻️ Suggested simplification
// ToBifrostPredictionStatus converts Replicate status to Bifrost status func ToBifrostPredictionStatus(status ReplicatePredictionStatus) string { - switch status { - case ReplicatePredictionStatusStarting: - return "starting" - case ReplicatePredictionStatusProcessing: - return "processing" - case ReplicatePredictionStatusSucceeded: - return "succeeded" - case ReplicatePredictionStatusFailed: - return "failed" - case ReplicatePredictionStatusCanceled: - return "canceled" - default: - return string(status) - } + return string(status) }
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
b5a13a3 to
8a8c9d6
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
transports/bifrost-http/handlers/inference.go (1)
2011-2108:⚠️ Potential issue | 🟡 MinorDon’t drop additional
image_urlentries for variations.
With URL inputs, only the first URL is used and any additional URLs are ignored (and an empty first entry can cause a false validation failure). Normalize non-empty URLs and forward the remaining ones (e.g., via ExtraParams) so multi-image URL inputs aren’t silently lost.🔧 Suggested fix
// Check for image URLs if imageURLValues := form.Value["image_url[]"]; len(imageURLValues) > 0 { imageURLs = imageURLValues } else if imageURLSingle := form.Value["image_url"]; len(imageURLSingle) > 0 { imageURLs = imageURLSingle } + + // Normalize URL inputs to skip empty entries + nonEmptyImageURLs := make([]string, 0, len(imageURLs)) + for _, url := range imageURLs { + if url != "" { + nonEmptyImageURLs = append(nonEmptyImageURLs, url) + } + } // Validate at least one image source is provided - if len(imageFiles) == 0 && len(imageURLs) == 0 { + if len(imageFiles) == 0 && len(nonEmptyImageURLs) == 0 { SendError(ctx, fasthttp.StatusBadRequest, "At least one image (file or URL) is required") return } @@ - } else if len(imageURLs) > 0 && imageURLs[0] != "" { + } else if len(nonEmptyImageURLs) > 0 { // Use URL-based input req.ImageVariationInput = &schemas.ImageVariationInput{ Image: schemas.ImageInput{ - URL: imageURLs[0], + URL: nonEmptyImageURLs[0], }, } @@ // Store additional images (after the first one) in ExtraParams for providers that support multiple images if len(images) > 1 { req.ImageVariationParameters.ExtraParams["images"] = images[1:] } + if len(nonEmptyImageURLs) > 1 { + req.ImageVariationParameters.ExtraParams["image_urls"] = nonEmptyImageURLs[1:] + }
🤖 Fix all issues with AI agents
In `@core/providers/replicate/chat.go`:
- Around line 145-154: The current branching checks
strings.HasPrefix(bifrostReq.Model, string(schemas.OpenAI)) before testing
isGPT5Structured, which makes the isGPT5Structured branch unreachable and
prevents setting input.MaxOutputTokens; update the conditional order in the
block that handles params.MaxCompletionTokens so you first check
isGPT5Structured and assign input.MaxOutputTokens = params.MaxCompletionTokens,
then else-if for strings.HasPrefix(bifrostReq.Model, string(schemas.OpenAI)) to
set input.MaxCompletionTokens, and finally the fallback that sets
input.MaxTokens.
In `@docs/providers/supported-providers/replicate.mdx`:
- Around line 20-28: The "List Models" row in the operations table currently
shows the endpoint `/v1/deployments` but the implementation and section `#6` use
`/v1/models`; update the table so the "List Models" row endpoint matches the
implementation by changing `/v1/deployments` to `/v1/models` in the table (refer
to the "List Models" row in the operations table).
🧹 Nitpick comments (5)
core/providers/replicate/errors.go (1)
12-34: Prefer pointer helpers forStatusCode.Using
&statusCodediverges from the repo’s pointer-helper convention;schemas.Ptr(statusCode)keeps consistency and avoids taking the address of a local.♻️ Suggested change
- StatusCode: &statusCode, + StatusCode: schemas.Ptr(statusCode), @@ - StatusCode: &statusCode, + StatusCode: schemas.Ptr(statusCode),Based on learnings: prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically.
core/internal/llmtests/account.go (1)
1279-1306: Enable Replicate text completion scenarios to exercise new support.This PR adds text completion support and
TextModelis set, butTextCompletion/TextCompletionStreamremain disabled, so the new path isn’t covered in the stack’s test matrix.✅ Suggested change
- TextCompletion: false, // Not typical + TextCompletion: true, + TextCompletionStream: true,core/providers/replicate/responses.go (1)
201-250: Preferbifrost.Ptrover address-of for new pointers.
There are several new&valueusages (e.g.,completedAt,joined, message type/role) — please align with the repo’s pointer helper for consistency.♻️ Example adjustment
import ( "fmt" "strings" "time" + bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" ) @@ - completedAt = &completed + completedAt = bifrost.Ptr(completed)Based on learnings: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically.
framework/configstore/rdb.go (1)
269-327: Preferbifrost.Ptrfor Weight pointer creation (apply consistently in other key upserts).♻️ Suggested refactor
- Weight: &key.Weight, + Weight: bifrost.Ptr(key.Weight),Based on learnings, prefer using bifrost.Ptr() instead of the address operator for pointer creation.
core/providers/replicate/images.go (1)
225-253: Preferbifrost.Ptrfor prompt/image pointers (apply across image edit/generation).♻️ Suggested refactor
- input := &ReplicatePredictionRequestInput{ - Prompt: &bifrostReq.Input.Prompt, - } + input := &ReplicatePredictionRequestInput{ + Prompt: bifrost.Ptr(bifrostReq.Input.Prompt), + }Based on learnings, prefer using bifrost.Ptr() instead of the address operator for pointer creation.
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
8a8c9d6 to
e3f9ed5
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
transports/bifrost-http/handlers/inference.go (2)
1732-1794:⚠️ Potential issue | 🟡 MinorPrevent image_url fields from leaking into ExtraParams.
Now thatimage_url/image_url[]are first-class inputs, they should be treated as known fields so they don’t get forwarded to providers via ExtraParams.🛠️ Suggested update
var imageEditParamsKnownFields = map[string]bool{ "image": true, "image[]": true, + "image_url": true, + "image_url[]": true, "mask": true, // ... } var imageVariationParamsKnownFields = map[string]bool{ "image": true, "image[]": true, + "image_url": true, + "image_url[]": true, "n": true, // ... }
2011-2073:⚠️ Potential issue | 🟡 MinorHandle empty/multiple image_url values for variations.
OnlyimageURLs[0]is used today; if it’s empty, valid later URLs are ignored, and additional URLs aren’t forwarded. Consider filtering empties and passing remaining URLs via ExtraParams (or another key aligned with provider expectations).🛠️ One possible adjustment
// After reading imageURLValues / imageURLSingle +filtered := make([]string, 0, len(imageURLs)) +for _, url := range imageURLs { + if url != "" { + filtered = append(filtered, url) + } +} +imageURLs = filtered -} else if len(imageURLs) > 0 && imageURLs[0] != "" { +} else if len(imageURLs) > 0 { // After ExtraParams init +if len(imageURLs) > 1 { + req.ImageVariationParameters.ExtraParams["image_urls"] = imageURLs[1:] +}
🤖 Fix all issues with AI agents
In `@core/providers/replicate/chat.go`:
- Around line 31-125: The prompt/system_prompt extraction runs unconditionally
after setting input.Messages/InputItemList for OpenAI models, causing both
messages and prompt fields to be sent; change the flow so that the entire block
that extracts systemPrompt, conversationParts, imageInput and sets
input.Prompt/input.SystemPrompt is executed only when the OpenAI branch did not
run—i.e., attach an else to the existing condition (the if using
len(bifrostReq.Input) && strings.HasPrefix(bifrostReq.Model,
string(schemas.OpenAI))) and move the code that builds
systemPrompt/conversationParts and sets input.Prompt/input.SystemPrompt into
that else branch to enforce mutual exclusivity between
input.Messages/InputItemList and input.Prompt/input.SystemPrompt.
In `@core/providers/replicate/responses.go`:
- Around line 18-36: The gpt-5-structured branch drops provider-specific
ExtraParams; update the handling inside the strings.HasPrefix(...) &&
strings.Contains(...) block to copy bifrostReq.Params.ExtraParams into the
structured request input (e.g., into input.ExtraParams or the equivalent field
on the request object) when bifrostReq.Params != nil and
bifrostReq.Params.ExtraParams != nil so provider-specific settings are
preserved; locate the code manipulating
input.Instructions/Tools/MaxOutputTokens/JsonSchema and add a similar
conditional to assign ExtraParams from bifrostReq.Params.ExtraParams to the
target input structure.
In `@core/providers/replicate/utils.go`:
- Around line 192-207: The docstring for buildPredictionURL incorrectly states
it returns a bool along with the URL; update it to match the actual function
signature (which only returns the URL string) or change the function to also
return the boolean if intended. Locate the buildPredictionURL function and
either remove the mention of the bool from the comment (adjust the first line
and description) or modify the function signature to return (url string,
includeVersion bool) and update all call sites that use buildPredictionURL (and
related helpers like isVersionID and providerUtils.GetRequestPath) to handle the
additional boolean accordingly.
In `@framework/configstore/rdb.go`:
- Around line 319-325: The persisted ReplicateSigningSecret is being stored with
.Redacted(), which masks literal secrets and breaks signature generation; remove
the .Redacted() call wherever the code assigns to dbKey.ReplicateSigningSecret
(e.g., change assignments like dbKey.ReplicateSigningSecret =
key.ReplicateKeyConfig.SigningSecret.Redacted() to store the raw key value
instead), and do the same for the two other spots indicated in the review; keep
handling of env.* references as-is (they are resolved via
os.LookupEnv()/Scan/Value/GetValue), and reserve Redacted() only for display
contexts rather than persistence.
🧹 Nitpick comments (9)
core/schemas/images.go (1)
133-141: Consider validating mutually exclusive image inputs.Line 139-Line 140 now allow both
imagebytes andurlon the same item. If not already enforced upstream, consider validation to ensure exactly one is set to avoid ambiguous provider mappings.core/providers/replicate/types.go (1)
110-176: Reduce drift risk between struct fields andknownFields.Line 128-Line 160 requires manual sync with the struct. Consider a small test or shared list to prevent omissions when adding fields in the future.
core/providers/replicate/utils.go (1)
52-87: Normalize Prefer header handling to avoid case-sensitive misses.
extraHeadersuses raw keys; HTTP headers are case-insensitive, so a lowercasepreferkey or mixed-case value can bypass sync detection and stripping. Consider normalizing keys/values defensively.♻️ Suggested normalization
func parsePreferHeader(extraHeaders map[string]string) bool { - if preferValue, exists := extraHeaders["Prefer"]; exists { - if strings.HasPrefix(preferValue, "wait") { - return true - } - return false - } - return false + for key, value := range extraHeaders { + if strings.EqualFold(key, "Prefer") { + v := strings.ToLower(strings.TrimSpace(value)) + return strings.HasPrefix(v, "wait") + } + } + return false } @@ func stripPreferHeader(extraHeaders map[string]string) map[string]string { if extraHeaders == nil { return nil } - // Check if Prefer header exists - if _, exists := extraHeaders["Prefer"]; !exists { - // No Prefer header, return original map - return extraHeaders - } - - // Create new map without Prefer header - filtered := make(map[string]string, len(extraHeaders)-1) + // Create new map without Prefer header (case-insensitive) + filtered := make(map[string]string, len(extraHeaders)) for key, value := range extraHeaders { - if key != "Prefer" { + if !strings.EqualFold(key, "Prefer") { filtered[key] = value } } - return filtered + if len(filtered) == len(extraHeaders) { + return extraHeaders + } + return filtered }core/providers/replicate/chat.go (2)
21-24: Consider removing redundant nil guards here.
Core-layer validation already rejects nil request/input; if this helper must defend against external callers, add a brief comment explaining the exception.Based on learnings: Nil validation for request objects should occur at the Bifrost core layer (core/bifrost.go). Provider-layer methods can assume the request has been validated and should not include redundant nil checks. Apply this pattern across all provider files under core/providers, and keep nil checks only where directly arising from external boundaries. If a provider method must guard against unexpected nil input due to external calls, add a focused defensive check with a clear comment.
103-120: Usebifrost.Ptrfor new pointer values for consistency.
Several new pointers are created via&(systemPrompt, prompt, version, joined output, finishReason). Align with the repo convention by usingbifrost.Ptr(...)and add the import where needed.Based on learnings: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically. Apply this consistently across all code paths, including test utilities, to improve consistency and readability. Replace occurrences of &value where a *T is expected with bifrost.Ptr(value) (or an equivalent call) and ensure the function is in scope and used correctly for the target pointer type.♻️ Example change
- reason := "stop" - finishReason = &reason + finishReason = bifrost.Ptr("stop")Also applies to: 199-201, 244-275
core/providers/replicate/files.go (1)
39-45: Usebifrost.Ptrfor ExpiresAt pointers.
Align with the repo’s pointer-creation convention and avoid raw&usage in these new response builders.Based on learnings: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically. Apply this consistently across all code paths, including test utilities, to improve consistency and readability. Replace occurrences of &value where a *T is expected with bifrost.Ptr(value) (or an equivalent call) and ensure the function is in scope and used correctly for the target pointer type.♻️ Example change
- resp.ExpiresAt = &expiresAt + resp.ExpiresAt = bifrost.Ptr(expiresAt)Also applies to: 76-82
framework/configstore/rdb.go (1)
270-277: Preferbifrost.Ptrfor new weight pointers.
These hunks now touch weight pointer creation; please align with the repo convention and usebifrost.Ptr(key.Weight)(apply across the same struct literals).Based on learnings: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically. Apply this consistently across all code paths, including test utilities, to improve consistency and readability. Replace occurrences of &value where a *T is expected with bifrost.Ptr(value) (or an equivalent call) and ensure the function is in scope and used correctly for the target pointer type.♻️ Example change
- Weight: &key.Weight, + Weight: bifrost.Ptr(key.Weight),Also applies to: 432-439, 555-562, 679-684, 721-726
core/providers/replicate/images.go (1)
35-59: Usebifrost.Ptrfor new pointer values in request builders.
These new request constructors take addresses of fields/slice elements; align with the repo convention and usebifrost.Ptr(...)instead (add the import as needed).Based on learnings: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically. Apply this consistently across all code paths, including test utilities, to improve consistency and readability. Replace occurrences of &value where a *T is expected with bifrost.Ptr(value) (or an equivalent call) and ensure the function is in scope and used correctly for the target pointer type.♻️ Example change
- Prompt: &bifrostReq.Input.Prompt, + Prompt: bifrost.Ptr(bifrostReq.Input.Prompt),Also applies to: 116-118, 225-254
core/providers/replicate/replicate_test.go (1)
33-36: Consider making the file expiry relative to now.
The hard‑coded Unix timestamp will eventually expire; deriving it fromtime.Now()keeps the integration test valid long‑term.🧪 Example change
import ( "os" "strings" "testing" + "time" @@ FileExtraParams: map[string]interface{}{ "owner": os.Getenv("REPLICATE_OWNER"), - "expiry": 1830297599, + "expiry": time.Now().Add(24 * time.Hour).Unix(), },
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
e3f9ed5 to
1ca3b75
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
framework/configstore/clientconfig.go (1)
405-466:⚠️ Potential issue | 🟠 MajorPrefix env vs literal values in
GenerateKeyHashto avoid collisions.Right now,
key.Valuehashing can collide when an env var name equals a literal value. Since this function is being modified, please addenv:/val:prefixes to match the established hashing scheme.Suggested fix
- if key.Value.IsFromEnv() { - hash.Write([]byte(key.Value.EnvVar)) - } else { - hash.Write([]byte(key.Value.Val)) - } + if key.Value.IsFromEnv() { + hash.Write([]byte("env:" + key.Value.EnvVar)) + } else { + hash.Write([]byte("val:" + key.Value.Val)) + }Based on learnings: In framework/configstore/clientconfig.go, when implementing GenerateKeyHash, prefix hashed keys with distinct identifiers: use 'env:' for environment-sourced values and 'val:' for literal values, mirroring the existing 'GenerateMCPClientHash' pattern to prevent hash collisions.
transports/bifrost-http/handlers/inference.go (2)
1732-1794:⚠️ Potential issue | 🟡 MinorTreat
image_urlas a known field to avoid leaking it into ExtraParams.Now that URLs are first-class inputs, they shouldn’t be forwarded again via ExtraParams.
🔧 Suggested fix
var imageEditParamsKnownFields = map[string]bool{ "model": true, "prompt": true, "fallbacks": true, "image": true, "image[]": true, + "image_url": true, + "image_url[]": true, "mask": true, // ... } var imageVariationParamsKnownFields = map[string]bool{ "model": true, "fallbacks": true, "image": true, "image[]": true, + "image_url": true, + "image_url[]": true, "n": true, // ... }
2011-2108:⚠️ Potential issue | 🟡 MinorPreserve additional URL images in variations.
Right now only extra file images are stored in
ExtraParams["images"]; additional URLs are dropped.🛠️ Suggested fix
// Store additional images (after the first one) in ExtraParams for providers that support multiple images if len(images) > 1 { req.ImageVariationParameters.ExtraParams["images"] = images[1:] } +if len(imageURLs) > 1 { + req.ImageVariationParameters.ExtraParams["image_urls"] = imageURLs[1:] +}
🤖 Fix all issues with AI agents
In `@docs/providers/supported-providers/replicate.mdx`:
- Around line 286-291: The table mapping "Replicate Status" to "Responses
Status" uses inconsistent spelling for canceled/cancelled: update the row where
Replicate has `canceled` and Responses has `cancelled` so both use the same
variant (choose either `canceled` or `cancelled`) to make the
"canceled"/"cancelled" cell pair consistent with the other rows in the table.
In `@framework/configstore/rdb.go`:
- Around line 269-277: The loop is taking the address of the range variable's
Weight (e.g., &key.Weight) which captures the reused range variable and causes
all pointers to point to the last iteration; update each construction of
tables.TableKey (and any similar structs where Weight is set) to allocate a
fresh pointer per iteration by replacing &key.Weight with
bifrost.Ptr(key.Weight) so each dbKey.Weight is a distinct pointer (apply the
same change in the other blocks that build tables.TableKey instances).
🧹 Nitpick comments (5)
core/providers/replicate/files.go (1)
58-93: Consider extracting shared logic for response conversion.Both
ToBifrostFileUploadResponseandToBifrostFileRetrieveResponseshare nearly identical logic forExpiresAthandling and raw request/response embedding. While acceptable for clarity, you could extract a helper to reduce duplication if this pattern extends to more file operations.core/providers/replicate/responses.go (1)
214-219: Usebifrost.Ptrinstead of&for pointer creation.Keeps pointer creation consistent with repo style.
♻️ Suggested refactor
import ( "fmt" "strings" "time" + bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" ) @@ - completedAt = &completed + completedAt = bifrost.Ptr(completed)Based on learnings: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically.
core/providers/replicate/chat.go (1)
21-24: Drop provider-layer nil guards and rely on core validation.
Core/bifrost already validates request objects, so this guard duplicates logic and forces tests to encode provider-specific nil semantics. If you still need a defensive guard for external entry points, add a brief comment explaining why. Also update the NilRequest/NilInput cases incore/providers/replicate/replicate_test.goif you remove the guard.♻️ Suggested refactor
func ToReplicateChatRequest(bifrostReq *schemas.BifrostChatRequest) (*ReplicatePredictionRequest, error) { - if bifrostReq == nil || bifrostReq.Input == nil { - return nil, fmt.Errorf("bifrost request is nil or input is nil") - } + // Request validation is handled in core/bifrost.go. // Build the input from messages input := &ReplicatePredictionRequestInput{}Based on learnings: Nil validation for request objects should occur at the Bifrost core layer (core/bifrost.go).
core/providers/replicate/images.go (1)
30-33: Rely on core validation instead of returning nil for nil requests.
Core/bifrost already validates image requests; dropping these guards keeps provider converters consistent and simplifies tests (e.g., NilRequest/NilInput cases inreplicate_test.go). If you keep the guard for external entry points, add a short justification comment.♻️ Suggested refactor
- if bifrostReq == nil || bifrostReq.Input == nil { - return nil - } + // Request validation is handled in core/bifrost.go.Based on learnings: Nil validation for request objects should occur at the Bifrost core layer (core/bifrost.go).
Also applies to: 220-223
core/providers/replicate/types.go (1)
284-289: Consider stricter validation for object output shape.The object unmarshal branch will succeed for any JSON object, including ones that don't have
response_idortextfields. This could silently produce aReplicateOutputTextwith all nil fields for unexpected object structures.If Replicate might return other object shapes in the future, this is fine. Otherwise, consider validating that at least one expected field is present, or logging unexpected shapes for debugging.
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
1ca3b75 to
28950dd
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Fix all issues with AI agents
In `@core/providers/replicate/chat.go`:
- Around line 233-236: The response metadata currently hardcodes Provider:
schemas.Replicate inside schemas.BifrostResponseExtraFields; change it to use
the actual provider key variable (e.g., providerKey, p.Key, or similar) by
passing that provider identifier into the converter/build-response call and set
ExtraFields.Provider = providerKey instead of schemas.Replicate so custom
provider aliases are preserved (update the converter function signature if
needed to accept the providerKey and propagate it into
schemas.BifrostResponseExtraFields).
In `@core/providers/replicate/text.go`:
- Around line 95-102: The response metadata currently hardcodes
ExtraFields.Provider to schemas.Replicate when building the bifrostResponse
(schemas.BifrostTextCompletionResponse), which strips custom provider aliases;
update the converter that constructs bifrostResponse to accept and use the
actual provider key passed into the conversion (instead of schemas.Replicate),
mirroring the pattern used by the file converters—locate the converter that
builds bifrostResponse and replace the hardcoded Provider assignment with the
dynamic provider parameter so the ExtraFields.Provider reflects the real
provider alias.
In `@core/providers/replicate/utils.go`:
- Around line 143-167: The function parseDataURIImage mutates the input by
trimming the "data:" prefix before validating the presence of the comma, causing
malformed data URIs to return a truncated string; fix it by preserving the
original input or performing the comma-split/validation before calling
strings.TrimPrefix so that when the URI is malformed (no comma or wrong format)
the function returns the original input unchanged; locate parseDataURIImage and
either store the original dataURI in a local variable for safe return or move
the HasPrefix/parts check ahead of TrimPrefix, then return the original input
when validation fails.
In `@docs/providers/supported-providers/replicate.mdx`:
- Line 112: The sentence "It is the default mode of replicate predictions,
Bifrost automatically polls the prediction URL every 2 seconds until
completion." is a comma splice; change it by either replacing the comma with a
semicolon ("It is the default mode of replicate predictions; Bifrost
automatically polls...") or split into two sentences ("It is the default mode of
replicate predictions. Bifrost automatically polls...") to correct the
punctuation.
🧹 Nitpick comments (6)
core/providers/replicate/errors.go (1)
1-35: Usebifrost.Ptr()instead of&statusCodefor pointer creation.Per repository convention, prefer
bifrost.Ptr()over the address operator for creating pointers. Also, theschemasimport alias is redundant since the package is already namedschemas.Suggested fix
import ( + bifrost "github.com/maximhq/bifrost/core" "github.com/bytedance/sonic" - schemas "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/core/schemas" ) func parseReplicateError(body []byte, statusCode int) *schemas.BifrostError { var replicateErr ReplicateError if err := sonic.Unmarshal(body, &replicateErr); err == nil && replicateErr.Detail != "" { return &schemas.BifrostError{ IsBifrostError: false, - StatusCode: &statusCode, + StatusCode: bifrost.Ptr(statusCode), Error: &schemas.ErrorField{ Message: replicateErr.Detail, }, ExtraFields: schemas.BifrostErrorExtraFields{ Provider: schemas.Replicate, }, } } return &schemas.BifrostError{ IsBifrostError: false, - StatusCode: &statusCode, + StatusCode: bifrost.Ptr(statusCode), Error: &schemas.ErrorField{ Message: string(body), },Based on learnings: "In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically."
core/providers/replicate/files.go (1)
40-44: Use pointer helper forExpiresAtto match repo convention.This keeps pointer creation consistent with the rest of the codebase.
Based on learnings: “prefer using bifrost.Ptr() to create pointers instead of the address operator (&) … (or an equivalent call).”♻️ Proposed change
- if expiresAt > 0 { - resp.ExpiresAt = &expiresAt - } + if expiresAt > 0 { + resp.ExpiresAt = schemas.Ptr(expiresAt) + }- if expiresAt > 0 { - resp.ExpiresAt = &expiresAt - } + if expiresAt > 0 { + resp.ExpiresAt = schemas.Ptr(expiresAt) + }Also applies to: 77-81
core/providers/replicate/text.go (1)
10-13: Avoid redundant nil checks in provider converters.Core already validates request/input; keep this only if it’s a defensive boundary (with a brief comment).
Based on learnings: “Nil validation for request objects should occur at the Bifrost core layer (core/bifrost.go). Provider-layer methods … should not include redundant nil checks.”♻️ Proposed change
- if bifrostReq == nil || bifrostReq.Input == nil { - return nil, fmt.Errorf("bifrost request is nil or prompt is nil") - }core/providers/replicate/chat.go (1)
21-24: Avoid redundant nil checks in provider converters.Core already validates request/input; keep this only if it’s a defensive boundary (with a brief comment).
Based on learnings: “Nil validation for request objects should occur at the Bifrost core layer (core/bifrost.go). Provider-layer methods … should not include redundant nil checks.”♻️ Proposed change
- if bifrostReq == nil || bifrostReq.Input == nil { - return nil, fmt.Errorf("bifrost request is nil or input is nil") - }core/providers/replicate/types.go (2)
284-289: Useschemas.Ptr(objectContent)instead of&objectContentfor consistency.Line 273 already uses
schemas.Ptr(stringContent)for the string case. The object case should follow the same pattern.Proposed fix
var objectContent ReplicateOutputText if err := sonic.Unmarshal(data, &objectContent); err == nil { - mc.OutputObject = &objectContent + mc.OutputObject = schemas.Ptr(objectContent) return nil }Based on learnings: "prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically."
486-496:ExpiresAtshould be*stringto distinguish absent from empty.Other optional timestamp fields in this file (e.g.,
StartedAt,CompletedAtonReplicatePredictionResponse) use*string.ExpiresAtonReplicateFileResponseis a barestringwithomitempty, which means an empty string is indistinguishable from a missing field on unmarshal, and a zero-value""will be silently dropped on marshal. If the Replicate API can omit this field,*stringis more appropriate.Proposed fix
- ExpiresAt string `json:"expires_at,omitempty"` // ISO 8601 timestamp + ExpiresAt *string `json:"expires_at,omitempty"` // ISO 8601 timestamp
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
28950dd to
c1a2dad
Compare
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
7 times, most recently
from
75e9091 to
e6306cd
Compare
TejasGhatte
force-pushed
the
01-20-feat_replicate_provider
branch
from
e6306cd to
adab48a
Compare
Merge activity
|
2 participants
Add Replicate provider support
This PR adds support for the Replicate API, enabling access to Replicate's models through Bifrost. Replicate offers a wide range of open and closed-source models through a prediction-based API architecture.
Changes
Type of change
Affected areas
How to test
Set the
REPLICATE_API_KEYenvironment variable with your Replicate API key, then run the tests:You can also test with the HTTP gateway:
Breaking changes
Related issues
N/A
Security considerations
The implementation handles API keys securely and follows the same security patterns as other providers.
Checklist