From f2b55ebabb7a4b0ca7fcfa726cb1b89e546a489c Mon Sep 17 00:00:00 2001
From: reanon <>
Date: Mon, 5 May 2025 13:59:53 +0200
Subject: [PATCH] small cohere fix

---
 src/proxy/cohere.ts                                      | 4 ++--
 .../request/preprocessors/validate-context-size.ts       | 9 ++++++---
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/proxy/cohere.ts b/src/proxy/cohere.ts
index 913d6ae..e4d6de0 100644
--- a/src/proxy/cohere.ts
+++ b/src/proxy/cohere.ts
@@ -156,8 +156,8 @@ function removeUnsupportedParameters(req: Request) {
     // Transform to Cohere's format if needed
     const jsonSchema = req.body.response_format.schema;
     req.body.response_format = {
-      type: "json_schema",
-      json_schema: jsonSchema
+      type: "json_object",
+      schema: jsonSchema
     };
   }
 
diff --git a/src/proxy/middleware/request/preprocessors/validate-context-size.ts b/src/proxy/middleware/request/preprocessors/validate-context-size.ts
index d7b781c..65c06f6 100644
--- a/src/proxy/middleware/request/preprocessors/validate-context-size.ts
+++ b/src/proxy/middleware/request/preprocessors/validate-context-size.ts
@@ -119,11 +119,14 @@ export const validateContextSize: RequestPreprocessor = async (req) => {
     modelMax = 64000;
   } else if (model.match(/^grok/)) {
     modelMax = 128000;
-  } else if (model.match(/command-a-03/)) {
+  } else if (model.match(/command-a-03-202[0-9]/)) {
     // Cohere's command-a-03 models have 256k context window
     modelMax = 256000;
-  } else if (model.match(/command/)) {
-    // Other Cohere command models have 128k context window
+  } else if (model.match(/command[\w-]*-03-202[0-9]/)) {
+    // Other Command models with -03- pattern (including r, r-plus) have 128k context window
+    modelMax = 128000;
+  } else if (model.match(/command/) || model.match(/cohere/)) {
+    // Default for all other Cohere models
     modelMax = 128000;
   } else if (model.match(/tral/)) {
     // catches mistral, mixtral, codestral, mathstral, etc. mistral models have