third-party-mirrors/ollama
1
0
Fork 1
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge 091037857f4d5f0c15d8055b74792f74ef052b4a into d7eb05b9361febead29a74e71ddffc2ebeff5302

Browse Source
This commit is contained in:
Kemal Elmizan 2024-11-14 13:55:49 +08:00 committed by GitHub
commit c2202bb8f4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 52 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
envconfig/config.go
View File

@ -180,6 +180,8 @@ var (
RocrVisibleDevices = String("ROCR_VISIBLE_DEVICES")
GpuDeviceOrdinal = String("GPU_DEVICE_ORDINAL")
HsaOverrideGfxVersion = String("HSA_OVERRIDE_GFX_VERSION")
BasicAuthKey = String("OLLAMA_BASIC_AUTH_KEY")
)
func Uint(key string, defaultValue uint) func() uint {
@ -249,6 +251,7 @@ func AsMap() map[string]EnvVar {
"OLLAMA_SCHED_SPREAD": {"OLLAMA_SCHED_SPREAD", SchedSpread(), "Always schedule model across all GPUs"},
"OLLAMA_TMPDIR": {"OLLAMA_TMPDIR", TmpDir(), "Location for temporary files"},
"OLLAMA_MULTIUSER_CACHE": {"OLLAMA_MULTIUSER_CACHE", MultiUserCache(), "Optimize prompt caching for multi-user scenarios"},
"OLLAMA_BASIC_AUTH_KEY": {"OLLAMA_BASIC_AUTH_KEY", BasicAuthKey(), "Basic auth key for user ollama, will allow all request if empty"},
// Informational
"HTTP_PROXY": {"HTTP_PROXY", String("HTTP_PROXY")(), "HTTP proxy"},

9
server/routes.go
View File

@ -1145,6 +1145,15 @@ func (s *Server) GenerateRoutes() http.Handler {
allowedHostsMiddleware(s.addr),
)
ollamaAuthKey := envconfig.BasicAuthKey()
if ollamaAuthKey != "" {
r.Use(
gin.BasicAuth(gin.Accounts{
"ollama": ollamaAuthKey,
}),
)
}
r.POST("/api/pull", s.PullHandler)
r.POST("/api/generate", s.GenerateHandler)
r.POST("/api/chat", s.ChatHandler)

40
server/routes_test.go
View File

@ -438,6 +438,46 @@ func Test_Routes(t *testing.T) {
}
},
},
{
Name: "Authenticated Route Success",
Method: http.MethodGet,
Path: "/api/version",
Setup: func(t *testing.T, req *http.Request) {
t.Setenv("OLLAMA_BASIC_AUTH_KEY", "password")
req.SetBasicAuth("ollama", "password")
},
Expected: func(t *testing.T, resp *http.Response) {
assert.Equal(t, http.StatusOK, resp.StatusCode)
_, err := io.ReadAll(resp.Body)
require.NoError(t, err)
},
},
{
Name: "Authenticated Route Failure",
Method: http.MethodGet,
Path: "/api/version",
Setup: func(t *testing.T, req *http.Request) {
t.Setenv("OLLAMA_BASIC_AUTH_KEY", "password")
req.SetBasicAuth("ollama", "wrongpassword")
},
Expected: func(t *testing.T, resp *http.Response) {
assert.Equal(t, http.StatusUnauthorized, resp.StatusCode)
},
},
{
Name: "BasicAuthKey not set",
Method: http.MethodGet,
Path: "/api/version",
Setup: func(t *testing.T, req *http.Request) {
t.Setenv("OLLAMA_BASIC_AUTH_KEY", "")
req.SetBasicAuth("ollama", "wrongpassword")
},
Expected: func(t *testing.T, resp *http.Response) {
assert.Equal(t, http.StatusOK, resp.StatusCode)
_, err := io.ReadAll(resp.Body)
require.NoError(t, err)
},
},
}
t.Setenv("OLLAMA_MODELS", t.TempDir())