fix: make single step forcing optional on cohere endpoint type

huggingface · Dec 17, 2024 · 6dd3ae0 · 6dd3ae0
1 parent 71aae58
commit 6dd3ae0
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/lib/server/endpoints/cohere/endpointCohere.ts b/src/lib/server/endpoints/cohere/endpointCohere.ts
@@ -15,12 +15,14 @@ export const endpointCohereParametersSchema = z.object({
 	apiKey: z.string().default(env.COHERE_API_TOKEN),
 	clientName: z.string().optional(),
 	raw: z.boolean().default(false),
+	forceSingleStep: z.boolean().default(true),
 });
 
 export async function endpointCohere(
 	input: z.input<typeof endpointCohereParametersSchema>
 ): Promise<Endpoint> {
-	const { apiKey, clientName, model, raw } = endpointCohereParametersSchema.parse(input);
+	const { apiKey, clientName, model, raw, forceSingleStep } =
+		endpointCohereParametersSchema.parse(input);
 
 	let cohere: CohereClient;
 
@@ -62,7 +64,7 @@ export async function endpointCohere(
 				});
 
 				stream = await cohere.chatStream({
-					forceSingleStep: true,
+					forceSingleStep,
 					message: prompt,
 					rawPrompting: true,
 					model: model.id ?? model.name,