feat: checking users limits when AI scans their images

This commit is contained in:
2025-08-30 11:31:45 +01:00
parent 48579267b5
commit 84a0996be9
2 changed files with 17 additions and 4 deletions

View File

@ -3,8 +3,10 @@ package agents
import (
"context"
"encoding/json"
"fmt"
"screenmark/screenmark/.gen/haystack/haystack/model"
"screenmark/screenmark/agents/client"
"screenmark/screenmark/limits"
"screenmark/screenmark/models"
"github.com/charmbracelet/log"
@ -174,7 +176,7 @@ type addToListArguments struct {
Schema []models.IDValue
}
func NewListAgent(log *log.Logger, listModel models.ListModel) client.AgentClient {
func NewListAgent(log *log.Logger, listModel models.ListModel, limitsMethods limits.LimitsManagerMethods) client.AgentClient {
agentClient := client.CreateAgentClient(client.CreateAgentClientOptions{
SystemPrompt: listPrompt,
JsonTools: listTools,
@ -193,6 +195,16 @@ func NewListAgent(log *log.Logger, listModel models.ListModel) client.AgentClien
return "", err
}
hasReachedLimit, err := limitsMethods.HasReachedStackLimit(info.UserId)
if err != nil {
return "", fmt.Errorf("error checking stack limits: %w", err)
}
if hasReachedLimit {
log.Warn("User has reached limits", "userID", info.UserId)
return "", fmt.Errorf("reached stack limits")
}
ctx := context.Background()
savedList, err := listModel.Save(ctx, info.UserId, args.Name, args.Desription, args.Schema)
@ -201,8 +213,6 @@ func NewListAgent(log *log.Logger, listModel models.ListModel) client.AgentClien
return "", err
}
log.Debug(savedList)
return savedList, nil
})

View File

@ -8,6 +8,7 @@ import (
"net/http"
"os"
"screenmark/screenmark/agents"
"screenmark/screenmark/limits"
"screenmark/screenmark/middleware"
"screenmark/screenmark/models"
"strconv"
@ -86,6 +87,8 @@ func ListenNewImageEvents(db *sql.DB) {
imageModel := models.NewImageModel(db)
listModel := models.NewListModel(db)
limits := limits.CreateLimitsManager(db)
databaseEventLog := createLogger("Database Events 🤖", os.Stdout)
databaseEventLog.SetLevel(log.DebugLevel)
@ -116,7 +119,7 @@ func ListenNewImageEvents(db *sql.DB) {
}
descriptionAgent := agents.NewDescriptionAgent(createLogger("Description 📝", splitWriter), imageModel)
listAgent := agents.NewListAgent(createLogger("Lists 🖋️", splitWriter), listModel)
listAgent := agents.NewListAgent(createLogger("Lists 🖋️", splitWriter), listModel, limits)
var wg sync.WaitGroup
wg.Add(2)