feat: support json_schema for ollama models (#11449)

langgenius · Dec 8, 2024 · 7e1184c · 7e1184c
1 parent 1ce51e5
commit 7e1184c
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/api/core/model_runtime/model_providers/ollama/llm/llm.py b/api/core/model_runtime/model_providers/ollama/llm/llm.py
@@ -181,9 +181,11 @@ def _generate(
         # prepare the payload for a simple ping to the model
         data = {"model": model, "stream": stream}
 
-        if "format" in model_parameters:
-            data["format"] = model_parameters["format"]
-            del model_parameters["format"]
+        if format_schema := model_parameters.pop("format", None):
+            try:
+                data["format"] = format_schema if format_schema == "json" else json.loads(format_schema)
+            except json.JSONDecodeError as e:
+                raise InvokeBadRequestError(f"Invalid format schema: {str(e)}")
 
         if "keep_alive" in model_parameters:
             data["keep_alive"] = model_parameters["keep_alive"]
@@ -733,12 +735,12 @@ def get_customizable_model_schema(self, model: str, credentials: dict) -> AIMode
                 ParameterRule(
                     name="format",
                     label=I18nObject(en_US="Format", zh_Hans="返回格式"),
-                    type=ParameterType.STRING,
+                    type=ParameterType.TEXT,
+                    default="json",
                     help=I18nObject(
-                        en_US="the format to return a response in. Currently the only accepted value is json.",
-                        zh_Hans="返回响应的格式。目前唯一接受的值是json。",
+                        en_US="the format to return a response in. Format can be `json` or a JSON schema.",
+                        zh_Hans="返回响应的格式。目前接受的值是字符串`json`或JSON schema.",
                     ),
-                    options=["json"],
                 ),
             ],
             pricing=PriceConfig(