alpha version

workers/openai.go

@@ -0,0 +1,160 @@
+package workers
+
+import (
+	"context"
+	"log/slog"
+	"sync/atomic"
+
+	"github.com/alex-savin/go-receipt-tracker/bus"
+	"github.com/alex-savin/go-receipt-tracker/config"
+	"github.com/alex-savin/go-receipt-tracker/utils"
+	"github.com/sashabaranov/go-openai"
+)
+
+type OpenAi struct {
+	id            string
+	model         string
+	opts          *config.Parser
+	client        *openai.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 NewOpenAi(cfg *config.Parser, b *bus.Bus) *OpenAi {
+
+	ai := &OpenAi{
+		id:            cfg.Type,  // "openai",
+		model:         cfg.Model, // "gpt-4o-mini"
+		opts:          cfg,
+		subscriptions: make(map[string]chan bus.Event),
+		bus:           b,
+	}
+	return ai
+}
+
+func (ai *OpenAi) ID() string {
+	return ai.id
+}
+
+func (ai *OpenAi) Type() string {
+	return "parser"
+}
+
+func (ai *OpenAi) Init(log *slog.Logger) error {
+	ai.client = openai.NewClient(ai.opts.ApiKey)
+	ai.log = log
+
+	return nil
+}
+
+func (ai *OpenAi) Serve() {
+	if atomic.LoadUint32(&ai.end) == 1 {
+		return
+	}
+
+	ai.subscribe("parser:" + ai.ID())
+
+	ctx, cancel := context.WithCancel(context.Background())
+	ai.cancel = cancel
+
+	go ai.eventLoop(ctx)
+}
+
+func (ai *OpenAi) OneTime() error {
+	return nil
+}
+
+func (ai *OpenAi) Stop() {
+}
+
+func (ai *OpenAi) Close() {
+	ai.cancel()
+}
+
+func (ai *OpenAi) 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 *OpenAi) eventLoop(ctx context.Context) {
+	ai.log.Debug(ai.ID() + " communication event loop is started")
+	defer ai.log.Debug(ai.ID() + " communication event loop halted")
+
+	for {
+		ai.log.Debug("looping over the subscriptions", "channel", "parser:"+ai.ID(), "worker", ai.ID())
+
+		select {
+		case event := <-ai.subscriptions["parser:"+ai.ID()]:
+			ai.log.Debug("worker got a new message from a channel", "channel", "parser:"+ai.ID(), "type", ai.Type(), "id", ai.ID())
+
+			switch event.Payload.(type) {
+			case *bus.Message:
+				msg := event.Payload.(*bus.Message)
+				res, err := ai.recognize("data:" + event.Payload.(*bus.Message).Image.Type + ";base64," + event.Payload.(*bus.Message).Image.Base64)
+				if err != nil {
+					ai.log.Error("got an error from parser ("+ai.ID()+")", "error", err)
+				}
+
+				msg.Image.Parsed[ai.ID()] = res
+				err = ai.bus.Publish("processor:receipt_add", msg)
+				if err != nil {
+					ai.log.Error("couldn't publish to a channel", "channel", "processor:receipt_add", "error", err.Error())
+				}
+			}
+		case <-ctx.Done():
+			ai.log.Info("stopping " + ai.ID() + " communication event loop")
+			return
+		}
+	}
+}
+
+func (ai *OpenAi) recognize(img string) (res string, err error) {
+	resp, err := ai.client.CreateChatCompletion(
+		context.Background(),
+		openai.ChatCompletionRequest{
+			Model:     openai.GPT4oMini,
+			MaxTokens: 1000,
+			Messages: []openai.ChatCompletionMessage{
+				{
+					Role: openai.ChatMessageRoleUser,
+					MultiContent: []openai.ChatMessagePart{
+						{
+							Type: openai.ChatMessagePartTypeText,
+							Text: config.Request,
+						},
+						{
+							Type: openai.ChatMessagePartTypeImageURL,
+							ImageURL: &openai.ChatMessageImageURL{
+								URL:    img,
+								Detail: openai.ImageURLDetailHigh,
+							},
+						},
+					},
+				},
+			},
+		},
+	)
+
+	if err != nil {
+		ai.log.Error("error during recognition process", "%+v", resp.Choices[0].Message.Content)
+		return "", err
+	}
+	if !utils.IsJSON(resp.Choices[0].Message.Content) {
+		ai.log.Error("returned not valid JSON", "%+v", resp.Choices[0].Message.Content)
+		return "", err
+	}
+
+	ai.log.Debug("recognition result", "worker", ai.ID(), "output", resp.Choices[0].Message.Content)
+
+	return resp.Choices[0].Message.Content, nil
+}