Compare commits
2 Commits
de96f12b55
...
76924a0332
Author | SHA1 | Date | |
---|---|---|---|
76924a0332 | |||
d97593d487 |
@ -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,
|
||||
|
@ -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",
|
||||
|
@ -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",
|
||||
|
@ -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
61
backend/limits/limits.go
Normal 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,
|
||||
}
|
||||
}
|
36
backend/middleware/limits.go
Normal file
36
backend/middleware/limits.go
Normal 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)
|
||||
}
|
||||
}
|
@ -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)
|
||||
|
||||
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
Reference in New Issue
Block a user