diff --git a/client/ollama.go b/client/ollama.go new file mode 100644 index 0000000..04c597e --- /dev/null +++ b/client/ollama.go @@ -0,0 +1,64 @@ +package client + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" +) + +type Ollama struct { + baseURL string + c *http.Client +} + +func NewOllama(baseURL string) *Ollama { + return &Ollama{ + baseURL: baseURL, + c: &http.Client{}, + } +} + +func (o *Ollama) Generate(model, prompt string) (string, error) { + url := fmt.Sprintf("%s/api/generate", o.baseURL) + reqBody := struct { + Model string + Prompt string + Format string + Stream bool + }{ + Model: model, + Prompt: prompt, + Format: "json", + Stream: false, + } + reqBodyJSON, err := json.Marshal(reqBody) + if err != nil { + return "", err + } + + req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(reqBodyJSON)) + if err != nil { + return "", err + } + res, err := o.c.Do(req) + if err != nil { + return "", err + } + + body, err := io.ReadAll(res.Body) + if err != nil { + return "", err + } + defer res.Body.Close() + + resBody := struct { + Response string + }{} + if err := json.Unmarshal(body, &resBody); err != nil { + return "", err + } + + return resBody.Response, nil +} diff --git a/storage/reviewpg.go b/storage/reviewpg.go index fe6c9d1..09cad90 100644 --- a/storage/reviewpg.go +++ b/storage/reviewpg.go @@ -12,13 +12,6 @@ const ( type ReviewSource string -type Titles struct { - Movies []string `json:"movies"` - TVShows []string `json:"tvShows"` - Games []string `json:"games"` - Books []string `json:"books"` -} - type Review struct { ID string MovieID string @@ -27,7 +20,7 @@ type Review struct { Review string MovieRating int Quality int - Titles Titles + Titles []string } type ReviewRepository struct { diff --git a/worker-client/main.go b/worker-client/main.go index b296837..c4c5d7e 100644 --- a/worker-client/main.go +++ b/worker-client/main.go @@ -42,8 +42,9 @@ func main() { movieRepo := storage.NewMovieRepository(dbPostgres) reviewRepo := storage.NewReviewRepository(dbPostgres) jobQueue := job.NewJobQueue(dbPostgres, logger) + ollama := client.NewOllama("http://localhost:11434") - w := worker.NewWorker(jobQueue, movieRepo, reviewRepo, client.NewIMDB(), logger) + w := worker.NewWorker(jobQueue, movieRepo, reviewRepo, client.NewIMDB(), ollama, logger) go w.Run() diff --git a/worker-client/worker/findtitles.go b/worker-client/worker/findtitles.go index cce564b..40ee690 100644 --- a/worker-client/worker/findtitles.go +++ b/worker-client/worker/findtitles.go @@ -1,27 +1,22 @@ package worker import ( - "context" "encoding/json" "fmt" - - "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. + 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. ---- -{{.review}} +%s ---- -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: +If you found any movie titles other than %s, 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.` +Just answer with the JSON and nothing else. If you don't see any other movie titles, just answer with an empty JSON array.` ) func (w *Worker) FindTitles(jobID int, reviewID string) { @@ -40,53 +35,20 @@ func (w *Worker) FindTitles(jobID int, reviewID string) { w.jq.MarkFailed(jobID) return } - llm, err := ollama.New(ollama.WithModel("mistral")) - if err != nil { - logger.Error("could not create llm", "error", err) - w.jq.MarkFailed(jobID) - return - } - - 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, - }) + prompt := fmt.Sprintf(mentionsTemplate, movieTitle, review.Review, movieTitle) + titles, err := w.ollama.Generate("mistral", prompt) if err != nil { - logger.Error("could not call chain", "error", err) - w.jq.MarkFailed(jobID) - return + logger.Error("could not find titles: %w", err) } - - out, ok := outputValues[llmChain.OutputKey].(string) - if !ok { - logger.Error("chain output is not valid") - w.jq.MarkFailed(jobID) - 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 { + logger.Info("checked review", "found", titles) + var resp []string + if err := json.Unmarshal([]byte(titles), &resp); err != nil { logger.Error("could not unmarshal llm response", "error", err) w.jq.MarkFailed(jobID) return diff --git a/worker-client/worker/worker.go b/worker-client/worker/worker.go index 331a445..35ab0ca 100644 --- a/worker-client/worker/worker.go +++ b/worker-client/worker/worker.go @@ -20,15 +20,17 @@ type Worker struct { movieRepo *storage.MovieRepository reviewRepo *storage.ReviewRepository imdb *client.IMDB + ollama *client.Ollama logger *slog.Logger } -func NewWorker(jq *job.JobQueue, movieRepo *storage.MovieRepository, reviewRepo *storage.ReviewRepository, imdb *client.IMDB, logger *slog.Logger) *Worker { +func NewWorker(jq *job.JobQueue, movieRepo *storage.MovieRepository, reviewRepo *storage.ReviewRepository, imdb *client.IMDB, ollama *client.Ollama, logger *slog.Logger) *Worker { return &Worker{ jq: jq, movieRepo: movieRepo, reviewRepo: reviewRepo, imdb: imdb, + ollama: ollama, logger: logger.With("service", "worker"), } } @@ -57,7 +59,7 @@ func (w *Worker) Run() { case job.ActionRefreshAllIMDBReviews: w.RefreshAllReviews(j.ID) case job.ActionFindTitles: - //w.FindTitles(j.ID, j.ActionID) + w.FindTitles(j.ID, j.ActionID) case job.ActionFindAllTitles: w.FindAllTitles(j.ID) default: