Merge pull request #819 from bjwswang/main

feat: convert message.Files to full Documents

apiserver/docs/docs.go

@@ -1941,12 +1941,19 @@ const docTemplate = `{
                     "example": "旷工最小计算单位为0.5天。"
                 },
                 "documents": {
-                    "description": "Docs uploaded in this message",
+                    "description": "For Action Upload",
                     "type": "array",
                     "items": {
                         "$ref": "#/definitions/storage.Document"
                     }
                 },
+                "files": {
+                    "description": "Files that shall be used in this Chat",
+                    "type": "array",
+                    "items": {
+                        "type": "string"
+                    }
+                },
                 "id": {
                     "type": "string",
                     "example": "4f3546dd-5404-4bf8-a3bc-4fa3f9a7ba24"
@@ -1956,6 +1963,7 @@ const docTemplate = `{
                     "example": 1000
                 },
                 "query": {
+                    "description": "For Action Chat",
                     "type": "string",
                     "example": "旷工最小计算单位为多少天？"
                 },

apiserver/docs/swagger.json

@@ -1935,12 +1935,19 @@
                     "example": "旷工最小计算单位为0.5天。"
                 },
                 "documents": {
-                    "description": "Docs uploaded in this message",
+                    "description": "For Action Upload",
                     "type": "array",
                     "items": {
                         "$ref": "#/definitions/storage.Document"
                     }
                 },
+                "files": {
+                    "description": "Files that shall be used in this Chat",
+                    "type": "array",
+                    "items": {
+                        "type": "string"
+                    }
+                },
                 "id": {
                     "type": "string",
                     "example": "4f3546dd-5404-4bf8-a3bc-4fa3f9a7ba24"
@@ -1950,6 +1957,7 @@
                     "example": 1000
                 },
                 "query": {
+                    "description": "For Action Chat",
                     "type": "string",
                     "example": "旷工最小计算单位为多少天？"
                 },

apiserver/docs/swagger.yaml

@@ -453,17 +453,23 @@ definitions:
         example: 旷工最小计算单位为0.5天。
         type: string
       documents:
-        description: Docs uploaded in this message
+        description: For Action Upload
         items:
           $ref: '#/definitions/storage.Document'
         type: array
+      files:
+        description: Files that shall be used in this Chat
+        items:
+          type: string
+        type: array
       id:
         example: 4f3546dd-5404-4bf8-a3bc-4fa3f9a7ba24
         type: string
       latency:
         example: 1000
         type: integer
       query:
+        description: For Action Chat
         example: 旷工最小计算单位为多少天？
         type: string
       references:

apiserver/pkg/chat/chat_docs.go

@@ -69,6 +69,10 @@ func (cs *ChatServer) ReceiveConversationFile(ctx context.Context, messageID str
 			User:         currentUser,
 			Debug:        req.Debug,
 		}
+		// create before upload documents
+		if err := cs.Storage().UpdateConversation(conversation); err != nil {
+			return nil, err
+		}
 	}
 
 	// upload files to system datasource

apiserver/pkg/chat/chat_server.go

@@ -136,6 +136,10 @@ func (cs *ChatServer) AppRun(ctx context.Context, req ChatReqBody, respStream ch
 			User:         currentUser,
 			Debug:        req.Debug,
 		}
+		// create before do AppRun
+		if err := cs.Storage().UpdateConversation(conversation); err != nil {
+			return nil, err
+		}
 	}
 	conversation.Messages = append(conversation.Messages, storage.Message{
 		ID:     messageID,
@@ -157,6 +161,10 @@ func (cs *ChatServer) AppRun(ctx context.Context, req ChatReqBody, respStream ch
 	conversation.Messages[len(conversation.Messages)-1].Answer = out.Answer
 	conversation.Messages[len(conversation.Messages)-1].References = out.References
 	conversation.Messages[len(conversation.Messages)-1].Latency = time.Since(req.StartTime).Milliseconds()
+	if req.Files != nil && len(req.Files) > 0 {
+		conversation.Messages[len(conversation.Messages)-1].RawFiles = strings.Join(req.Files, ",")
+	}
+
 	if err := cs.Storage().UpdateConversation(conversation); err != nil {
 		return nil, err
 	}

apiserver/pkg/chat/storage/storage.go

@@ -18,6 +18,7 @@ package storage
 
 import (
 	"errors"
+	"strings"
 	"time"
 
 	"gorm.io/gorm"
@@ -44,21 +45,31 @@ type Conversation struct {
 
 // Message represent a message in storage
 type Message struct {
-	ID string `gorm:"column:id;primaryKey;type:uuid;comment:message id" json:"id" example:"4f3546dd-5404-4bf8-a3bc-4fa3f9a7ba24"`
+	ID             string `gorm:"column:id;primaryKey;type:uuid;comment:message id" json:"id" example:"4f3546dd-5404-4bf8-a3bc-4fa3f9a7ba24"`
+	ConversationID string `gorm:"column:conversation_id;type:uuid;comment:conversation id" json:"-"`
+	Latency        int64  `gorm:"column:latency;type:int;comment:request latency, in ms" json:"latency" example:"1000"`
+
 	// Action indicates what is this message for
 	// Chat(by default),UPLOAD,etc...
 	Action string `gorm:"column:action;type:string;comment:user action" json:"action" example:"UPLOAD"`
 
-	Query          string     `gorm:"column:query;type:string;comment:user input" json:"query" example:"旷工最小计算单位为多少天？"`
-	Answer         string     `gorm:"column:answer;type:string;comment:ai response" json:"answer" example:"旷工最小计算单位为0.5天。"`
-	References     References `gorm:"column:references;type:json;comment:references" json:"references,omitempty"`
-	ConversationID string     `gorm:"column:conversation_id;type:uuid;comment:conversation id" json:"-"`
-	Latency        int64      `gorm:"column:latency;type:int;comment:request latency, in ms" json:"latency" example:"1000"`
+	// For Action Chat
+	Query string `gorm:"column:query;type:string;comment:user input" json:"query" example:"旷工最小计算单位为多少天？"`
+	// Files that shall be used in this Chat
+	Files      []string   `gorm:"-" json:"files"`
+	RawFiles   string     `gorm:"column:files;type:text[];comment:input files" json:"-"`
+	Answer     string     `gorm:"column:answer;type:string;comment:ai response" json:"answer" example:"旷工最小计算单位为0.5天。"`
+	References References `gorm:"column:references;type:json;comment:references" json:"references,omitempty"`
 
-	// Docs uploaded in this message
+	// For Action Upload
 	Documents []Document `gorm:"foreignKey:MessageID" json:"documents"`
 }
 
+func (m *Message) AfterFind(tx *gorm.DB) error {
+	m.Files = strings.Split(m.RawFiles, ",")
+	return nil
+}
+
 type Document struct {
 	ID             string `gorm:"column:id;primaryKey;type:uuid;comment:document id" json:"id" example:"4f3546dd-5404-4bf8-a3bc-4fa3f9a7ba24"`
 	Name           string `gorm:"column:name;type:string;comment:document name" json:"name" example:"kaoqin.pdf"`
@@ -85,6 +96,7 @@ func (Document) TableName() string {
 type Storage interface {
 	ConversationStorage
 	MessageStorage
+	DocumentStorage
 }
 
 // ConversationStorage interface
@@ -119,8 +131,5 @@ type MessageStorage interface {
 }
 
 type DocumentStorage interface {
-	// FindExistingDocuments finds a document in the message.
-	//
-	// It takes messageID, documentID string parameters and returns *Document, error.
-	FindExistingDocument(conversationID, messageID, documentID string, opts ...SearchOption) (*Document, error)
+	// TO BE DEFINED
 }

apiserver/pkg/chat/storage/storage_postgresql.go

@@ -103,7 +103,7 @@ func (p *PostgreSQLStorage) ListConversations(opts ...SearchOption) ([]Conversat
 	conversationQuery.Debug = false
 	conversationQuery.DeletedAt.Valid = false
 	res := make([]Conversation, 0)
-	tx := p.db.Preload("Messages").Order("updated_at DESC").Find(&res, conversationQuery)
+	tx := p.db.Preload("Messages.Documents").Order("updated_at DESC").Find(&res, conversationQuery)
 	if tx.Error != nil {
 		return nil, tx.Error
 	}
@@ -161,9 +161,30 @@ func (p *PostgreSQLStorage) FindExistingConversation(conversationID string, opts
 	if tx.Error != nil {
 		return nil, tx.Error
 	}
+
+	for index, message := range res.Messages {
+		// search document info based on object which is also a primary key in Document
+		if message.Action != "UPLOAD" && message.Files != nil && len(message.Files) > 0 {
+			documents, err := p.findMessageRelevantDocuments(message)
+			if err == nil {
+				message.Documents = documents
+				res.Messages[index] = message
+			}
+		}
+	}
+
 	return res, nil
 }
 
+func (p *PostgreSQLStorage) findMessageRelevantDocuments(message Message) ([]Document, error) {
+	var documents []Document
+	err := p.db.Where("object IN ?", message.Files).Find(&documents).Error
+	if err != nil {
+		return nil, err
+	}
+	return documents, nil
+}
+
 func (p *PostgreSQLStorage) Delete(opts ...SearchOption) error {
 	searchOpt := applyOptions(nil, opts...)
 	c := &Conversation{}