Merge pull request #98 from gen-mind/fix_sources

fix_sources
gen-mind · May 13, 2024 · 787c781 · 787c781
2 parents b638957 + 0284dd2
commit 787c781
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 16 deletions.
diff --git a/backend/core/bll/chat.go b/backend/core/bll/chat.go
@@ -63,7 +63,7 @@ func (b *chatBL) SendMessage(ctx *gin.Context, user *model.User, param *paramete
 	aiClient := b.aiBuilder.New(chatSession.Persona.LLM)
 	resp := responder.NewManager(
 		responder.NewAIResponder(aiClient, b.chatRepo),
-		responder.NewEmbeddingResponder())
+	)
 
 	go resp.Send(ctx, &message)
 	return resp, nil
@@ -81,14 +81,15 @@ func (b *chatBL) GetSessionByID(ctx context.Context, user *model.User, id int64)
 	docs := make([]model.DocumentResponse, 0)
 	for i := 0; i < 4; i++ {
 		docs = append(docs, model.DocumentResponse{
-			ID:         decimal.NewFromInt(int64(i)),
-			DocumentID: "11",
-			Link:       fmt.Sprintf("link for document %d", i),
-			Content:    fmt.Sprintf("content of document %d", i),
+			ID:          decimal.NewFromInt(int64(i)),
+			DocumentID:  "11",
+			Link:        fmt.Sprintf("link for document %d", i),
+			Content:     fmt.Sprintf("content of document %d", i),
+			UpdatedDate: time.Now().UTC().Add(-48 * time.Hour),
 		})
 	}
 	for _, msg := range result.Messages {
-		if msg.MessageType == model.MessageTypeUser {
+		if msg.MessageType == model.MessageTypeAssistant {
 			for _, d := range docs {
 				md := d
 				md.MessageID = msg.ID

diff --git a/backend/core/model/document.go b/backend/core/model/document.go
@@ -32,11 +32,12 @@ type Document struct {
 }
 
 type DocumentResponse struct {
-	ID         decimal.Decimal `json:"id,omitempty"`
-	MessageID  decimal.Decimal `json:"message_id,omitempty"`
-	Link       string          `json:"link,omitempty"`
-	DocumentID string          `json:"document_id,omitempty"`
-	Content    string          `json:"content,omitempty"`
+	ID          decimal.Decimal `json:"id,omitempty"`
+	MessageID   decimal.Decimal `json:"message_id,omitempty"`
+	Link        string          `json:"link,omitempty"`
+	DocumentID  string          `json:"document_id,omitempty"`
+	Content     string          `json:"content,omitempty"`
+	UpdatedDate time.Time       `json:"updated_date,omitempty"`
 }
 type DocumentFeedback struct {
 	tableName    struct{}        `pg:"document_feedbacks"`

diff --git a/backend/core/responder/ai.go b/backend/core/responder/ai.go
@@ -5,6 +5,8 @@ import (
 	"cognix.ch/api/v2/core/model"
 	"cognix.ch/api/v2/core/repository"
 	"context"
+	"fmt"
+	"github.com/shopspring/decimal"
 	"sync"
 	"time"
 )
@@ -42,6 +44,23 @@ func (r *aiResponder) Send(ctx context.Context, ch chan *Response, wg *sync.Wait
 		payload.Type = ResponseError
 	}
 	ch <- payload
+	time.Sleep(10 * time.Millisecond)
+	for i := 0; i < 4; i++ {
+		ch <- &Response{
+			IsValid: true,
+			Type:    ResponseDocument,
+			Message: nil,
+			Document: &model.DocumentResponse{
+				ID:          decimal.NewFromInt(int64(i)),
+				DocumentID:  "11",
+				Link:        fmt.Sprintf("link for document %d", i),
+				Content:     fmt.Sprintf("content of document %d", i),
+				UpdatedDate: time.Now().UTC().Add(-48 * time.Hour),
+				MessageID:   message.ID,
+			},
+		}
+	}
+
 	wg.Done()
 }
 

diff --git a/backend/core/responder/embedding.go b/backend/core/responder/embedding.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"github.com/shopspring/decimal"
 	"sync"
+	"time"
 )
 
 type embedding struct {
@@ -19,11 +20,12 @@ func (r *embedding) Send(ctx context.Context, ch chan *Response, wg *sync.WaitGr
 			Type:    ResponseDocument,
 			Message: nil,
 			Document: &model.DocumentResponse{
-				ID:         decimal.NewFromInt(int64(i)),
-				DocumentID: "11",
-				Link:       fmt.Sprintf("link for document %d", i),
-				Content:    fmt.Sprintf("content of document %d", i),
-				MessageID:  parentMessage.ID,
+				ID:          decimal.NewFromInt(int64(i)),
+				DocumentID:  "11",
+				Link:        fmt.Sprintf("link for document %d", i),
+				Content:     fmt.Sprintf("content of document %d", i),
+				UpdatedDate: time.Now().UTC().Add(-48 * time.Hour),
+				MessageID:   parentMessage.ID,
 			},
 		}
 	}