Add sources with chat messages (#4478)

cozy · Oct 23, 2024 · 232c522 · 232c522
2 parents 669cd69 + a67bdcb
commit 232c522
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 11 deletions.
diff --git a/docs/ai.md b/docs/ai.md
@@ -114,6 +114,14 @@ server > {"event": "CREATED",
                       "type": "io.cozy.ai.chat.events",
                       "doc": {"object": "delta", "content": "sky ", "position": 1}}}
 [...]
+server > {"event": "CREATED",
+          "payload": {"id": "eb17c3205bf1013ddea018c04daba326",
+                      "type": "io.cozy.ai.chat.events",
+                      "doc": {"object": "generated"}}}
+server > {"event": "CREATED",
+          "payload": {"id": "eb17c3205bf1013ddea018c04daba326",
+                      "type": "io.cozy.ai.chat.events",
+                      "doc": {"object": "sources", "content": [{"id": "827f0fbb928b375cc457c732a4013aa7", "doctype": "io.cozy.files"}]}}}
 server > {"event": "CREATED",
           "payload": {"id": "eb17c3205bf1013ddea018c04daba326",
                       "type": "io.cozy.ai.chat.events",

diff --git a/model/rag/chat.go b/model/rag/chat.go
@@ -37,10 +37,11 @@ type ChatConversation struct {
 }
 
 type ChatMessage struct {
-	ID        string    `json:"id"`
-	Role      string    `json:"role"`
-	Content   string    `json:"content"`
-	CreatedAt time.Time `json:"createdAt"`
+	ID        string        `json:"id"`
+	Role      string        `json:"role"`
+	Content   string        `json:"content"`
+	Sources   []interface{} `json:"sources,omitempty"`
+	CreatedAt time.Time     `json:"createdAt"`
 }
 
 const (
@@ -130,10 +131,11 @@ func Query(inst *instance.Instance, logger logger.Logger, query QueryMessage) er
 	myself, _ := contact.GetMyself(inst)
 	relatives, _ := getRelatives(inst, myself)
 	payload := map[string]interface{}{
-		"messages":  chat.Messages,
-		"myself":    myself,
-		"relatives": relatives,
-		"stream":    true,
+		"messages":     chat.Messages,
+		"myself":       myself,
+		"relatives":    relatives,
+		"stream":       true,
+		"with_sources": true,
 	}
 
 	res, err := callRAGQuery(inst, payload)
@@ -148,9 +150,10 @@ func Query(inst *instance.Instance, logger logger.Logger, query QueryMessage) er
 	msg := chat.Messages[len(chat.Messages)-1]
 	position := 0
 	var completion string
+	var sources []interface{}
 	err = foreachSSE(res.Body, func(event map[string]interface{}) {
 		switch event["object"] {
-		case "delta", "done":
+		case "delta":
 			event["position"] = position
 			position++
 			content, _ := event["content"].(string)
@@ -161,8 +164,21 @@ func Query(inst *instance.Instance, logger logger.Logger, query QueryMessage) er
 			}
 			delta.SetID(msg.ID)
 			go realtime.GetHub().Publish(inst, realtime.EventCreate, &delta, nil)
-		default:
-			// We can ignore done events
+		case "sources":
+			sources, _ = event["content"].([]interface{})
+			doc := couchdb.JSONDoc{
+				Type: consts.ChatEvents,
+				M:    event,
+			}
+			doc.SetID(msg.ID)
+			go realtime.GetHub().Publish(inst, realtime.EventCreate, &doc, nil)
+		default: // done, generated, etc.
+			doc := couchdb.JSONDoc{
+				Type: consts.ChatEvents,
+				M:    event,
+			}
+			doc.SetID(msg.ID)
+			go realtime.GetHub().Publish(inst, realtime.EventCreate, &doc, nil)
 		}
 	})
 	if err != nil {
@@ -174,6 +190,7 @@ func Query(inst *instance.Instance, logger logger.Logger, query QueryMessage) er
 		ID:        uuidv7.String(),
 		Role:      AssistantRole,
 		Content:   completion,
+		Sources:   sources,
 		CreatedAt: time.Now().UTC(),
 	}
 	chat.Messages = append(chat.Messages, answer)