from volcenginesdkarkruntime import Ark
import os
from pydantic import BaseModel  # 用于定义响应解析模型

# 初始化方舟SDK客户端
client = Ark(
    # 从环境变量获取方舟API Key（需提前设置环境变量）
    api_key=os.environ.get("ARK_API_KEY"),
)

# 定义分步解析模型（对应业务场景的结构化响应）
class Step(BaseModel):
    explanation: str  # 步骤说明
    output: str       # 步骤计算结果

# 定义最终响应模型（包含分步过程和最终答案）
class MathResponse(BaseModel):
    steps: list[Step]       # 解题步骤列表
    final_answer: str       # 最终答案

# 调用方舟模型生成响应（自动解析为指定模型）
completion = client.beta.chat.completions.parse(
    model="doubao-seed-1-6-250615",  # 具体模型需替换为实际可用模型
    messages=[
        {"role": "system", "content": "你是一位数学辅导老师，需详细展示解题步骤"},
        {"role": "user", "content": "用中文解方程组：8x + 9 = 32 和 x + y = 1"}
    ],
    response_format=MathResponse,  # 指定响应解析模型
    extra_body={
         "thinking": {
             "type": "disabled" # 不使用深度思考能力
             # "type": "enabled" # 使用深度思考能力
         }
     }
)

# 提取解析后的结构化响应
resp = completion.choices[0].message.parsed

# 打印格式化的JSON结果
print(resp.model_dump_json(indent=2))

from openai import OpenAI
import os
from pydantic import BaseModel

client = OpenAI(
    # 从环境变量中获取方舟 API Key
    api_key=os.environ.get("ARK_API_KEY"),
    base_url = "https://ark.cn-beijing.volces.com/api/v3",
)

class Step(BaseModel):
    explanation: str
    output: str
class MathResponse(BaseModel):
    steps: list[Step]
    final_answer: str
    
completion = client.beta.chat.completions.parse(
    model = "doubao-seed-1-6-250615",  # 替换为您需要使用的模型
    messages = [
        {"role": "system", "content": "你是一位数学辅导老师。"},
        {"role": "user", "content": "使用中文解题: 8x + 9 = 32 and x + y = 1"},
    ],
    response_format=MathResponse,
    extra_body={
         "thinking": {
             "type": "disabled" # 不使用深度思考能力
             # "type": "enabled" # 使用深度思考能力
         }
     }
)
resp = completion.choices[0].message.parsed
# 打印 JSON 格式结果
print(resp.model_dump_json(indent=2))

package main

import (
    "context"
    "encoding/json"
    "fmt"
    "os"
    "github.com/invopop/jsonschema" // required go1.18+
    "github.com/volcengine/volcengine-go-sdk/service/arkruntime"
    "github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
    "github.com/volcengine/volcengine-go-sdk/volcengine"
)

// 定义分步解析模型（对应业务场景的结构化响应）
type Step struct {
    Explanation string `json:"explanation" jsonschema_description:"步骤说明"`
    Output      string `json:"output" jsonschema_description:"步骤计算结果"`
}

// 定义最终响应模型（包含分步过程和最终答案）
type MathResponse struct {
    Steps       []Step `json:"steps" jsonschema_description:"解题步骤列表"`
    FinalAnswer string `json:"final_answer" jsonschema_description:"最终答案"`
}

// 复用原有 Schema 生成函数（已优化返回类型）
func GenerateSchema[T any]() *jsonschema.Schema { // <-- 优化返回类型为具体 Schema 类型
    reflector := jsonschema.Reflector{
        AllowAdditionalProperties: false,
        DoNotReference:            true,
    }
    return reflector.Reflect(new(T)) // 使用 new(T) 避免空值问题
}

// 生成数学响应的 JSON Schema
var MathResponseSchema = GenerateSchema[MathResponse]()

func main() {
    client := arkruntime.NewClientWithApiKey(
        os.Getenv("ARK_API_KEY"),
        )
    ctx := context.Background()

    // 构造请求消息（包含 system 和 user 角色）
    messages := []*model.ChatCompletionMessage{
        {
            Role: model.ChatMessageRoleSystem,
            Content: &model.ChatCompletionMessageContent{
                StringValue: volcengine.String("你是一位数学辅导老师，需详细展示解题步骤"),
            },
        },
        {
            Role: model.ChatMessageRoleUser,
            Content: &model.ChatCompletionMessageContent{
                StringValue: volcengine.String("用中文解方程组：8x + 9 = 32 和 x + y = 1"),
            },
        },
    }

    // 配置响应格式（使用 MathResponse 的 Schema）
    schemaParam := model.ResponseFormatJSONSchemaJSONSchemaParam{
        Name:        "math_response", // 对应 Python 中的响应名称
        Description: "数学题解答的结构化响应",
        Schema:      MathResponseSchema,
        Strict:      true,
    }

    // 构造请求（包含 thinking 配置）
    req := model.CreateChatCompletionRequest{
        Model:    "doubao-seed-1-6-250615", // 需替换为实际可用模型
        Messages: messages,
        ResponseFormat: &model.ResponseFormat{
            Type:       model.ResponseFormatJSONSchema,
            JSONSchema: &schemaParam,
        },
        Thinking: &model.Thinking{
            // Type: model.ThinkingTypeDisabled, // 关闭深度思考能力
            Type: model.ThinkingTypeEnabled, //开启深度思考能力
        },
    }


    // 调用 API
    resp, err := client.CreateChatCompletion(ctx, req)
    if err != nil {
        fmt.Printf("structured output chat error: %v\n", err)
        return
    }


    // 解析结构化响应（关键差异：Go 需要手动反序列化）
    var mathResp MathResponse
    err = json.Unmarshal([]byte(*resp.Choices[0].Message.Content.StringValue), &mathResp)
    if err != nil {
        panic(err.Error())
    }


    // 打印格式化结果（使用 json.MarshalIndent 实现缩进）
    prettyJSON, _ := json.MarshalIndent(mathResp, "", "  ")
    fmt.Println(string(prettyJSON))
}

package com.example;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest.ChatCompletionRequestResponseFormat;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.model.completion.chat.ResponseFormatJSONSchemaJSONSchemaParam;
import com.volcengine.ark.runtime.service.ArkService;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

public class ChatCompletionsStructuredOutputsExamplev4 {
    static String apiKey = System.getenv("ARK_API_KEY");
    static ArkService service = ArkService.builder()
            .connectionPool(new ConnectionPool(5, 1, TimeUnit.SECONDS))
            .dispatcher(new Dispatcher())
            .apiKey(apiKey)
            .build();

    public static void main(String[] args) throws JsonProcessingException {
        ObjectMapper mapper = new ObjectMapper();

        // 构造消息列表（包含 system 和 user 角色）
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content("你是一位数学辅导老师，需详细展示解题步骤")
                .build());
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content("用中文解方程组：8x + 9 = 32 和 x + y = 1")
                .build());

        // 生成 JSON Schema
        String schemaJson = "{\n" +
                "  \"type\": \"object\",\n" +
                "  \"properties\": {\n" +
                "    \"steps\": {\n" +
                "      \"type\": \"array\",\n" +
                "      \"items\": {\n" +
                "        \"$ref\": \"#/definitions/Step\"\n" +
                "      }\n" +
                "    },\n" +
                "    \"finalAnswer\": {\n" +
                "      \"type\": \"string\"\n" +
                "    }\n" +
                "  },\n" +
                "  \"definitions\": {\n" +
                "    \"Step\": {\n" +
                "      \"type\": \"object\",\n" +
                "      \"properties\": {\n" +
                "        \"explanation\": {\n" +
                "          \"type\": \"string\"\n" +
                "        },\n" +
                "        \"output\": {\n" +
                "          \"type\": \"string\"\n" +
                "        }\n" +
                "      }\n" +
                "    }\n" +
                "  }\n" +
                "}";
        JsonNode schemaNode = mapper.readTree(schemaJson);

        // 配置响应格式
        ChatCompletionRequestResponseFormat responseFormat = new ChatCompletionRequestResponseFormat(
                "json_schema",
                new ResponseFormatJSONSchemaJSONSchemaParam(
                        "math_response",
                        "数学题解答的结构化响应",
                        schemaNode,
                        true
                )
        );

        // 构造请求（包含 thinking 配置）
        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model("doubao-seed-1-6-250615") // 替换为实际使用模型
                .messages(messages)
                .responseFormat(responseFormat)
                .thinking(new ChatCompletionRequest.ChatCompletionRequestThinking("disabled")) // 关闭模型深度思考能力
                .build();

        // 调用 API 并解析响应
        var response = service.createChatCompletion(request);
        if (!response.getChoices().isEmpty()) {
            String content = String.valueOf(response.getChoices().get(0).getMessage().getContent());
            JsonNode jsonNode = mapper.readTree(content);
            // 打印格式化结果
            System.out.println(mapper.writerWithDefaultPrettyPrinter().writeValueAsString(jsonNode));
        }

        service.shutdownExecutor();
    }
}