Skip to content

Commit 1a2c6b9

Browse files
committed
Fixed the integration issues with the Replicate provider
1 parent 70d1bfa commit 1a2c6b9

File tree

3 files changed

+87
-4
lines changed

3 files changed

+87
-4
lines changed

core/llm/llms/Replicate.ts

Lines changed: 85 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import ReplicateClient from "replicate";
22

3-
import { CompletionOptions, LLMOptions } from "../../index.js";
3+
import { ChatMessage, CompletionOptions, LLMOptions } from "../../index.js";
4+
import { renderChatMessage } from "../../util/messageContent.js";
45
import { BaseLLM } from "../index.js";
56

67
class Replicate extends BaseLLM {
@@ -33,6 +34,7 @@ class Replicate extends BaseLLM {
3334
"tomasmcm/neural-chat-7b-v3-1:acb450496b49e19a1e410b50c574a34acacd54820bc36c19cbfe05148de2ba57",
3435
"deepseek-7b": "kcaverly/deepseek-coder-33b-instruct-gguf" as any,
3536
"phind-codellama-34b": "kcaverly/phind-codellama-34b-v2-gguf" as any,
37+
"claude-4-sonnet-latest": "anthropic/claude-4-sonnet" as any,
3638
};
3739

3840
static providerName = "replicate";
@@ -42,12 +44,55 @@ class Replicate extends BaseLLM {
4244
options: CompletionOptions,
4345
prompt: string,
4446
signal: AbortSignal,
45-
): [`${string}/${string}:${string}`, { input: any; signal: AbortSignal }] {
47+
): [`${string}/${string}:${string}`, { input: any }] {
4648
return [
4749
Replicate.MODEL_IDS[options.model] || (options.model as any),
4850
{
4951
input: { prompt, message: prompt },
50-
signal,
52+
},
53+
];
54+
}
55+
56+
private _convertChatArgs(
57+
options: CompletionOptions,
58+
messages: ChatMessage[],
59+
signal: AbortSignal,
60+
): [`${string}/${string}:${string}`, { input: any }] {
61+
let prompt = "";
62+
let system_prompt = "";
63+
64+
for (const message of messages) {
65+
const content =
66+
typeof message.content === "string"
67+
? message.content
68+
: renderChatMessage(message);
69+
70+
if (message.role === "system") {
71+
system_prompt += `System: ${content}\n\n`;
72+
} else if (message.role === "user") {
73+
prompt += `Human: ${content}\n\n`;
74+
} else if (message.role === "assistant") {
75+
prompt += `Assistant: ${content}\n\n`;
76+
}
77+
}
78+
79+
if (!prompt.endsWith("Assistant: ")) {
80+
prompt += "Assistant: ";
81+
}
82+
83+
// Construct the input
84+
const input: any = {
85+
prompt: prompt,
86+
system_prompt: system_prompt,
87+
max_tokens: options.maxTokens || 2048,
88+
extended_thinking: options.reasoning,
89+
thinking_budget_tokens: options.reasoningBudgetTokens || 1024,
90+
};
91+
92+
return [
93+
Replicate.MODEL_IDS[options.model] || (options.model as any),
94+
{
95+
input,
5196
},
5297
];
5398
}
@@ -80,6 +125,43 @@ class Replicate extends BaseLLM {
80125
}
81126
}
82127
}
128+
129+
protected async *_streamChat(
130+
messages: ChatMessage[],
131+
signal: AbortSignal,
132+
options: CompletionOptions,
133+
): AsyncGenerator<ChatMessage> {
134+
if (!this.apiKey || this.apiKey === "") {
135+
throw new Error("You need to use an API key");
136+
}
137+
138+
const [model, args] = this._convertChatArgs(options, messages, signal);
139+
140+
try {
141+
for await (const event of this._replicate.stream(model, args)) {
142+
if (event.event === "output") {
143+
yield {
144+
role: "assistant",
145+
content: event.data,
146+
};
147+
}
148+
}
149+
} catch (error) {
150+
if (error instanceof Error) {
151+
if (error.message.includes("authentication")) {
152+
throw new Error(
153+
"Replicate API authentication failed. Please check your API key",
154+
);
155+
}
156+
if (error.message.includes("model not found")) {
157+
throw new Error(
158+
`Model "${options.model}" not found on Replicate. Please check the model name or use another model.`,
159+
);
160+
}
161+
}
162+
throw error;
163+
}
164+
}
83165
}
84166

85167
export default Replicate;

gui/src/pages/AddNewModel/configs/models.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1155,7 +1155,7 @@ export const models: { [key: string]: ModelPackage } = {
11551155
title: "Claude 4 Sonnet",
11561156
apiKey: "",
11571157
},
1158-
providerOptions: ["anthropic"],
1158+
providerOptions: ["anthropic", "replicate"],
11591159
icon: "anthropic.png",
11601160
isOpenSource: false,
11611161
},

gui/src/pages/AddNewModel/configs/providers.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -725,6 +725,7 @@ Select the \`GPT-4o\` model below to complete your provider configuration, but n
725725
models.codeLlamaInstruct,
726726
models.wizardCoder,
727727
models.mistralOs,
728+
models.claude4Sonnet,
728729
],
729730
apiKeyUrl: "https://replicate.com/account/api-tokens",
730731
},

0 commit comments

Comments
 (0)