emdb/worker-client/worker/findtitles.go

72 lines
1.9 KiB
Go
Raw Permalink Normal View History

2024-03-09 14:59:18 +01:00
package worker
import (
"encoding/json"
"fmt"
2024-03-10 10:21:01 +01:00
2024-09-17 07:43:12 +02:00
"go-mod.ewintr.nl/emdb/storage"
2024-03-09 14:59:18 +01:00
)
const (
2024-03-09 21:13:32 +01:00
mentionsTemplate = `The following text is a user comment about the movie %s. In it, the user may have referenced other movie titles. List them if you see any.
2024-03-09 14:59:18 +01:00
----
2024-03-09 21:13:32 +01:00
%s
2024-03-09 14:59:18 +01:00
----
2024-03-10 10:21:01 +01:00
If you found any movie titles other than %s, list them below in a JSON array. If there are titles of other media, like TV shows, books or games, ignore them. The format is as follows:
2024-03-09 14:59:18 +01:00
2024-03-10 10:21:01 +01:00
---
{"titles": ["movie title 1", "movie title 2"]}
---
2024-03-09 14:59:18 +01:00
2024-03-10 10:21:01 +01:00
Just answer with the JSON and nothing else. If you don't see any other movie titles, just use an empty JSON array.`
2024-03-09 14:59:18 +01:00
)
func (w *Worker) FindTitles(jobID int, reviewID string) {
logger := w.logger.With("method", "findTitles", "jobID", jobID)
review, err := w.reviewRepo.FindOne(reviewID)
if err != nil {
logger.Error("could not get review", "error", err)
w.jq.MarkFailed(jobID)
return
}
movie, err := w.movieRepo.FindOne(review.MovieID)
if err != nil {
logger.Error("could not get movie", "error", err)
w.jq.MarkFailed(jobID)
return
}
movieTitle := movie.Title
if movie.EnglishTitle != "" && movie.EnglishTitle != movie.Title {
movieTitle = fmt.Sprintf("%s (English title: %s)", movieTitle, movie.EnglishTitle)
}
2024-03-09 21:13:32 +01:00
prompt := fmt.Sprintf(mentionsTemplate, movieTitle, review.Review, movieTitle)
2024-03-10 10:21:01 +01:00
resp, err := w.ollama.Generate("mistral", prompt)
2024-03-09 14:59:18 +01:00
if err != nil {
2024-03-09 21:13:32 +01:00
logger.Error("could not find titles: %w", err)
2024-03-09 14:59:18 +01:00
}
2024-03-10 10:21:01 +01:00
logger.Info("checked review", "found", resp)
var mentions storage.TitleMentions
if err := json.Unmarshal([]byte(resp), &mentions); err != nil {
2024-03-09 14:59:18 +01:00
logger.Error("could not unmarshal llm response", "error", err)
w.jq.MarkFailed(jobID)
return
}
2024-03-10 10:21:01 +01:00
review.Mentions = mentions
2024-03-09 14:59:18 +01:00
if err := w.reviewRepo.Store(review); err != nil {
logger.Error("could not update review", "error", err)
w.jq.MarkFailed(jobID)
return
}
2024-03-10 10:21:01 +01:00
logger.Info("done finding title mentions", "count", len(mentions.Titles))
w.jq.MarkDone(jobID)
2024-03-09 14:59:18 +01:00
}