WIP: ai deepseek integration

pkg/deepseek/deepseek.go

@@ -0,0 +1,171 @@
+package deepseek
+
+import (
+	"context"
+	"fmt"
+	"polaris/log"
+	"time"
+
+	"github.com/goccy/go-json"
+	"github.com/openai/openai-go"
+	"github.com/openai/openai-go/option"
+)
+
+func NewClient(apiKey string) *Client {
+	r := openai.NewClient(option.WithAPIKey(apiKey), option.WithBaseURL("https://api.deepseek.com"))
+	return &Client{openai: &r, model: "deepseek-chat"}
+}
+
+type Client struct {
+	openai *openai.Client
+	model  string
+}
+
+func (c *Client) Test() error {
+
+	question := `What computer ran the first neural network?
+	EXAMPLE JSON OUTPUT:
+	{
+		"origin": "The origin of the computer",
+		"full_name": "The name of the device model",
+		"legacy": "Its influence on the field of computing",
+		"notable_facts": "A few key facts about the computer
+	}
+	`
+
+	chat, err_ := c.openai.Chat.Completions.New(context.Background(), openai.ChatCompletionNewParams{
+		// ...
+		ResponseFormat: openai.ChatCompletionNewParamsResponseFormatUnion{
+			OfJSONSchema: &openai.ResponseFormatJSONSchemaParam{
+				Type: "json_object",
+			},
+		},
+		Messages: []openai.ChatCompletionMessageParamUnion{
+			openai.UserMessage(question),
+		},
+
+		// only certain models can perform structured outputs
+		Model: c.model,
+	})
+	if err_ != nil {
+		return err_
+	}
+
+	log.Infof("%+v", chat.Choices[0].Message.Content)
+	// extract into a well-typed struct
+	return nil
+}
+
+type Movies struct {
+	Movies []struct {
+		Name  string `json:"name"`
+		Match string `json:"match"`
+	} `json:"movies"`
+}
+
+type Tvs struct {
+	Tvs []Tv `json:"tvs"`
+}
+
+type Tv struct {
+	Name            string `json:"name"`
+	FileName        string `json:"file_name"`
+	Match           string `json:"match"`
+	Season          string `json:"season"`
+	StartEpisode    string `json:"start_episode"`
+	EndEpisode      string `json:"end_episode"`
+	Quality         string `json:"quality"`
+	IsSeasonPackage string `json:"is_season_package"`
+}
+
+func (c *Client) AssessMovieNames(movieName string, releaseYear int, torrentNames []string) (*Movies, error) {
+	q := `用户输入的是一些文件名称，你需要判断哪些文件可能属于 %d 年的电影 %s，哪些可能不是。
+
+	EXAMPLE JSON OUTPUT:
+	{
+		"movies": [
+			{
+				"name": "The name of the movie",
+				"match": "true or false"
+			},
+		]
+	}
+	`
+
+	q = fmt.Sprintf(q, releaseYear, movieName)
+	chat, err_ := c.openai.Chat.Completions.New(context.Background(), openai.ChatCompletionNewParams{
+		//
+		ResponseFormat: openai.ChatCompletionNewParamsResponseFormatUnion{
+			OfJSONSchema: &openai.ResponseFormatJSONSchemaParam{
+				Type: "json_object",
+			},
+		},
+		Messages: []openai.ChatCompletionMessageParamUnion{
+			openai.SystemMessage(q),
+			openai.UserMessage(fmt.Sprintf("文件名称: %v", torrentNames)),
+		},
+
+		// only certain models can perform structured outputs
+		Model: c.model,
+	})
+	if err_ != nil {
+		return nil, err_
+	}
+
+	log.Infof("%+v", chat.Choices[0].Message.Content)
+	var res Movies
+	if err := json.Unmarshal([]byte(chat.Choices[0].Message.Content), &res); err != nil {
+		return nil, err
+	}
+
+	// extract into a well-typed struct
+	return &res, nil
+}
+
+func (c *Client) AssessTvNames(tvName string, releaseYear int, torrentNames []string) ([]Tv, error) {
+	log.Debugf("deepseek tv name: %s, year: %d, torrent name len: %v", tvName, releaseYear, len(torrentNames))
+	t := time.Now()
+	defer func() {
+		log.Infof("deepseek assess tv name cost: %v", time.Since(t))
+	}()
+
+	q := `用户输入的是一些文件名称，你需要判断哪些文件可能属于 %d 年的电视剧 %s，哪些可能不是，并返回匹配的文件名。
+
+	EXAMPLE JSON OUTPUT:	
+	{
+		"tvs": [
+		"matched file name 1", "matched file name 2", ...
+		]	
+	}`
+	q = fmt.Sprintf(q, releaseYear, tvName)
+
+	var res []Tv
+
+	chat, err_ := c.openai.Chat.Completions.New(context.Background(), openai.ChatCompletionNewParams{
+		MaxTokens: openai.Opt(int64(4096)),
+		//...
+		ResponseFormat: openai.ChatCompletionNewParamsResponseFormatUnion{
+			OfJSONSchema: &openai.ResponseFormatJSONSchemaParam{
+				Type: "json_object",
+			},
+		},
+		Messages: []openai.ChatCompletionMessageParamUnion{
+			openai.SystemMessage(q),
+			openai.UserMessage(fmt.Sprintf("文件名称: %v", torrentNames)),
+		},
+
+		// only certain models can perform structured outputs
+		Model: c.model,
+	})
+	if err_ != nil {
+		return nil, err_
+	}
+	log.Infof("%+v", chat.Choices[0].Message.Content)
+	var tvs Tvs
+	if err := json.Unmarshal([]byte(chat.Choices[0].Message.Content), &tvs); err != nil {
+		return nil, err
+	}
+	res = append(res, tvs.Tvs...)
+
+	return res, nil
+}

pkg/deepseek/deepseek_test.go

@@ -0,0 +1,11 @@
+package deepseek
+
+import "testing"
+
+func TestDeepseek(t *testing.T) {
+	r := NewClient("sk-")
+	_, err := r.AssessTvNames("基督山伯爵", 2025, []string{"The Count of Monte Cristo 2024 S01 1080p WEB-DL DD 5.1 H.264-playWEB", "The Count of Monte Cristo 2024 S01E06-08 MULTi 1080p WEB H264-AMB3R"})
+	if err != nil {
+		t.Fatal(err)
+	}
+}