Add auth for RAG server

cozy · Sep 25, 2024 · 6ebb03b · 6ebb03b
1 parent 97919af
commit 6ebb03b
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 2 deletions.
diff --git a/cozy.example.yaml b/cozy.example.yaml
@@ -208,8 +208,10 @@ rag:
   # declared, for default.
   default:
     url: http://localhost:8000
+    api_key: $3cr3t
   beta:
     url: http://localhost:8001
+    api_key: $3cr3t
 
 # mail service parameters for sending email via SMTP
 mail:

diff --git a/model/rag/chat.go b/model/rag/chat.go
@@ -190,6 +190,7 @@ func callRAGQuery(inst *instance.Instance, payload map[string]interface{}) (*htt
 	if err != nil {
 		return nil, err
 	}
+	req.Header.Add(echo.HeaderAuthorization, "Bearer "+ragServer.APIKey)
 	req.Header.Add("Content-Type", echo.MIMEApplicationJSON)
 	return http.DefaultClient.Do(req)
 }

diff --git a/model/rag/index.go b/model/rag/index.go
@@ -16,6 +16,7 @@ import (
 	"github.com/cozy/cozy-stack/pkg/couchdb"
 	"github.com/cozy/cozy-stack/pkg/couchdb/revision"
 	"github.com/cozy/cozy-stack/pkg/logger"
+	"github.com/labstack/echo/v4"
 )
 
 // BatchSize is the maximal number of documents manipulated at once by the
@@ -86,6 +87,7 @@ func callRAGIndexer(inst *instance.Instance, doctype string, change couchdb.Chan
 		if err != nil {
 			return err
 		}
+		req.Header.Add(echo.HeaderAuthorization, "Bearer "+ragServer.APIKey)
 		res, err := http.DefaultClient.Do(req)
 		if err != nil {
 			return err
@@ -100,6 +102,7 @@ func callRAGIndexer(inst *instance.Instance, doctype string, change couchdb.Chan
 		if err != nil {
 			return err
 		}
+		req.Header.Add(echo.HeaderAuthorization, "Bearer "+ragServer.APIKey)
 		res, err := http.DefaultClient.Do(req)
 		if err != nil {
 			return err
@@ -154,6 +157,7 @@ func callRAGIndexer(inst *instance.Instance, doctype string, change couchdb.Chan
 		if err != nil {
 			return err
 		}
+		req.Header.Add(echo.HeaderAuthorization, "Bearer "+ragServer.APIKey)
 		req.Header.Add("Content-Type", mime)
 		res, err = http.DefaultClient.Do(req)
 		if err != nil {

diff --git a/pkg/config/config/config.go b/pkg/config/config/config.go
@@ -232,7 +232,8 @@ type Office struct {
 
 // RAGServer contains the configuration for a RAG server (AI features).
 type RAGServer struct {
-	URL string
+	URL    string
+	APIKey string
 }
 
 // Notifications contains the configuration for the mobile push-notification
@@ -1026,8 +1027,10 @@ func makeRAGServers(v *viper.Viper) (map[string]RAGServer, error) {
 					"should be a map, got %#v", v)
 		}
 		url, _ := m["url"].(string)
+		key, _ := m["api_key"].(string)
 		servers[k] = RAGServer{
-			URL: url,
+			URL:    url,
+			APIKey: key,
 		}
 	}
 	return servers, nil