Skip to content

Commit a6e5bf5

Browse files
Dev AgentQinYuuuu
authored andcommitted
add query to list model
1 parent 99232a1 commit a6e5bf5

File tree

4 files changed

+447
-12
lines changed

4 files changed

+447
-12
lines changed

.tool-versions

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,5 @@
11
golang 1.25.5
2-
mockery 2.53.5
2+
mockery 2.53.5
3+
minikube 1.34.0
4+
kubectl 1.28.3
5+
argo 3.6.10

aigateway/handler/openai.go

Lines changed: 70 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"log/slog"
1010
"net/http"
1111
"net/url"
12+
"strconv"
1213
"strings"
1314
"time"
1415

@@ -95,10 +96,11 @@ type OpenAIHandlerImpl struct {
9596
// ListModels godoc
9697
// @Security ApiKey
9798
// @Summary List available models
98-
// @Description Returns a list of available models
99+
// @Description Returns a list of available models, supports fuzzy search by model_id query parameter
99100
// @Tags AIGateway
100101
// @Accept json
101102
// @Produce json
103+
// @Param model_id query string false "Model ID for fuzzy search"
102104
// @Success 200 {object} types.ModelList "OK"
103105
// @Failure 500 {object} error "Internal server error"
104106
// @Router /v1/models [get]
@@ -116,9 +118,74 @@ func (h *OpenAIHandlerImpl) ListModels(c *gin.Context) {
116118
return
117119
}
118120

121+
// Apply fuzzy search filter if model_id query parameter is provided
122+
searchQuery := c.Query("model_id")
123+
if searchQuery != "" {
124+
filteredModels := make([]types.Model, 0)
125+
for _, model := range models {
126+
if strings.Contains(strings.ToLower(model.ID), strings.ToLower(searchQuery)) {
127+
filteredModels = append(filteredModels, model)
128+
}
129+
}
130+
models = filteredModels
131+
}
132+
133+
// Parse pagination parameters
134+
perStr := c.Query("per")
135+
pageStr := c.Query("page")
136+
137+
// Set default values
138+
per := 20 // default per page
139+
page := 1 // default page (1-based)
140+
141+
if perStr != "" {
142+
if parsedPerPage, err := strconv.Atoi(perStr); err == nil && parsedPerPage > 0 {
143+
per = parsedPerPage
144+
// Cap the per_page to prevent excessive requests
145+
if per > 100 {
146+
per = 100
147+
}
148+
}
149+
}
150+
151+
if pageStr != "" {
152+
if parsedPage, err := strconv.Atoi(pageStr); err == nil && parsedPage > 0 {
153+
page = parsedPage
154+
}
155+
}
156+
157+
totalCount := len(models)
158+
159+
// Apply pagination
160+
offset := (page - 1) * per
161+
startIndex := offset
162+
if startIndex > totalCount {
163+
startIndex = totalCount
164+
}
165+
166+
endIndex := startIndex + per
167+
if endIndex > totalCount {
168+
endIndex = totalCount
169+
}
170+
171+
paginatedModels := models[startIndex:endIndex]
172+
173+
// Set pagination metadata
174+
var firstID, lastID *string
175+
if len(paginatedModels) > 0 {
176+
firstID = &paginatedModels[0].ID
177+
lastID = &paginatedModels[len(paginatedModels)-1].ID
178+
}
179+
180+
hasMore := endIndex < totalCount
181+
119182
response := types.ModelList{
120-
Object: "list",
121-
Data: models,
183+
Object: "list",
184+
Data: paginatedModels,
185+
FirstID: firstID,
186+
LastID: lastID,
187+
HasMore: hasMore,
188+
TotalCount: totalCount,
122189
}
123190

124191
c.PureJSON(http.StatusOK, response)

0 commit comments

Comments
 (0)