feat: checking users limits when AI scans their images
This commit is contained in:
@ -3,8 +3,10 @@ package agents
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
"screenmark/screenmark/.gen/haystack/haystack/model"
|
"screenmark/screenmark/.gen/haystack/haystack/model"
|
||||||
"screenmark/screenmark/agents/client"
|
"screenmark/screenmark/agents/client"
|
||||||
|
"screenmark/screenmark/limits"
|
||||||
"screenmark/screenmark/models"
|
"screenmark/screenmark/models"
|
||||||
|
|
||||||
"github.com/charmbracelet/log"
|
"github.com/charmbracelet/log"
|
||||||
@ -174,7 +176,7 @@ type addToListArguments struct {
|
|||||||
Schema []models.IDValue
|
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{
|
agentClient := client.CreateAgentClient(client.CreateAgentClientOptions{
|
||||||
SystemPrompt: listPrompt,
|
SystemPrompt: listPrompt,
|
||||||
JsonTools: listTools,
|
JsonTools: listTools,
|
||||||
@ -193,6 +195,16 @@ func NewListAgent(log *log.Logger, listModel models.ListModel) client.AgentClien
|
|||||||
return "", err
|
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()
|
ctx := context.Background()
|
||||||
savedList, err := listModel.Save(ctx, info.UserId, args.Name, args.Desription, args.Schema)
|
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
|
return "", err
|
||||||
}
|
}
|
||||||
|
|
||||||
log.Debug(savedList)
|
|
||||||
|
|
||||||
return savedList, nil
|
return savedList, nil
|
||||||
})
|
})
|
||||||
|
|
||||||
|
@ -8,6 +8,7 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
"screenmark/screenmark/agents"
|
"screenmark/screenmark/agents"
|
||||||
|
"screenmark/screenmark/limits"
|
||||||
"screenmark/screenmark/middleware"
|
"screenmark/screenmark/middleware"
|
||||||
"screenmark/screenmark/models"
|
"screenmark/screenmark/models"
|
||||||
"strconv"
|
"strconv"
|
||||||
@ -86,6 +87,8 @@ func ListenNewImageEvents(db *sql.DB) {
|
|||||||
imageModel := models.NewImageModel(db)
|
imageModel := models.NewImageModel(db)
|
||||||
listModel := models.NewListModel(db)
|
listModel := models.NewListModel(db)
|
||||||
|
|
||||||
|
limits := limits.CreateLimitsManager(db)
|
||||||
|
|
||||||
databaseEventLog := createLogger("Database Events 🤖", os.Stdout)
|
databaseEventLog := createLogger("Database Events 🤖", os.Stdout)
|
||||||
databaseEventLog.SetLevel(log.DebugLevel)
|
databaseEventLog.SetLevel(log.DebugLevel)
|
||||||
|
|
||||||
@ -116,7 +119,7 @@ func ListenNewImageEvents(db *sql.DB) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
descriptionAgent := agents.NewDescriptionAgent(createLogger("Description 📝", splitWriter), imageModel)
|
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
|
var wg sync.WaitGroup
|
||||||
wg.Add(2)
|
wg.Add(2)
|
||||||
|
Reference in New Issue
Block a user