Skip to content

Commit f2e61bf

Browse files
0uep0uep
committed
feat(proxy): add /d1 /a1 /A1 routes
1 parent f23a93f commit f2e61bf

3 files changed

Lines changed: 161 additions & 108 deletions

File tree

‎infer/infer.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ func (inf *Infer) completionHandler(c echo.Context) error {
103103
ginCtx.Request = ginCtx.Request.WithContext(ctx)
104104
}
105105

106-
inf.ProxyMan.ProxyOAIHandler(ginCtx)
106+
inf.ProxyMan.ProxyOAIHandler(ginCtx, false, false)
107107
return nil
108108
}
109109

@@ -117,7 +117,7 @@ func (inf *Infer) chatCompletionsHandler(c echo.Context) error {
117117
}
118118

119119
ginCtx := echo2ginWithBody(c, body)
120-
inf.ProxyMan.ProxyOAIHandler(ginCtx)
120+
inf.ProxyMan.ProxyOAIHandler(ginCtx, false, false)
121121
return nil
122122
}
123123

@@ -129,7 +129,7 @@ func (inf *Infer) proxyOAIHandler(c echo.Context) error {
129129
}
130130

131131
ginCtx := echo2ginWithBody(c, body)
132-
inf.ProxyMan.ProxyOAIHandler(ginCtx)
132+
inf.ProxyMan.ProxyOAIHandler(ginCtx, false, false)
133133
return nil
134134
}
135135

@@ -141,7 +141,7 @@ func (inf *Infer) proxyOAIPostFormHandler(c echo.Context) error {
141141
}
142142

143143
ginCtx := echo2ginWithBody(c, body)
144-
inf.ProxyMan.ProxyOAIHandler(ginCtx)
144+
inf.ProxyMan.ProxyOAIHandler(ginCtx, false, false)
145145
return nil
146146
}
147147

‎proxy/proxymanager.go‎

Lines changed: 130 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020

2121
"github.com/LM4eu/goinfer/conf"
2222
"github.com/LM4eu/goinfer/event"
23+
"github.com/LM4eu/goinfer/gie"
2324
"github.com/LM4eu/goinfer/proxy/config"
2425
"github.com/gin-gonic/gin"
2526
"github.com/tidwall/gjson"
@@ -242,30 +243,61 @@ func (pm *ProxyManager) setupGinEngine() {
242243
})
243244

244245
// Set up routes using the Gin engine
245-
pm.ginEngine.POST("/v1/chat/completions", pm.ProxyOAIHandler)
246+
pm.ginEngine.POST("/v1/chat/completions", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, false) })
247+
pm.ginEngine.POST("/d1/chat/completions", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, false) })
248+
pm.ginEngine.POST("/a1/chat/completions", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, true) })
249+
pm.ginEngine.POST("/A1/chat/completions", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, true) })
246250
// Support legacy /v1/completions api, see issue #12
247-
pm.ginEngine.POST("/v1/completions", pm.ProxyOAIHandler)
251+
pm.ginEngine.POST("/v1/completions", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, false) })
252+
pm.ginEngine.POST("/d1/completions", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, false) })
253+
pm.ginEngine.POST("/a1/completions", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, true) })
254+
pm.ginEngine.POST("/A1/completions", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, true) })
248255

249256
// Support embeddings and reranking
250-
pm.ginEngine.POST("/v1/embeddings", pm.ProxyOAIHandler)
257+
pm.ginEngine.POST("/v1/embeddings", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, false) })
258+
pm.ginEngine.POST("/d1/embeddings", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, false) })
259+
pm.ginEngine.POST("/a1/embeddings", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, true) })
260+
pm.ginEngine.POST("/A1/embeddings", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, true) })
251261

252262
// llama-server's /reranking endpoint + aliases
253-
pm.ginEngine.POST("/reranking", pm.ProxyOAIHandler)
254-
pm.ginEngine.POST("/rerank", pm.ProxyOAIHandler)
255-
pm.ginEngine.POST("/v1/rerank", pm.ProxyOAIHandler)
256-
pm.ginEngine.POST("/v1/reranking", pm.ProxyOAIHandler)
263+
pm.ginEngine.POST("/reranking", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, false) })
264+
pm.ginEngine.POST("/rerank", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, false) })
265+
pm.ginEngine.POST("/v1/rerank", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, false) })
266+
pm.ginEngine.POST("/d1/rerank", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, false) })
267+
pm.ginEngine.POST("/a1/rerank", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, true) })
268+
pm.ginEngine.POST("/A1/rerank", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, true) })
269+
pm.ginEngine.POST("/v1/reranking", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, false) })
270+
pm.ginEngine.POST("/d1/reranking", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, false) })
271+
pm.ginEngine.POST("/a1/reranking", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, true) })
272+
pm.ginEngine.POST("/A1/reranking", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, true) })
257273

258274
// llama-server's /infill endpoint for code infilling
259-
pm.ginEngine.POST("/infill", pm.ProxyOAIHandler)
275+
pm.ginEngine.POST("/v1/infill", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, false) })
276+
pm.ginEngine.POST("/d1/infill", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, false) })
277+
pm.ginEngine.POST("/a1/infill", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, true) })
278+
pm.ginEngine.POST("/A1/infill", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, true) })
260279

261280
// llama-server's /completion endpoint
262-
pm.ginEngine.POST("/completion", pm.ProxyOAIHandler)
281+
pm.ginEngine.POST("/completion", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, false) })
282+
pm.ginEngine.POST("/d/completion", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, false) })
283+
pm.ginEngine.POST("/a/completion", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, true) })
284+
pm.ginEngine.POST("/A/completion", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, true) })
263285

264286
// Support audio/speech endpoint
265-
pm.ginEngine.POST("/v1/audio/speech", pm.ProxyOAIHandler)
266-
pm.ginEngine.POST("/v1/audio/transcriptions", pm.ProxyOAIPostFormHandler)
287+
pm.ginEngine.POST("/v1/audio/speech", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, false) })
288+
pm.ginEngine.POST("/d1/audio/speech", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, false) })
289+
pm.ginEngine.POST("/a1/audio/speech", func(c *gin.Context) { pm.ProxyOAIHandler(c, false, true) })
290+
pm.ginEngine.POST("/A1/audio/speech", func(c *gin.Context) { pm.ProxyOAIHandler(c, true, true) })
291+
pm.ginEngine.POST("/v1/audio/transcriptions", func(c *gin.Context) { pm.ProxyOAIPostFormHandler(c, false, false) })
292+
pm.ginEngine.POST("/d1/audio/transcriptions", func(c *gin.Context) { pm.ProxyOAIPostFormHandler(c, true, false) })
293+
pm.ginEngine.POST("/a1/audio/transcriptions", func(c *gin.Context) { pm.ProxyOAIPostFormHandler(c, false, true) })
294+
pm.ginEngine.POST("/A1/audio/transcriptions", func(c *gin.Context) { pm.ProxyOAIPostFormHandler(c, true, true) })
267295

268296
pm.ginEngine.GET("/v1/models", pm.ListModelsHandler)
297+
pm.ginEngine.GET("/d1/models", pm.ListModelsHandler)
298+
pm.ginEngine.GET("/a1/models", pm.ListModelsHandler)
299+
pm.ginEngine.GET("/A1/models", pm.ListModelsHandler)
300+
pm.ginEngine.GET("/gi/models", pm.GiListModelsHandler)
269301

270302
// in proxymanager_loghandlers.go
271303
pm.ginEngine.GET("/logs", pm.sendLogsHandlers)
@@ -470,6 +502,18 @@ func (pm *ProxyManager) ListModelsHandler(c *gin.Context) {
470502
})
471503
}
472504

505+
func (pm *ProxyManager) GiListModelsHandler(c *gin.Context) {
506+
models := pm.cfg.ListModels()
507+
508+
if len(models) == 0 {
509+
c.String(http.StatusNoContent, `{"count":0, "models":[]}`)
510+
return
511+
}
512+
513+
// Use gin's JSON method which handles content-type and encoding
514+
c.JSON(http.StatusOK, gin.H{"count": len(models), "models": models})
515+
}
516+
473517
func (pm *ProxyManager) proxyToUpstream(c *gin.Context) {
474518
upstreamPath := c.Param("upstreamPath")
475519

@@ -553,55 +597,35 @@ func (pm *ProxyManager) proxyToUpstream(c *gin.Context) {
553597
}
554598
}
555599

556-
func (pm *ProxyManager) ProxyOAIHandler(c *gin.Context) {
557-
bodyBytes, err := io.ReadAll(c.Request.Body)
558-
if err != nil {
559-
pm.sendErrorResponse(c, http.StatusBadRequest, "could not ready request body")
560-
return
561-
}
562-
563-
requestedModel := gjson.GetBytes(bodyBytes, "model").String()
564-
if requestedModel == "" {
565-
// fallback: the first running process we find
566-
for _, processGroup := range pm.processGroups {
567-
for _, process := range processGroup.processes {
568-
if process.CurrentState() == StateReady {
569-
requestedModel = process.ID
570-
}
571-
}
572-
}
573-
if requestedModel == "" {
574-
pm.sendErrorResponse(c, http.StatusBadRequest, "missing or invalid 'model' key")
575-
return
600+
func (pm *ProxyManager) ProxyOAIHandler(c *gin.Context, download, agentSmith bool) {
601+
if download || agentSmith {
602+
if c.Request.URL.Path[2] == '/' {
603+
c.Request.URL.Path = c.Request.URL.Path[2:]
604+
} else {
605+
c.Request.URL.Path = "/v" + c.Request.URL.Path[2:]
576606
}
577607
}
578608

579-
realModelName, found := pm.config.RealModelName(requestedModel)
580-
if !found {
581-
pm.sendErrorResponse(c, http.StatusBadRequest, "could not find real modelID for "+requestedModel)
609+
bodyBytes, realModelName, download, agentSmith, err := pm.getSetModel(c.Request.Body, download, agentSmith)
610+
if err != nil {
611+
c.JSON(http.StatusBadRequest, err)
582612
return
583613
}
584-
585-
processGroup, _, err := pm.swapProcessGroup(realModelName)
586-
if err != nil {
587-
pm.sendErrorResponse(c, http.StatusInternalServerError, "error swapping process group: "+err.Error())
614+
if realModelName == "" {
615+
c.JSON(http.StatusBadRequest, gie.New(gie.Invalid, "no model provided and no model loaded"))
588616
return
589617
}
590618

591-
// issue #69 allow custom model names to be sent to upstream
592-
useModelName := pm.config.Models[realModelName].UseModelName
593-
if useModelName != "" {
594-
bodyBytes, err = sjson.SetBytes(bodyBytes, "model", useModelName)
595-
if err != nil {
596-
pm.sendErrorResponse(c, http.StatusInternalServerError, "error rewriting model name in JSON: "+err.Error())
597-
return
598-
}
619+
processGroup, realModelName, err := pm.swapProcessGroup(realModelName)
620+
if err != nil {
621+
c.JSON(http.StatusInternalServerError, gie.Wrap(err, gie.Invalid, "error swapping process group", "model", realModelName))
622+
return
599623
}
600624

601625
// issue #174 strip parameters from the JSON body
602626
stripParams, err := pm.config.Models[realModelName].Filters.SanitizedStripParams()
603627
if err != nil { // just log it and continue
604-
pm.proxyLogger.Errorf("Error sanitizing strip params string: %s, %s", pm.config.Models[realModelName].Filters.StripParams, err.Error())
628+
pm.proxyLogger.Warnf("Error sanitizing strip params string: %s, %s", pm.config.Models[realModelName].Filters.StripParams, err.Error())
605629
} else {
606630
for _, param := range stripParams {
607631
pm.proxyLogger.Debugf("<%s> stripping param: %s", realModelName, param)
@@ -631,19 +655,19 @@ func (pm *ProxyManager) ProxyOAIHandler(c *gin.Context) {
631655
if err != nil {
632656
pm.sendErrorResponse(c, http.StatusInternalServerError, "error proxying metrics wrapped request: "+err.Error())
633657
pm.proxyLogger.Errorf("Error Proxying Metrics Wrapped Request for processGroup %s and model %s", processGroup.id, realModelName)
634-
return
635-
}
636-
} else {
637-
err := processGroup.ProxyRequest(realModelName, c.Writer, c.Request)
638-
if err != nil {
639-
pm.sendErrorResponse(c, http.StatusInternalServerError, "error proxying request: "+err.Error())
640-
pm.proxyLogger.Errorf("Error Proxying Request for processGroup %s and model %s", processGroup.id, realModelName)
641-
return
642658
}
659+
return
660+
}
661+
662+
err = processGroup.ProxyRequest(realModelName, c.Writer, c.Request)
663+
if err != nil {
664+
pm.sendErrorResponse(c, http.StatusInternalServerError, "error proxying request: "+err.Error())
665+
pm.proxyLogger.Errorf("Error Proxying Request for processGroup %s and model %s", processGroup.id, realModelName)
666+
return
643667
}
644668
}
645669

646-
func (pm *ProxyManager) ProxyOAIPostFormHandler(c *gin.Context) {
670+
func (pm *ProxyManager) ProxyOAIPostFormHandler(c *gin.Context, download, agentSmith bool) {
647671
// Parse multipart form
648672
if err := c.Request.ParseMultipartForm(32 << 20); err != nil { // 32MB max memory, larger files go to tmp disk
649673
pm.sendErrorResponse(c, http.StatusBadRequest, "error parsing multipart form: "+err.Error())
@@ -652,22 +676,9 @@ func (pm *ProxyManager) ProxyOAIPostFormHandler(c *gin.Context) {
652676

653677
// Get model parameter from the form
654678
requestedModel := c.Request.FormValue("model")
655-
if requestedModel == "" {
656-
// fallback: the first running process we find
657-
for _, processGroup := range pm.processGroups {
658-
for _, process := range processGroup.processes {
659-
if process.CurrentState() == StateReady {
660-
requestedModel = process.ID
661-
}
662-
}
663-
}
664-
if requestedModel == "" {
665-
pm.sendErrorResponse(c, http.StatusBadRequest, "missing or invalid 'model' parameter in form data")
666-
return
667-
}
668-
}
679+
fixed, download, agentSmith := pm.fixModelName(requestedModel, download, agentSmith)
669680

670-
processGroup, realModelName, err := pm.swapProcessGroup(requestedModel)
681+
processGroup, realModelName, err := pm.swapProcessGroup(fixed)
671682
if err != nil {
672683
pm.sendErrorResponse(c, http.StatusInternalServerError, "error swapping process group: "+err.Error())
673684
return
@@ -780,7 +791,18 @@ func (pm *ProxyManager) UnloadAllModelsHandler(c *gin.Context) {
780791

781792
func (pm *ProxyManager) ListRunningProcessesHandler(context *gin.Context) {
782793
context.Header("Content-Type", "application/json")
783-
runningProcesses := make([]gin.H, 0) // Default to an empty response.
794+
795+
// Put the results under the `running` key.
796+
response := gin.H{
797+
"running": pm.listRunningProcesses(),
798+
}
799+
800+
context.JSON(http.StatusOK, response) // Always return 200 OK
801+
}
802+
803+
func (pm *ProxyManager) listRunningProcesses() []gin.H {
804+
count := pm.countRunningProcesses()
805+
runningProcesses := make([]gin.H, count) // Default to an empty response.
784806

785807
for _, processGroup := range pm.processGroups {
786808
for _, process := range processGroup.processes {
@@ -792,13 +814,21 @@ func (pm *ProxyManager) ListRunningProcessesHandler(context *gin.Context) {
792814
}
793815
}
794816
}
817+
return runningProcesses
818+
}
795819

796-
// Put the results under the `running` key.
797-
response := gin.H{
798-
"running": runningProcesses,
820+
func (pm *ProxyManager) countRunningProcesses() int {
821+
count := 0
822+
823+
for _, processGroup := range pm.processGroups {
824+
for _, process := range processGroup.processes {
825+
if process.CurrentState() == StateReady {
826+
count++
827+
}
828+
}
799829
}
800830

801-
context.JSON(http.StatusOK, response) // Always return 200 OK
831+
return count
802832
}
803833

804834
func (pm *ProxyManager) findGroupByModelName(modelName string) *ProcessGroup {
@@ -818,15 +848,37 @@ func (pm *ProxyManager) SetVersion(buildDate, commit, version string) {
818848
pm.version = version
819849
}
820850

821-
// ProxyToFirstRunningProcess forwards the request to the any running process (llama-server).
851+
func (pm *ProxyManager) firstRunningProcess() string {
852+
var starting string
853+
for _, processGroup := range pm.processGroups {
854+
for _, process := range processGroup.processes {
855+
if process.state == StateReady {
856+
return process.ID
857+
}
858+
if process.state == StateStarting {
859+
starting = process.ID
860+
}
861+
}
862+
}
863+
return starting
864+
}
865+
866+
// ProxyToFirstRunningProcess forwards the request to a running process (llama-server).
822867
func (pm *ProxyManager) ProxyToFirstRunningProcess(c *gin.Context) {
868+
var starting *Process
823869
for _, processGroup := range pm.processGroups {
824870
for _, process := range processGroup.processes {
825-
if process.CurrentState() == StateReady {
871+
if process.state == StateReady {
826872
process.ProxyRequest(c.Writer, c.Request)
827873
return
828874
}
875+
if process.state == StateStarting {
876+
starting = process
877+
}
829878
}
830879
}
880+
if starting != nil {
881+
starting.ProxyRequest(c.Writer, c.Request)
882+
}
831883
pm.sendErrorResponse(c, http.StatusInternalServerError, "No model currently running. Please select a model.")
832884
}

0 commit comments

Comments
 (0)