Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion internal/ent/schema/request.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package schema
import (
"entgo.io/contrib/entgql"
"entgo.io/ent"
"entgo.io/ent/dialect"
"entgo.io/ent/schema"
"entgo.io/ent/schema/edge"
"entgo.io/ent/schema/field"
Expand Down Expand Up @@ -88,7 +89,9 @@ func (Request) Fields() []ent.Field {
// External ID for tracking requests in external systems
field.String("external_id").
Optional().
MaxLen(512),
SchemaType(map[string]string{
dialect.MySQL: "text",
}),
// The status of the request.
field.Enum("status").Values("pending", "processing", "completed", "failed", "canceled"),
// Whether the request is a streaming request
Expand Down
5 changes: 4 additions & 1 deletion internal/ent/schema/request_execution.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package schema
import (
"entgo.io/contrib/entgql"
"entgo.io/ent"
"entgo.io/ent/dialect"
"entgo.io/ent/schema"
"entgo.io/ent/schema/edge"
"entgo.io/ent/schema/field"
Expand Down Expand Up @@ -46,7 +47,9 @@ func (RequestExecution) Fields() []ent.Field {
// External ID for tracking requests in external systems
field.String("external_id").
Optional().
MaxLen(512),
SchemaType(map[string]string{
dialect.MySQL: "text",
}),
field.String("model_id").Immutable(),
// The format of the request, e.g: openai/chat_completions, claude/messages, openai/response.
field.String("format").Immutable().Default("openai/chat_completions"),
Expand Down
22 changes: 13 additions & 9 deletions internal/server/api/openai.go
Original file line number Diff line number Diff line change
Expand Up @@ -614,15 +614,19 @@ func (handlers *OpenAIHandlers) ListModels(c *gin.Context) {
return m.ID
})

dbModels, err := handlers.EntClient.Model.Query().
Where(
model.StatusEQ(model.StatusEnabled),
model.ModelIDIn(visibleIDs...),
).
All(ctx)
if err != nil {
handlers.writeOpenAIInternalError(c, requestID, err)
return
var dbModels []*ent.Model
for _, chunk := range lo.Chunk(visibleIDs, biz.SQLiteMaxVariableLimit) {
chunkModels, err := handlers.EntClient.Model.Query().
Where(
model.StatusEQ(model.StatusEnabled),
model.ModelIDIn(chunk...),
).
All(ctx)
if err != nil {
handlers.writeOpenAIInternalError(c, requestID, err)
return
}
dbModels = append(dbModels, chunkModels...)
}

dbModelMap := make(map[string]*ent.Model, len(dbModels))
Expand Down
136 changes: 89 additions & 47 deletions internal/server/biz/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,12 @@ func NewModelService(params ModelServiceParams) *ModelService {
}
}

// SQLiteMaxVariableLimit is a safe chunk size for SQL IN clauses.
// SQLite versions prior to 3.32.0 have a limit of 999 variables;
// newer versions support 32766. We use 500 to stay safely below both limits
// and leave room for other query parameters.
const SQLiteMaxVariableLimit = 500

type ModelService struct {
*AbstractService

Expand Down Expand Up @@ -312,22 +318,31 @@ func (svc *ModelService) BulkCreateModels(ctx context.Context, inputs []*ent.Cre
inputMap[key] = true
}

// Check if any models already exist
existingModels, err := svc.entFromContext(ctx).Model.Query().
Where(func(s *sql.Selector) {
var predicates []*sql.Predicate
for _, input := range inputs {
predicates = append(predicates, sql.And(
sql.EQ(model.FieldDeveloper, input.Developer),
sql.EQ(model.FieldModelID, input.ModelID),
))
}
// Check if any models already exist (chunk inputs to avoid SQLite variable limit;
// each predicate uses 2 variables: developer + modelID).
chunkSize := SQLiteMaxVariableLimit / 2

s.Where(sql.Or(predicates...))
}).
All(ctx)
if err != nil {
return nil, fmt.Errorf("failed to check existing models: %w", err)
var existingModels []*ent.Model

for _, chunk := range lo.Chunk(inputs, chunkSize) {
chunkModels, err := svc.entFromContext(ctx).Model.Query().
Where(func(s *sql.Selector) {
var predicates []*sql.Predicate
for _, input := range chunk {
predicates = append(predicates, sql.And(
sql.EQ(model.FieldDeveloper, input.Developer),
sql.EQ(model.FieldModelID, input.ModelID),
))
}

s.Where(sql.Or(predicates...))
}).
All(ctx)
if err != nil {
return nil, fmt.Errorf("failed to check existing models: %w", err)
}

existingModels = append(existingModels, chunkModels...)
}

if len(existingModels) > 0 {
Expand Down Expand Up @@ -430,51 +445,63 @@ func (svc *ModelService) DeleteModel(ctx context.Context, id int) error {
}

// BulkArchiveModels archives multiple models by their IDs.
// IDs are chunked to avoid SQLite "too many SQL variables" error.
func (svc *ModelService) BulkArchiveModels(ctx context.Context, ids []int) error {
_, err := svc.entFromContext(ctx).Model.Update().
Where(model.IDIn(ids...)).
SetStatus(model.StatusArchived).
Save(ctx)
if err != nil {
return fmt.Errorf("failed to bulk archive models: %w", err)
for _, chunk := range lo.Chunk(ids, SQLiteMaxVariableLimit) {
_, err := svc.entFromContext(ctx).Model.Update().
Where(model.IDIn(chunk...)).
SetStatus(model.StatusArchived).
Save(ctx)
if err != nil {
return fmt.Errorf("failed to bulk archive models: %w", err)
}
}

return nil
}

// BulkDisableModels disables multiple models by their IDs.
// IDs are chunked to avoid SQLite "too many SQL variables" error.
func (svc *ModelService) BulkDisableModels(ctx context.Context, ids []int) error {
_, err := svc.entFromContext(ctx).Model.Update().
Where(model.IDIn(ids...)).
SetStatus(model.StatusDisabled).
Save(ctx)
if err != nil {
return fmt.Errorf("failed to bulk disable models: %w", err)
for _, chunk := range lo.Chunk(ids, SQLiteMaxVariableLimit) {
_, err := svc.entFromContext(ctx).Model.Update().
Where(model.IDIn(chunk...)).
SetStatus(model.StatusDisabled).
Save(ctx)
if err != nil {
return fmt.Errorf("failed to bulk disable models: %w", err)
}
}

return nil
}

// BulkEnableModels enables multiple models by their IDs.
// IDs are chunked to avoid SQLite "too many SQL variables" error.
func (svc *ModelService) BulkEnableModels(ctx context.Context, ids []int) error {
_, err := svc.entFromContext(ctx).Model.Update().
Where(model.IDIn(ids...)).
SetStatus(model.StatusEnabled).
Save(ctx)
if err != nil {
return fmt.Errorf("failed to bulk enable models: %w", err)
for _, chunk := range lo.Chunk(ids, SQLiteMaxVariableLimit) {
_, err := svc.entFromContext(ctx).Model.Update().
Where(model.IDIn(chunk...)).
SetStatus(model.StatusEnabled).
Save(ctx)
if err != nil {
return fmt.Errorf("failed to bulk enable models: %w", err)
}
}

return nil
}

// BulkDeleteModels deletes multiple models by their IDs.
// IDs are chunked to avoid SQLite "too many SQL variables" error.
func (svc *ModelService) BulkDeleteModels(ctx context.Context, ids []int) error {
_, err := svc.entFromContext(ctx).Model.Delete().
Where(model.IDIn(ids...)).
Exec(ctx)
if err != nil {
return fmt.Errorf("failed to bulk delete models: %w", err)
for _, chunk := range lo.Chunk(ids, SQLiteMaxVariableLimit) {
_, err := svc.entFromContext(ctx).Model.Delete().
Where(model.IDIn(chunk...)).
Exec(ctx)
if err != nil {
return fmt.Errorf("failed to bulk delete models: %w", err)
}
}

return nil
Expand Down Expand Up @@ -651,17 +678,32 @@ func (svc *ModelService) ListEnabledModels(ctx context.Context) ([]ModelFacade,
// queryConfiguredModelFacades queries enabled Model entities and returns them as ModelFacades
// filtered by allowed model IDs and channel associations.
func (svc *ModelService) queryConfiguredModelFacades(ctx context.Context, allowedModelIDs []string, channels []*Channel) ([]ModelFacade, error) {
query := svc.entFromContext(ctx).
Model.
Query().
Where(model.StatusEQ(model.StatusEnabled))
var enabledModels []*ent.Model

if len(allowedModelIDs) > 0 {
query = query.Where(model.ModelIDIn(allowedModelIDs...))
}
// Chunk allowedModelIDs to avoid SQLite "too many SQL variables" error.
for _, chunk := range lo.Chunk(allowedModelIDs, SQLiteMaxVariableLimit) {
chunkModels, err := svc.entFromContext(ctx).Model.Query().
Where(
model.StatusEQ(model.StatusEnabled),
model.ModelIDIn(chunk...),
).
All(ctx)
if err != nil {
return nil, fmt.Errorf("failed to list configured models: %w", err)
}

enabledModels, err := query.All(ctx)
if err != nil {
return nil, fmt.Errorf("failed to list configured models: %w", err)
enabledModels = append(enabledModels, chunkModels...)
}
} else {
var err error

enabledModels, err = svc.entFromContext(ctx).Model.Query().
Where(model.StatusEQ(model.StatusEnabled)).
All(ctx)
if err != nil {
return nil, fmt.Errorf("failed to list configured models: %w", err)
}
}

var models []ModelFacade
Expand Down
Loading