package main import ( "context" "encoding/json" "fmt" "log/slog" "os" "os/signal" "syscall" "code.ewintr.nl/emdb/job" "code.ewintr.nl/emdb/storage" "github.com/tmc/langchaingo/chains" "github.com/tmc/langchaingo/llms/ollama" "github.com/tmc/langchaingo/prompts" ) const ( mentionsTemplate = `The following text is a user comment about the movie {{.title}}. In it, the user may have referenced other movie titles. List them if you see any. ---- {{.review}} ---- If you found any movie titles other than {{.title}}, list them below in a JSON array. If there are other titles, like TV shows, books or games, ignore them. The format is as follows: ["movie title 1", "movie title 2"] Just answer with the JSON and nothing else. If you don't see any other movie titles, just answer with an empty array.` ) func main() { logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) dbHost := os.Getenv("EMDB_DB_HOST") dbName := os.Getenv("EMDB_DB_NAME") dbUser := os.Getenv("EMDB_DB_USER") dbPassword := os.Getenv("EMDB_DB_PASSWORD") pgConnStr := fmt.Sprintf("host=%s user=%s password=%s dbname=%s sslmode=disable", dbHost, dbUser, dbPassword, dbName) dbPostgres, err := storage.NewPostgres(pgConnStr) if err != nil { fmt.Printf("could not create new postgres repo: %s", err.Error()) os.Exit(1) } movieRepo := storage.NewMovieRepository(dbPostgres) reviewRepo := storage.NewReviewRepository(dbPostgres) jobQueue := job.NewJobQueue(dbPostgres, logger) go Work(movieRepo, reviewRepo, jobQueue) c := make(chan os.Signal, 1) signal.Notify(c, syscall.SIGINT, syscall.SIGTERM) <-c } func Work(movieRepo *storage.MovieRepository, reviewRepo *storage.ReviewRepository, jobQueue *job.JobQueue) { for { j, err := jobQueue.Next() if err != nil { fmt.Println(err) os.Exit(1) } review, err := reviewRepo.FindOne(j.ActionID) if err != nil { fmt.Println(err) os.Exit(1) } movie, err := movieRepo.FindOne(review.MovieID) if err != nil { fmt.Println(err) os.Exit(1) } llm, err := ollama.New(ollama.WithModel("mistral")) if err != nil { fmt.Println(err) os.Exit(1) } ctx := context.Background() prompt := prompts.NewPromptTemplate( mentionsTemplate, []string{"title", "review"}, ) llmChain := chains.NewLLMChain(llm, prompt) movieTitle := movie.Title if movie.EnglishTitle != "" && movie.EnglishTitle != movie.Title { movieTitle = fmt.Sprintf("%s (English title: %s)", movieTitle, movie.EnglishTitle) } fmt.Printf("Processing review for movie: %s\n", movieTitle) fmt.Printf("Review: %s\n", review.Review) outputValues, err := chains.Call(ctx, llmChain, map[string]any{ "title": movieTitle, "review": review.Review, }) if err != nil { fmt.Println(err) os.Exit(1) } out, ok := outputValues[llmChain.OutputKey].(string) if !ok { fmt.Println("invalid chain return") } fmt.Println(out) resp := struct { Movies []string `json:"movies"` TVShows []string `json:"tvShows"` Games []string `json:"games"` Books []string `json:"books"` }{} if err := json.Unmarshal([]byte(out), &resp); err != nil { fmt.Printf("could not unmarshal llm response, skipping this one: %s", err) os.Exit(1) } review.Titles = resp if err := reviewRepo.Store(review); err != nil { fmt.Printf("could not update review: %s\n", err) os.Exit(1) } } }