andreas/llama-swap
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix model alias usage in upstream path (#230)

Browse Source 
Model alias values are not properly resolved and work in upstream/ path.

Related to #229.
This commit is contained in:
Benson Wong
2025-08-07 20:16:56 -07:00
committed by GitHub
parent 5b10b3c23f
commit 10569ed546
2 changed files with 65 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
proxy/proxymanager.go
View File

@@ -361,7 +361,7 @@ func (pm *ProxyManager) proxyToUpstream(c *gin.Context) {
return
}
processGroup, _, err := pm.swapProcessGroup(requestedModel)
processGroup, realModelName, err := pm.swapProcessGroup(requestedModel)
if err != nil {
pm.sendErrorResponse(c, http.StatusInternalServerError, fmt.Sprintf("error swapping process group: %s", err.Error()))
return
@@ -369,7 +369,7 @@ func (pm *ProxyManager) proxyToUpstream(c *gin.Context) {
// rewrite the path
c.Request.URL.Path = c.Param("upstreamPath")
processGroup.ProxyRequest(requestedModel, c.Writer, c.Request)
processGroup.ProxyRequest(realModelName, c.Writer, c.Request)
}
func (pm *ProxyManager) proxyOAIHandler(c *gin.Context) {

28
proxy/proxymanager_test.go
View File

@@ -9,6 +9,7 @@ import (
"net/http"
"net/http/httptest"
"strconv"
"strings"
"sync"
"testing"
"time"
@@ -656,21 +657,34 @@ func TestProxyManager_CORSOptionsHandler(t *testing.T) {
}
func TestProxyManager_Upstream(t *testing.T) {
config := AddDefaultGroupToConfig(Config{
HealthCheckTimeout: 15,
Models: map[string]ModelConfig{
"model1": getTestSimpleResponderConfig("model1"),
},
LogLevel: "error",
})
configStr := fmt.Sprintf(`
logLevel: error
models:
model1:
cmd: %s -port ${PORT} -silent -respond model1
aliases: [model-alias]
`, getSimpleResponderPath())
config, err := LoadConfigFromReader(strings.NewReader(configStr))
assert.NoError(t, err)
proxy := New(config)
defer proxy.StopProcesses(StopWaitForInflightRequest)
t.Run("main model name", func(t *testing.T) {
req := httptest.NewRequest("GET", "/upstream/model1/test", nil)
rec := httptest.NewRecorder()
proxy.ServeHTTP(rec, req)
assert.Equal(t, http.StatusOK, rec.Code)
assert.Equal(t, "model1", rec.Body.String())
})
t.Run("model alias", func(t *testing.T) {
req := httptest.NewRequest("GET", "/upstream/model-alias/test", nil)
rec := httptest.NewRecorder()
proxy.ServeHTTP(rec, req)
assert.Equal(t, http.StatusOK, rec.Code)
assert.Equal(t, "model1", rec.Body.String())
})
}
func TestProxyManager_ChatContentLength(t *testing.T) {