154 lines
3.5 KiB
Go
154 lines
3.5 KiB
Go
package parser
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log/slog"
|
|
"sync/atomic"
|
|
|
|
"git.savin.nyc/alex/go-receipt-tracker/bus"
|
|
"git.savin.nyc/alex/go-receipt-tracker/config"
|
|
"github.com/google/generative-ai-go/genai"
|
|
"google.golang.org/api/option"
|
|
)
|
|
|
|
type Gemini struct {
|
|
id string
|
|
model string
|
|
opts *config.Parser
|
|
client *genai.Client
|
|
subscriptions map[string]chan bus.Event //
|
|
bus *bus.Bus
|
|
log *slog.Logger
|
|
cancel context.CancelFunc //
|
|
end uint32 // ensure the close methods are only called once
|
|
}
|
|
|
|
func NewGemini(cfg *config.Parser, bus *bus.Bus) *Gemini {
|
|
ai := &Gemini{
|
|
id: cfg.Type, // "gemini",
|
|
model: cfg.Model, // "gemini-1.5-flash"
|
|
opts: cfg,
|
|
bus: bus,
|
|
}
|
|
// model := client.GenerativeModel("gemini-1.5-flash")
|
|
|
|
return ai
|
|
}
|
|
|
|
func (ai *Gemini) ID() string {
|
|
return ai.id
|
|
}
|
|
|
|
func (ai *Gemini) Type() string {
|
|
return "parser"
|
|
}
|
|
|
|
func (ai *Gemini) Init(log *slog.Logger) error {
|
|
ai.log = log
|
|
|
|
ctx := context.Background()
|
|
client, err := genai.NewClient(ctx, option.WithAPIKey(ai.opts.ApiKey))
|
|
if err != nil {
|
|
log.Error("cannot", "error", err)
|
|
}
|
|
ai.client = client
|
|
|
|
return nil
|
|
}
|
|
|
|
func (ai *Gemini) Serve() {
|
|
if atomic.LoadUint32(&ai.end) == 1 {
|
|
return
|
|
}
|
|
|
|
ai.subscribe("parser:" + ai.ID())
|
|
ai.subscribe("parser:*")
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
ai.cancel = cancel
|
|
|
|
go ai.eventLoop(ctx)
|
|
}
|
|
|
|
func (ai *Gemini) Stop() {
|
|
ai.client.Close()
|
|
}
|
|
|
|
func (ai *Gemini) subscribe(chn string) error {
|
|
s, err := ai.bus.Subscribe(chn, ai.Type()+":"+ai.ID())
|
|
if err != nil {
|
|
ai.log.Error("couldn't subscribe to a channel", "channel", chn, "error", err.Error())
|
|
return err
|
|
}
|
|
ai.subscriptions[chn] = s
|
|
|
|
return nil
|
|
}
|
|
|
|
// eventLoop loops forever
|
|
func (ai *Gemini) eventLoop(ctx context.Context) {
|
|
ai.log.Debug(ai.ID() + " communication event loop started")
|
|
defer ai.log.Debug(ai.ID() + " communication event loop halted")
|
|
|
|
for {
|
|
for chn, ch := range ai.subscriptions {
|
|
select {
|
|
case event := <-ch:
|
|
ai.log.Debug("got a new message to a channel", "channel", chn)
|
|
|
|
for _, msg := range event.Payload.([]*bus.Message) {
|
|
ai.log.Debug("publishing mqtt message", "message", msg)
|
|
}
|
|
case <-ctx.Done():
|
|
ai.log.Info("stopping " + ai.ID() + " communication event loop")
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (ai *Gemini) Recognize(img string) (res string, err error) {
|
|
ctx := context.Background()
|
|
|
|
image := "./" + img + ".jpg"
|
|
file, err := ai.client.UploadFileFromPath(ctx, image, nil)
|
|
if err != nil {
|
|
ai.log.Error("cannot upload a receipt file", "Messages error: %v\n", err)
|
|
}
|
|
defer ai.client.DeleteFile(ctx, file.Name)
|
|
|
|
gotFile, err := ai.client.GetFile(ctx, file.Name)
|
|
if err != nil {
|
|
ai.log.Error("cannot get uploaded file name", err)
|
|
}
|
|
fmt.Println("Got file:", gotFile.Name)
|
|
|
|
model := ai.client.GenerativeModel("gemini-1.5-flash")
|
|
resp, err := model.GenerateContent(ctx,
|
|
genai.FileData{URI: file.URI},
|
|
genai.Text(config.Request))
|
|
if err != nil {
|
|
ai.log.Error("cannot recognize a receipt", "Messages error: %v\n", err)
|
|
}
|
|
|
|
printResponse(resp)
|
|
// if !utils.IsJSON(resp.Choices[0].Message.Content) {
|
|
// ai.log.Error("OpenAI returned not valid JSON", "%+v", resp.Choices[0].Message.Content)
|
|
// return "", err
|
|
// }
|
|
|
|
return "", nil
|
|
}
|
|
|
|
func printResponse(resp *genai.GenerateContentResponse) {
|
|
for _, cand := range resp.Candidates {
|
|
if cand.Content != nil {
|
|
for _, part := range cand.Content.Parts {
|
|
fmt.Println(part)
|
|
}
|
|
}
|
|
}
|
|
fmt.Println("---")
|
|
}
|