11import 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" ;
45import { BaseLLM } from "../index.js" ;
56
67class 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
85167export default Replicate ;
0 commit comments