2 Commits

Author SHA1 Message Date
76924a0332 chore: using requesty policy instead of direct models 2025-08-30 10:46:37 +01:00
d97593d487 feat: adding limits 2025-08-30 10:45:53 +01:00
8 changed files with 124 additions and 18 deletions

View File

@ -246,7 +246,7 @@ func (client *AgentClient) RunAgent(userId uuid.UUID, imageId uuid.UUID, imageNa
request := AgentRequestBody{
Tools: &tools,
ToolChoice: &toolChoice,
Model: "google/gemini-2.5-flash",
Model: "policy/images",
RandomSeed: &seed,
Temperature: 0.3,
EndToolCall: client.Options.EndToolCall,
@ -284,7 +284,7 @@ func (client *AgentClient) RunAgentAlone(userID uuid.UUID, userReq string) error
request := AgentRequestBody{
Tools: &tools,
ToolChoice: &toolChoice,
Model: "google/gemini-2.5-flash",
Model: "policy/images",
RandomSeed: &seed,
Temperature: 0.3,
EndToolCall: client.Options.EndToolCall,

View File

@ -81,7 +81,7 @@ type CreateListAgent struct {
func (agent *CreateListAgent) CreateList(log *log.Logger, userID uuid.UUID, userReq string) error {
request := client.AgentRequestBody{
Model: "google/gemini-2.5-flash",
Model: "policy/images",
Temperature: 0.3,
ResponseFormat: client.ResponseFormat{
Type: "json_object",

View File

@ -28,7 +28,7 @@ type DescriptionAgent struct {
func (agent DescriptionAgent) Describe(log *log.Logger, imageId uuid.UUID, imageName string, imageData []byte) error {
request := client.AgentRequestBody{
Model: "google/gemini-2.5-flash-lite-preview-06-17",
Model: "policy/images",
Temperature: 0.3,
ResponseFormat: client.ResponseFormat{
Type: "text",

View File

@ -10,6 +10,7 @@ import (
"os"
"path/filepath"
"screenmark/screenmark/.gen/haystack/haystack/model"
"screenmark/screenmark/limits"
"screenmark/screenmark/middleware"
"screenmark/screenmark/models"
@ -18,9 +19,10 @@ import (
)
type ImageHandler struct {
logger *log.Logger
imageModel models.ImageModel
userModel models.UserModel
logger *log.Logger
imageModel models.ImageModel
userModel models.UserModel
limitsManager limits.LimitsManagerMethods
}
type ImagesReturn struct {
@ -153,18 +155,19 @@ func (h *ImageHandler) CreateRoutes(r chi.Router) {
r.Use(middleware.SetJson)
r.Get("/", h.listImages)
r.Post("/{name}", h.uploadImage)
r.Post("/{name}", middleware.WithLimit(h.logger, h.limitsManager.HasReachedImageLimit, h.uploadImage))
})
}
func CreateImageHandler(db *sql.DB) ImageHandler {
func CreateImageHandler(db *sql.DB, limitsManager limits.LimitsManagerMethods) ImageHandler {
imageModel := models.NewImageModel(db)
userModel := models.NewUserModel(db)
logger := log.New(os.Stdout).WithPrefix("Images")
return ImageHandler{
logger: logger,
imageModel: imageModel,
userModel: userModel,
logger: logger,
imageModel: imageModel,
userModel: userModel,
limitsManager: limitsManager,
}
}

61
backend/limits/limits.go Normal file
View File

@ -0,0 +1,61 @@
package limits
import (
"database/sql"
. "screenmark/screenmark/.gen/haystack/haystack/table"
. "github.com/go-jet/jet/v2/postgres"
"github.com/google/uuid"
)
const (
LISTS_LIMIT = 10
IMAGE_LIMIT = 50
)
type LimitsManager struct {
dbPool *sql.DB
}
type LimitsManagerMethods interface {
HasReachedStackLimit(userID uuid.UUID) (bool, error)
HasReachedImageLimit(userID uuid.UUID) (bool, error)
}
type listCount struct {
ListCount int `alias:"list_count"`
}
func (m *LimitsManager) HasReachedStackLimit(userID uuid.UUID) (bool, error) {
getStacks := Lists.
SELECT(COUNT(Lists.UserID).AS("listCount.ListCount")).
WHERE(Lists.UserID.EQ(UUID(userID)))
var count listCount
err := getStacks.Query(m.dbPool, &count)
return count.ListCount >= LISTS_LIMIT, err
}
type imageCount struct {
ImageCount int `alias:"image_count"`
}
func (m *LimitsManager) HasReachedImageLimit(userID uuid.UUID) (bool, error) {
getStacks := UserImages.
SELECT(COUNT(UserImages.UserID).AS("imageCount.ImageCount")).
WHERE(UserImages.UserID.EQ(UUID(userID)))
var count imageCount
err := getStacks.Query(m.dbPool, &count)
return count.ImageCount >= IMAGE_LIMIT, err
}
func CreateLimitsManager(db *sql.DB) *LimitsManager {
return &LimitsManager{
db,
}
}

View File

@ -0,0 +1,36 @@
package middleware
import (
"net/http"
"github.com/charmbracelet/log"
"github.com/google/uuid"
)
func WithLimit(logger *log.Logger, getLimit func(userID uuid.UUID) (bool, error), next func(w http.ResponseWriter, r *http.Request)) func(w http.ResponseWriter, r *http.Request) {
return func(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
userID, err := GetUserID(ctx, logger, w)
if err != nil {
w.WriteHeader(http.StatusUnauthorized)
return
}
hasReachedLimit, err := getLimit(userID)
if err != nil {
logger.Error("failed to image limit", "err", err)
w.WriteHeader(http.StatusInternalServerError)
return
}
logger.Info("Limits", "hasReachedLimit", hasReachedLimit)
if hasReachedLimit {
w.WriteHeader(http.StatusTooManyRequests)
return
}
next(w, r)
}
}

View File

@ -6,6 +6,7 @@ import (
"screenmark/screenmark/agents/client"
"screenmark/screenmark/auth"
"screenmark/screenmark/images"
"screenmark/screenmark/limits"
"screenmark/screenmark/models"
"screenmark/screenmark/stacks"
@ -27,9 +28,11 @@ func setupRouter(db *sql.DB) chi.Router {
imageModel := models.NewImageModel(db)
stackModel := models.NewListModel(db)
stackHandler := stacks.CreateStackHandler(db)
limitsManager := limits.CreateLimitsManager(db)
stackHandler := stacks.CreateStackHandler(db, limitsManager)
authHandler := auth.CreateAuthHandler(db)
imageHandler := images.CreateImageHandler(db)
imageHandler := images.CreateImageHandler(db, limitsManager)
notifier := NewNotifier[Notification](10)

View File

@ -6,6 +6,7 @@ import (
"net/http"
"os"
. "screenmark/screenmark/.gen/haystack/haystack/model"
"screenmark/screenmark/limits"
"screenmark/screenmark/middleware"
"screenmark/screenmark/models"
"strings"
@ -15,8 +16,9 @@ import (
)
type StackHandler struct {
logger *log.Logger
stackModel models.ListModel
logger *log.Logger
stackModel models.ListModel
limitsManager limits.LimitsManagerMethods
}
func (h *StackHandler) getAllStacks(w http.ResponseWriter, r *http.Request) {
@ -144,18 +146,19 @@ func (h *StackHandler) CreateRoutes(r chi.Router) {
r.Get("/", h.getAllStacks)
r.Get("/{listID}", h.getStackItems)
r.Post("/", middleware.WithValidatedPost(h.createStack))
r.Post("/", middleware.WithLimit(h.logger, h.limitsManager.HasReachedStackLimit, middleware.WithValidatedPost(h.createStack)))
r.Patch("/{listID}", middleware.WithValidatedPost(h.editStack))
r.Delete("/{listID}", h.deleteStack)
})
}
func CreateStackHandler(db *sql.DB) StackHandler {
func CreateStackHandler(db *sql.DB, limitsManager limits.LimitsManagerMethods) StackHandler {
stackModel := models.NewListModel(db)
logger := log.New(os.Stdout).WithPrefix("Stacks")
return StackHandler{
logger,
stackModel,
limitsManager,
}
}