AI intergration WIP

server/core/torrent.go

@@ -6,6 +6,7 @@ import (
 	"polaris/ent"
 	"polaris/ent/media"
 	"polaris/log"
+	"polaris/pkg/gemini"
 	"polaris/pkg/metadata"
 	"polaris/pkg/prowlarr"
 	"polaris/pkg/torznab"
@@ -72,6 +73,58 @@ func names2Query(media *ent.Media) []string {
 	return names
 }
 
+func filterBasedOnGemini(cfg db.AIConfig, res []torznab.Result, names ...string) []torznab.Result {
+
+	var torrentNames []string
+	for _, r := range res {
+		torrentNames = append(torrentNames, r.Name)
+	}
+	g, err := gemini.NewClient(cfg.GeminiApiKey, cfg.GeminiModelName)
+	if err != nil {
+		log.Warnf("create gemini client: %v", err)
+		return res
+	}
+	resf, err := g.FilterTvOrMovies(torrentNames, names...)
+	if err != nil {
+		log.Warnf("filter with gemini: %v", err)
+		return res
+	}
+	var newRes []torznab.Result
+	for _, r := range res {
+		if slices.Contains(resf, r.Name) {
+			newRes = append(newRes, r)
+		}
+	}
+	return newRes
+}
+
+func filterBasedOnRules(res []torznab.Result, names ...string) []torznab.Result {
+	var filtered []torznab.Result
+	for _, r := range res {
+		meta := metadata.ParseTv(r.Name)
+		if meta.IsAcceptable(names...) {
+			filtered = append(filtered, r)
+		}
+	}
+	return filtered
+}
+
+func filterResourceNames(db1 *db.Client, res []torznab.Result, names ...string) []torznab.Result {
+	n1 := len(res)
+	cfg, err := db1.GetAIConfig()
+	if err != nil {
+		log.Warnf("get ai config: %v", err)
+	}
+	if cfg.Enabled {
+		res = filterBasedOnGemini(cfg, res, names...)
+	} else {
+		res = filterBasedOnRules(res, names...)
+	}
+	log.Infof("resource before name filtering length is %d, after filtering length is %d", n1, len(res))
+	return res
+
+}
+
 func SearchTvSeries(db1 *db.Client, param *SearchParam) ([]torznab.Result, error) {
 	series := db1.GetMediaDetails(param.MediaId)
 	if series == nil {
@@ -88,6 +141,8 @@ func SearchTvSeries(db1 *db.Client, param *SearchParam) ([]torznab.Result, error
 
 	res := searchWithTorznab(db1, prowlarr.TV, names...)
 
+	res = filterResourceNames(db1, res, names...)
+
 	var filtered []torznab.Result
 lo:
 	for _, r := range res {
@@ -249,6 +304,8 @@ func SearchMovie(db1 *db.Client, param *SearchParam) ([]torznab.Result, error) {
 	names := names2Query(movieDetail.Media)
 
 	res := searchWithTorznab(db1, prowlarr.Movie, names...)
+	res = filterResourceNames(db1, res, names...)
+
 	if movieDetail.Extras.IsJav() {
 		res1 := searchWithTorznab(db1, prowlarr.Movie, movieDetail.Extras.JavId)
 		res = append(res, res1...)

server/server.go

@@ -80,6 +80,8 @@ func (s *Server) Serve() error {
 		setting.POST("/prowlarr", HttpHandler(s.SaveProwlarrSetting))
 		setting.GET("/limiter", HttpHandler(s.GetSizeLimiter))
 		setting.POST("/limiter", HttpHandler(s.SetSizeLimiter))
+		setting.GET("/ai", HttpHandler(s.GetAIConfig))
+		setting.POST("/ai", HttpHandler(s.SetAIConfig))
 	}
 	activity := api.Group("/activity")
 	{

server/setting.go

@@ -363,3 +363,22 @@ func (s *Server) SetSizeLimiter(c *gin.Context) (interface{}, error) {
 	}
 	return "success", nil
 }
+
+func (s *Server) GetAIConfig(c *gin.Context) (interface{}, error) {
+	aiConfig, err := s.db.GetAIConfig()
+	if err != nil {
+		return nil, errors.Wrap(err, "db")
+	}
+	return aiConfig, nil
+}
+
+func (s *Server) SetAIConfig(c *gin.Context) (interface{}, error) {
+	var in db.AIConfig
+	if err := c.ShouldBindJSON(&in); err != nil {
+		return nil, err
+	}
+	if err := s.db.SetAIConfig(&in); err != nil {
+		return nil, errors.Wrap(err, "db")
+	}
+	return "success", nil
+}