diff --git a/cmd/doc.go b/cmd/doc.go index eada3fe..d923b78 100644 --- a/cmd/doc.go +++ b/cmd/doc.go @@ -10,10 +10,8 @@ import ( func Start(sharedLib bool) { if sharedLib || os.Getenv("GIN_MODE") == "release" { log.InitLogger(true) - } else { - log.InitLogger(false) } - + log.Infof("------------------- Starting Polaris ---------------------") dbClient, err := db.Open() if err != nil { diff --git a/go.mod b/go.mod index 3b1349b..c684112 100644 --- a/go.mod +++ b/go.mod @@ -55,6 +55,7 @@ require ( github.com/bits-and-blooms/bitset v1.2.2 // indirect github.com/blinkbean/dingtalk v1.1.3 // indirect github.com/bradfitz/iter v0.0.0-20191230175014-e8f45d346db8 // indirect + github.com/buger/jsonparser v1.1.1 // indirect github.com/cespare/xxhash v1.1.0 // indirect github.com/dustin/go-humanize v1.0.0 // indirect github.com/edsrzf/mmap-go v1.1.0 // indirect @@ -72,12 +73,14 @@ require ( github.com/gregdel/pushover v1.3.1 // indirect github.com/huandu/xstrings v1.3.2 // indirect github.com/kennygrant/sanitize v1.2.4 // indirect + github.com/mailru/easyjson v0.7.7 // indirect github.com/minio/sha256-simd v1.0.0 // indirect github.com/mr-tron/base58 v1.2.0 // indirect github.com/mschoch/smat v0.2.0 // indirect github.com/multiformats/go-multihash v0.2.3 // indirect github.com/multiformats/go-varint v0.0.6 // indirect github.com/ncruces/julianday v1.0.0 // indirect + github.com/openai/openai-go v0.1.0-beta.10 // indirect github.com/pion/datachannel v1.5.9 // indirect github.com/pion/dtls/v3 v3.0.3 // indirect github.com/pion/ice/v4 v4.0.2 // indirect @@ -105,6 +108,11 @@ require ( github.com/temoto/robotstxt v1.1.2 // indirect github.com/tetratelabs/wazero v1.8.0 // indirect github.com/tidwall/btree v1.6.0 // indirect + github.com/tidwall/gjson v1.18.0 // indirect + github.com/tidwall/match v1.1.1 // indirect + github.com/tidwall/pretty v1.2.1 // indirect + github.com/tidwall/sjson v1.2.5 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect github.com/wlynxg/anet v0.0.3 // indirect go.etcd.io/bbolt v1.3.6 // indirect go.opentelemetry.io/otel v1.28.0 // indirect @@ -180,6 +188,7 @@ require ( github.com/cyruzin/golang-tmdb v1.6.3 github.com/gin-gonic/gin v1.10.0 github.com/hekmon/transmissionrpc/v3 v3.0.0 + github.com/invopop/jsonschema v0.13.0 github.com/json-iterator/go v1.1.12 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect diff --git a/go.sum b/go.sum index 5a0e2b2..ab82235 100644 --- a/go.sum +++ b/go.sum @@ -120,6 +120,8 @@ github.com/bradfitz/iter v0.0.0-20140124041915-454541ec3da2/go.mod h1:PyRFw1Lt2w github.com/bradfitz/iter v0.0.0-20190303215204-33e6a9893b0c/go.mod h1:PyRFw1Lt2wKX4ZVSQ2mk+PeDa1rxyObEDlApuIsUKuo= github.com/bradfitz/iter v0.0.0-20191230175014-e8f45d346db8 h1:GKTyiRCL6zVf5wWaqKnf+7Qs6GbEPfd4iMOitWzXJx8= github.com/bradfitz/iter v0.0.0-20191230175014-e8f45d346db8/go.mod h1:spo1JLcs67NmW1aVLEgtA8Yy1elc+X8y5SRW1sFW4Og= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= github.com/bytedance/sonic v1.11.6 h1:oUp34TzMlL+OY1OUWxHqsdkgC/Zfc85zGqw9siXjrc0= github.com/bytedance/sonic v1.11.6/go.mod h1:LysEHSvpvDySVdC2f87zGWf6CIKJcAvqab1ZaiQtds4= github.com/bytedance/sonic/loader v0.1.1 h1:c+e5Pt1k/cy5wMveRDyk2X4B9hF4g7an8N3zCYjJFNM= @@ -280,8 +282,11 @@ github.com/huandu/xstrings v1.3.2 h1:L18LIDzqlW6xN2rEkpdV8+oL/IXWJ1APd+vsdYy4Wdw github.com/huandu/xstrings v1.3.2/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq4ovT0aE= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= +github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= github.com/jordan-wright/email v4.0.1-0.20210109023952-943e75fe5223+incompatible h1:jdpOPRN1zP63Td1hDQbZW73xKmzDvZHzVdNYxhnTMDA= github.com/jordan-wright/email v4.0.1-0.20210109023952-943e75fe5223+incompatible/go.mod h1:1c7szIrayyPPB/987hsnvNzLushdWf4o/79s3P08L8A= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCVDaaPEHmU= github.com/json-iterator/go v1.1.9/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= @@ -315,6 +320,8 @@ github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY= github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-runewidth v0.0.9 h1:Lm995f3rfxdpd6TSmuVCHVb/QhupuXlYr8sCI/QdE+0= @@ -358,6 +365,8 @@ github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6 github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/ginkgo v1.7.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= +github.com/openai/openai-go v0.1.0-beta.10 h1:CknhGXe8aXQMRuqg255PFnWzgRY9nEryMxoNIBBM9tU= +github.com/openai/openai-go v0.1.0-beta.10/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y= github.com/openzipkin/zipkin-go v0.1.6/go.mod h1:QgAqvLzwWbR/WpD4A3cGpPtJrZXNIiJc5AZX7/PBEpw= github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM= github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs= @@ -492,6 +501,16 @@ github.com/tetratelabs/wazero v1.8.0 h1:iEKu0d4c2Pd+QSRieYbnQC9yiFlMS9D+Jr0LsRmc github.com/tetratelabs/wazero v1.8.0/go.mod h1:yAI0XTsMBhREkM/YDAK/zNou3GoiAce1P6+rp/wQhjs= github.com/tidwall/btree v1.6.0 h1:LDZfKfQIBHGHWSwckhXI0RPSXzlo+KYdjK7FWSqOzzg= github.com/tidwall/btree v1.6.0/go.mod h1:twD9XRA5jj9VUQGELzDO4HPQTNJsoWWfYEL+EUQ2cKY= +github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= +github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= +github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= +github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= +github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= +github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= github.com/tinylib/msgp v1.0.2/go.mod h1:+d+yLhGm8mzTaHzB+wgMYrodPfmZrzkirds8fDWklFE= github.com/tinylib/msgp v1.1.0/go.mod h1:+d+yLhGm8mzTaHzB+wgMYrodPfmZrzkirds8fDWklFE= github.com/tinylib/msgp v1.1.2/go.mod h1:+d+yLhGm8mzTaHzB+wgMYrodPfmZrzkirds8fDWklFE= @@ -503,6 +522,8 @@ github.com/vmihailenco/msgpack/v4 v4.3.12/go.mod h1:gborTTJjAo/GWTqqRjrLCn9pgNN+ github.com/vmihailenco/tagparser v0.1.1/go.mod h1:OeAg3pn3UbLjkWt+rN9oFYB6u/cQgqMEUPoW2WPyhdI= github.com/willf/bitset v1.1.9/go.mod h1:RjeCKbqT1RxIR/KWY6phxZiaY1IyutSBfGjNPySAYV4= github.com/willf/bitset v1.1.10/go.mod h1:RjeCKbqT1RxIR/KWY6phxZiaY1IyutSBfGjNPySAYV4= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= github.com/wlynxg/anet v0.0.3 h1:PvR53psxFXstc12jelG6f1Lv4MWqE0tI76/hHGjh9rg= github.com/wlynxg/anet v0.0.3/go.mod h1:eay5PRQr7fIVAMbTbchTnO9gG65Hg/uYGdc7mguHxoA= github.com/yuin/goldmark v1.4.1/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k= diff --git a/log/log.go b/log/log.go index 5a816a7..98b0a4e 100644 --- a/log/log.go +++ b/log/log.go @@ -15,6 +15,11 @@ var atom zap.AtomicLevel const dataPath = "./data" + +func init() { + InitLogger(false) +} + func InitLogger(toFile bool) { atom = zap.NewAtomicLevel() atom.SetLevel(zap.DebugLevel) diff --git a/pkg/deepseek/deepseek.go b/pkg/deepseek/deepseek.go new file mode 100644 index 0000000..c0ef9ba --- /dev/null +++ b/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 +} diff --git a/pkg/deepseek/deepseek_test.go b/pkg/deepseek/deepseek_test.go new file mode 100644 index 0000000..1f54118 --- /dev/null +++ b/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) + } +}