test/chat/api/api.go

package api

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"github.com/wecisecode/util/merrs"
)

type GenerateRequestOptions struct {
	Seed int64 `json:"seed,omitempty"`
}

// 定义请求结构体
type GenerateRequest struct {
	Model   string                 `json:"model"`
	Prompt  string                 `json:"prompt"`
	Stream  bool                   `json:"stream"`
	Context []int64                `json:"context,omitempty"`
	Options GenerateRequestOptions `json:"options,omitempty"`
}

type CallFunction struct {
	Name      string         `json:"name"`
	Arguments map[string]any `json:"arguments,omitempty"`
}

type ChatToolCall struct {
	Function CallFunction
}

type ChatFunctionParameter struct {
	Type        string   `json:"type"`
	Description string   `json:"description,omitempty"`
	Enum        []string `json:"enum,omitempty"`
}

type ChatFunctionParameters struct {
	Type       string                           `json:"type"` // object
	Properties map[string]ChatFunctionParameter `json:"properties,omitempty"`
	Required   []string                         `json:"required,omitempty"`
}

type ChatFunction struct {
	Name        string                  `json:"name"`
	Description string                  `json:"description,omitempty"`
	Parameters  *ChatFunctionParameters `json:"parameters,omitempty"`
}

type ChatTool struct {
	Type     string        `json:"type"` // "function"
	Function *ChatFunction `json:"function"`
}

type ChatMessage struct {
	Role      string          `json:"role"`
	Content   string          `json:"content"`
	Images    []string        `json:"images,omitempty"`
	ToolCalls []*ChatToolCall `json:"tool_calls,omitempty"`
}

// 定义请求结构体
type ChatRequest struct {
	Model    string         `json:"model"`
	Messages []*ChatMessage `json:"messages"`
	Stream   bool           `json:"stream"`
	Tools    []*ChatTool    `json:"tools,omitempty"`
}

type ChatResponse struct {
	Role    string      `json:"role"`
	Content string      `json:"content"`
	Images  []string    `json:"images,omitempty"`
	Tools   []*ChatTool `json:"tool_calls,omitempty"`
}

// 定义响应结构体
type GenerateResponse struct {
	Response string  `json:"response,omitempty"`
	Context  []int64 `json:"context,omitempty"`
	Done     bool    `json:"done,omitempty"`
}

type result struct {
	Error    error
	Response *GenerateResponse
}

type Result <-chan *result
type ChanResult chan *result

func newResult() ChanResult {
	return make(ChanResult, 10)
}

func (ret ChanResult) Error(e error) Result {
	ret <- &result{Error: e}
	return (chan *result)(ret)
}

func (ret ChanResult) Response(v *GenerateResponse) Result {
	ret <- &result{Response: v}
	return (chan *result)(ret)
}

func (ret ChanResult) Result() Result {
	return (chan *result)(ret)
}

func (ret ChanResult) Close() {
	close(ret)
}

func Request(context []int64, msg string) Result {
	ret := newResult()
	go func() {
		defer ret.Close()

		// Ollama 服务地址
		url := "http://127.0.0.1:11434/api/generate"

		// 创建请求体
		requestData := GenerateRequest{
			Model:   "deepseek-r1:7b", // 使用的模型名称
			Prompt:  msg,              // 输入的提示
			Stream:  true,             // 流式响应
			Context: context,          // 上下文
			Options: GenerateRequestOptions{
				Seed: 54321,
			},
		}

		// 将结构体转换为 JSON
		jsonData, err := json.Marshal(requestData)
		if err != nil {
			ret.Error(merrs.New("JSON 编码错误:", err))
			return
		}

		// 创建 HTTP 请求
		resp, err := http.Post(url, "application/json", bytes.NewBuffer(jsonData))
		if err != nil {
			ret.Error(merrs.New("请求失败:", err))
			return
		}
		defer resp.Body.Close()

		// 检查状态码
		if resp.StatusCode != http.StatusOK {
			body, _ := io.ReadAll(resp.Body)
			ret.Error(merrs.New(fmt.Sprintf("错误响应: %s\n状态码: %d\n", body, resp.StatusCode)))
			return
		}

		// 流式处理响应
		scanner := bufio.NewScanner(resp.Body)
		for scanner.Scan() {
			var chunk GenerateResponse
			if err := json.Unmarshal(scanner.Bytes(), &chunk); err != nil {
				ret.Error(merrs.New("解析分块失败:", err))
				break
			}
			// fmt.Print(chunk.Response) // 逐块打印响应
			ret.Response(&chunk)
			if chunk.Done {
				break
			}
		}
	}()
	return ret.Result()
}