From c2d03eb5a2708fc8ac814f9347655aa51fc2231f Mon Sep 17 00:00:00 2001
From: Bruno Michel <bruno@cozycloud.cc>
Date: Mon, 21 Oct 2024 11:23:47 +0200
Subject: [PATCH] Add sources with chat messages

---
 docs/ai.md        |  4 ++++
 model/rag/chat.go | 28 ++++++++++++++++++++--------
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/docs/ai.md b/docs/ai.md
index beaf403e48a..43263f38ae4 100644
--- a/docs/ai.md
+++ b/docs/ai.md
@@ -114,6 +114,10 @@ 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": "sources", "content": [{"doc_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
index 9980c9f98f4..f79ebc2562e 100644
--- 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,6 +150,7 @@ 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":
@@ -161,6 +164,14 @@ func Query(inst *instance.Instance, logger logger.Logger, query QueryMessage) er
 			}
 			delta.SetID(msg.ID)
 			go realtime.GetHub().Publish(inst, realtime.EventCreate, &delta, nil)
+		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:
 			// We can ignore done events
 		}
@@ -174,6 +185,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)