From 0e58042fa69eccd309181b9dc574f6121a75be6e Mon Sep 17 00:00:00 2001
From: Yang Beining <35399433+Suchun-sv@users.noreply.github.com>
Date: Thu, 22 Aug 2024 12:33:35 +0800
Subject: [PATCH] Support Openai structure output api  (#feat 1214) (#1217)

Co-authored-by: Kent Dong <ch3cho@qq.com>
---
 plugins/wasm-go/extensions/ai-proxy/README.md          |  1 +
 plugins/wasm-go/extensions/ai-proxy/provider/model.go  |  1 +
 plugins/wasm-go/extensions/ai-proxy/provider/openai.go |  4 ++++
 .../wasm-go/extensions/ai-proxy/provider/provider.go   | 10 ++++++++++
 4 files changed, 16 insertions(+)

diff --git a/plugins/wasm-go/extensions/ai-proxy/README.md b/plugins/wasm-go/extensions/ai-proxy/README.md
index 3f655a992..54ffa05b8 100644
--- a/plugins/wasm-go/extensions/ai-proxy/README.md
+++ b/plugins/wasm-go/extensions/ai-proxy/README.md
@@ -52,6 +52,7 @@ OpenAI 所对应的 `type` 为 `openai`。它特有的配置字段如下:
 | 名称              | 数据类型 | 填写要求 | 默认值 | 描述                                                                          |
 |-------------------|----------|----------|--------|-------------------------------------------------------------------------------|
 | `openaiCustomUrl` | string   | 非必填   | -      | 基于OpenAI协议的自定义后端URL，例如: www.example.com/myai/v1/chat/completions |
+| `responseJsonSchema` | object | 非必填 | - | 预先定义OpenAI响应需满足的Json Schema, 注意目前仅特定的几种模型支持该用法|
 
 
 #### Azure OpenAI
diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/model.go b/plugins/wasm-go/extensions/ai-proxy/provider/model.go
index 136671659..6cc5ae445 100644
--- a/plugins/wasm-go/extensions/ai-proxy/provider/model.go
+++ b/plugins/wasm-go/extensions/ai-proxy/provider/model.go
@@ -31,6 +31,7 @@ type chatCompletionRequest struct {
 	ToolChoice       *toolChoice    `json:"tool_choice,omitempty"`
 	User             string         `json:"user,omitempty"`
 	Stop             []string       `json:"stop,omitempty"`
+	ResponseFormat   map[string]interface{} `json:"response_format,omitempty"`
 }
 
 type streamOptions struct {
diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/openai.go b/plugins/wasm-go/extensions/ai-proxy/provider/openai.go
index dc989bf0c..9f34932c1 100644
--- a/plugins/wasm-go/extensions/ai-proxy/provider/openai.go
+++ b/plugins/wasm-go/extensions/ai-proxy/provider/openai.go
@@ -89,6 +89,10 @@ func (m *openaiProvider) OnRequestBody(ctx wrapper.HttpContext, apiName ApiName,
 	if err := decodeChatCompletionRequest(body, request); err != nil {
 		return types.ActionContinue, err
 	}
+	if m.config.responseJsonSchema != nil {
+		log.Debugf("[ai-proxy] set response format to %s", m.config.responseJsonSchema)
+		request.ResponseFormat = m.config.responseJsonSchema
+	}
 	if request.Stream {
 		// For stream requests, we need to include usage in the response.
 		if request.StreamOptions == nil {
diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/provider.go b/plugins/wasm-go/extensions/ai-proxy/provider/provider.go
index abcf65991..3d1d0217f 100644
--- a/plugins/wasm-go/extensions/ai-proxy/provider/provider.go
+++ b/plugins/wasm-go/extensions/ai-proxy/provider/provider.go
@@ -181,6 +181,9 @@ type ProviderConfig struct {
 	// @Title zh-CN 翻译服务需指定的目标语种
 	// @Description zh-CN 翻译结果的语种，目前仅适用于DeepL服务。
 	targetLang string `required:"false" yaml:"targetLang" json:"targetLang"`
+	// @Title zh-CN  指定服务返回的响应需满足的JSON Schema
+	// @Description zh-CN 目前仅适用于OpenAI部分模型服务。参考：https://platform.openai.com/docs/guides/structured-outputs
+	responseJsonSchema map[string]interface{} `required:"false" yaml:"responseJsonSchema" json:"responseJsonSchema"`
 }
 
 func (c *ProviderConfig) FromJson(json gjson.Result) {
@@ -229,6 +232,13 @@ func (c *ProviderConfig) FromJson(json gjson.Result) {
 		}
 	}
 	c.targetLang = json.Get("targetLang").String()
+	
+	if schemaValue, ok := json.Get("responseJsonSchema").Value().(map[string]interface{}); ok {
+		c.responseJsonSchema = schemaValue
+	} else {
+		c.responseJsonSchema = nil 
+	}
+	
 }
 
 func (c *ProviderConfig) Validate() error {