feat: support OPENAI json_schema (langgenius#7258)

MatriQ · Aug 15, 2024 · 6ff7fd8 · 6ff7fd8
1 parent 5aa373d
commit 6ff7fd8
Show file tree

Hide file tree

Showing 6 changed files with 44 additions and 10 deletions.
diff --git a/api/core/model_runtime/entities/defaults.py b/api/core/model_runtime/entities/defaults.py
@@ -1,4 +1,3 @@
-
 from core.model_runtime.entities.model_entities import DefaultParameterName
 
 PARAMETER_RULE_TEMPLATE: dict[DefaultParameterName, dict] = {
@@ -94,5 +93,16 @@
         },
         'required': False,
         'options': ['JSON', 'XML'],
-    }
-}
+    },
+    DefaultParameterName.JSON_SCHEMA: {
+        'label': {
+            'en_US': 'JSON Schema',
+        },
+        'type': 'text',
+        'help': {
+            'en_US': 'Set a response json schema will ensure LLM to adhere it.',
+            'zh_Hans': '设置返回的json schema，llm将按照它返回',
+        },
+        'required': False,
+    },
+}
diff --git a/api/core/model_runtime/entities/model_entities.py b/api/core/model_runtime/entities/model_entities.py
@@ -95,6 +95,7 @@ class DefaultParameterName(Enum):
     FREQUENCY_PENALTY = "frequency_penalty"
     MAX_TOKENS = "max_tokens"
     RESPONSE_FORMAT = "response_format"
+    JSON_SCHEMA = "json_schema"
 
     @classmethod
     def value_of(cls, value: Any) -> 'DefaultParameterName':
@@ -118,6 +119,7 @@ class ParameterType(Enum):
     INT = "int"
     STRING = "string"
     BOOLEAN = "boolean"
+    TEXT = "text"
 
 
 class ModelPropertyKey(Enum):

diff --git a/api/core/model_runtime/model_providers/openai/llm/gpt-4o-2024-08-06.yaml b/api/core/model_runtime/model_providers/openai/llm/gpt-4o-2024-08-06.yaml
@@ -37,6 +37,9 @@ parameter_rules:
     options:
       - text
       - json_object
+      - json_schema
+  - name: json_schema
+    use_template: json_schema
 pricing:
   input: '2.50'
   output: '10.00'

diff --git a/api/core/model_runtime/model_providers/openai/llm/gpt-4o-mini.yaml b/api/core/model_runtime/model_providers/openai/llm/gpt-4o-mini.yaml
@@ -37,6 +37,9 @@ parameter_rules:
     options:
       - text
       - json_object
+      - json_schema
+  - name: json_schema
+    use_template: json_schema
 pricing:
   input: '0.15'
   output: '0.60'

diff --git a/api/core/model_runtime/model_providers/openai/llm/llm.py b/api/core/model_runtime/model_providers/openai/llm/llm.py
@@ -1,3 +1,4 @@
+import json
 import logging
 from collections.abc import Generator
 from typing import Optional, Union, cast
@@ -544,13 +545,18 @@ def _chat_generate(self, model: str, credentials: dict,
 
         response_format = model_parameters.get("response_format")
         if response_format:
-            if response_format == "json_object":
-                response_format = {"type": "json_object"}
+            if response_format == "json_schema":
+                json_schema = model_parameters.get("json_schema")
+                if not json_schema:
+                    raise ValueError("Must define JSON Schema when the response format is json_schema")
+                try:
+                    schema = json.loads(json_schema)
+                except:
+                    raise ValueError(f"not currect json_schema format: {json_schema}")
+                model_parameters.pop("json_schema")
+                model_parameters["response_format"] = {"type": "json_schema", "json_schema": schema}
             else:
-                response_format = {"type": "text"}
-
-            model_parameters["response_format"] = response_format
-
+                model_parameters["response_format"] = {"type": response_format}
 
         extra_model_kwargs = {}
 

diff --git a/...nents/header/account-setting/model-provider-page/model-parameter-modal/parameter-item.tsx b/...nents/header/account-setting/model-provider-page/model-parameter-modal/parameter-item.tsx
@@ -100,7 +100,7 @@ const ParameterItem: FC<ParameterItemProps> = ({
     handleInputChange(v === 1)
   }
 
-  const handleStringInputChange = (e: React.ChangeEvent<HTMLInputElement>) => {
+  const handleStringInputChange = (e: React.ChangeEvent<HTMLInputElement | HTMLTextAreaElement>) => {
     handleInputChange(e.target.value)
   }
 
@@ -190,6 +190,16 @@ const ParameterItem: FC<ParameterItemProps> = ({
       )
     }
 
+    if (parameterRule.type === 'text') {
+      return (
+        <textarea
+          className='w-full h-20 ml-4 px-1 rounded-lg bg-gray-100 outline-none text-[12px] text-gray-900'
+          value={renderValue as string}
+          onChange={handleStringInputChange}
+        />
+      )
+    }
+
     if (parameterRule.type === 'string' && !!parameterRule?.options?.length) {
       return (
         <SimpleSelect