feat: processors for enriching documents

Author: SteveLauC

coco.yml

@@ -323,14 +323,27 @@ pipeline:
             group: enriched_documents
             fetch_max_messages: 10
           processor:
+            - read_local_file_content: {}
+            - file_extraction:
+                tika_endpoint: "http://127.0.0.1:9998"
+                timeout_in_seconds: 180
+                chunk_size: 7000
             - document_summarization:
-                model: $[[env.ENRICHMENT_MODEL]]
-                input_queue: "indexing_documents"
-                min_input_document_length: 500
+                model_provider: deepseek
+                model: deepseek-chat
+                model_context_length: 128000
+                max_summary_length: 300
+            - document_embedding:
+                model_provider: qianwen
+                model: text-embedding-v4
+                # Dimension of the vectors generated by your embedding model
+                embedding_dimension: 1024
+                # Chunk size in characters, it is recommended to use a value 
+                # lower than you embedding model's input limit, e.g., if the
+                # limit is 8192 tokens, then we can set chunk size to 7000.
+                chunk_size: 7000
                 output_queue:
                   name: "enriched_documents"
-                  label:
-                    tag: "enriched"
 
   - name: merge_documents
     auto_start: true
@@ -362,6 +375,7 @@ pipeline:
           queues:
             type: indexing_merge
             tag: "merged"
+
   - name: connector_dispatcher
     auto_start: true
     keep_running: true

core/document.go

@@ -5,6 +5,7 @@
 package core
 
 import (
+	"fmt"
 	"strings"
 	"time"
 )
@@ -41,8 +42,9 @@ type Document struct {
 	Title   string `json:"title,omitempty" elastic_mapping:"title:{type:text,copy_to:combined_fulltext,fields:{keyword: {type: keyword}, pinyin: {type: text, analyzer: pinyin_analyzer}}}"` // Document title
 	Summary string `json:"summary,omitempty" elastic_mapping:"summary:{type:text,copy_to:combined_fulltext}"`                                                                                // Brief summary or description of the document
 
-	Lang    string `json:"lang,omitempty" elastic_mapping:"lang:{type:keyword,copy_to:combined_fulltext}"`    // Language code (e.g., "en", "fr")
-	Content string `json:"content,omitempty" elastic_mapping:"content:{type:text,copy_to:combined_fulltext}"` // Document content for full-text indexing
+	Lang    string          `json:"lang,omitempty" elastic_mapping:"lang:{type:keyword,copy_to:combined_fulltext}"`    // Language code (e.g., "en", "fr")
+	Content string          `json:"content,omitempty" elastic_mapping:"content:{type:text,copy_to:combined_fulltext}"` // Document content for full-text indexing
+	Chunks  []DocumentChunk `json:"document_chunk,omitempty" elastic_mapping:"document_chunk:{type:nested}"`
 
 	Icon      string `json:"icon,omitempty" elastic_mapping:"icon:{enabled:false}"`           // Icon Key, need work with datasource's assets to get the icon url, if it is a full url, then use it directly
 	Thumbnail string `json:"thumbnail,omitempty" elastic_mapping:"thumbnail:{enabled:false}"` // Thumbnail image URL, for preview purposes
@@ -113,3 +115,78 @@ type UserInfo struct {
 	UserName   string `json:"username,omitempty" elastic_mapping:"username:{type:keyword,copy_to:combined_fulltext}"` // Login of the user
 	UserID     string `json:"userid,omitempty" elastic_mapping:"userid:{type:keyword,copy_to:combined_fulltext}"`     // Unique identifier for the user
 }
+
+type DocumentChunk struct {
+	Range     ChunkRange `json:"range" elastic_mapping:"range:{type:object}"`
+	Text      string     `json:"text" elastic_mapping:"text:{type:text}"`
+	Embedding Embedding  `json:"embedding" elastic_mapping:"embedding:{type:object}"`
+}
+
+// A `Embedding` stores a chunk's embedding.
+//
+// Only 1 field will be used, depending on the chosen embedding dimension, see
+// the `Dimension` field above.
+//
+// Having so many `EmbeddingXxx` fields is embarrasing, but we have no choice
+// since vector dimension is part of the type information and elastic mapping
+// has to be static.
+//
+// If you add or remove fields, please update variable "SupportedEmbeddingDimensions"
+// as well.
+type Embedding struct {
+	Embedding128  []float32 `json:"embedding128,omitempty" elastic_mapping:"embedding128:{type:knn_dense_float_vector,knn:{dims:128,model:lsh,similarity:cosine,L:99,k:1}}"`
+	Embedding256  []float32 `json:"embedding256,omitempty" elastic_mapping:"embedding256:{type:knn_dense_float_vector,knn:{dims:256,model:lsh,similarity:cosine,L:99,k:1}}"`
+	Embedding384  []float32 `json:"embedding384,omitempty" elastic_mapping:"embedding384:{type:knn_dense_float_vector,knn:{dims:384,model:lsh,similarity:cosine,L:99,k:1}}"`
+	Embedding512  []float32 `json:"embedding512,omitempty" elastic_mapping:"embedding512:{type:knn_dense_float_vector,knn:{dims:512,model:lsh,similarity:cosine,L:99,k:1}}"`
+	Embedding768  []float32 `json:"embedding768,omitempty" elastic_mapping:"embedding768:{type:knn_dense_float_vector,knn:{dims:768,model:lsh,similarity:cosine,L:99,k:1}}"`
+	Embedding1024 []float32 `json:"embedding1024,omitempty" elastic_mapping:"embedding1024:{type:knn_dense_float_vector,knn:{dims:1024,model:lsh,similarity:cosine,L:99,k:1}}"`
+	Embedding1536 []float32 `json:"embedding1536,omitempty" elastic_mapping:"embedding1536:{type:knn_dense_float_vector,knn:{dims:1536,model:lsh,similarity:cosine,L:99,k:1}}"`
+	Embedding2048 []float32 `json:"embedding2048,omitempty" elastic_mapping:"embedding2048:{type:knn_dense_float_vector,knn:{dims:2048,model:lsh,similarity:cosine,L:99,k:1}}"`
+	Embedding2560 []float32 `json:"embedding2560,omitempty" elastic_mapping:"embedding2560:{type:knn_dense_float_vector,knn:{dims:2560,model:lsh,similarity:cosine,L:99,k:1}}"`
+	Embedding4096 []float32 `json:"embedding4096,omitempty" elastic_mapping:"embedding4096:{type:knn_dense_float_vector,knn:{dims:4096,model:lsh,similarity:cosine,L:99,k:1}}"`
+}
+
+// Set the actual value of this "Embedding"
+func (e *Embedding) SetValue(embedding []float32) {
+	dimension := len(embedding)
+	switch dimension {
+	case 128:
+		e.Embedding128 = embedding
+	case 256:
+		e.Embedding256 = embedding
+	case 384:
+		e.Embedding384 = embedding
+	case 512:
+		e.Embedding512 = embedding
+	case 768:
+		e.Embedding768 = embedding
+	case 1024:
+		e.Embedding1024 = embedding
+	case 1536:
+		e.Embedding1536 = embedding
+	case 2048:
+		e.Embedding2048 = embedding
+	case 2560:
+		e.Embedding2560 = embedding
+	case 4096:
+		e.Embedding4096 = embedding
+	default:
+		panic(fmt.Sprintf("embedding's dimension is invalid, we accept %v", SupportedEmbeddingDimensions))
+	}
+}
+
+// Embedding dimensions supported by us, it should be kept sync with the
+// "EmbeddingXxx" fields of struct Embedding
+var SupportedEmbeddingDimensions = []int32{128, 256, 384, 512, 768, 1024, 1536, 2048, 2560, 4096}
+
+// Range of a chunk.
+//
+// A chunk contains roughly the same amount of tokens, say 8192 tokens. And
+// thus, a chunk can span many pages if these pages are small, or it is only
+// part of a page if the page is big.
+type ChunkRange struct {
+	// Start page of this chunk.
+	Start int `json:"start" elastic_mapping:"start:{type:integer}"`
+	// End page of this chuhk. This is **inclusive**.
+	End int `json:"end" elastic_mapping:"end:{type:integer}"`
+}

modules/assistant/langchain/llm.go

@@ -5,15 +5,16 @@
 package langchain
 
 import (
+	"net/http"
+	"net/http/httputil"
+
 	log "github.com/cihub/seelog"
 	"github.com/tmc/langchaingo/llms"
 	"github.com/tmc/langchaingo/llms/ollama"
 	"github.com/tmc/langchaingo/llms/openai"
 	"infini.sh/coco/core"
 	"infini.sh/coco/modules/common"
 	"infini.sh/framework/core/global"
-	"net/http"
-	"net/http/httputil"
 )
 
 type LoggingRoundTripper struct {
@@ -61,12 +62,14 @@ func GetLLM(endpoint, apiType, model, token string, keepalive string) llms.Model
 				openai.WithToken(token),
 				openai.WithBaseURL(endpoint),
 				openai.WithModel(model),
+				openai.WithEmbeddingModel(model),
 			)
 		} else {
 			llm, err = openai.New(
 				openai.WithToken(token),
 				openai.WithBaseURL(endpoint),
 				openai.WithModel(model),
+				openai.WithEmbeddingModel(model),
 			)
 		}
 

modules/attachment/attachment.go

@@ -7,15 +7,16 @@ package attachment
 import (
 	"errors"
 	"fmt"
-	"infini.sh/coco/core"
-	"infini.sh/framework/core/elastic"
 	"io"
 	"mime/multipart"
 	"net/http"
 	"path/filepath"
 	"strings"
 	"time"
 
+	"infini.sh/coco/core"
+	"infini.sh/framework/core/elastic"
+
 	log "github.com/cihub/seelog"
 	httprouter "infini.sh/framework/core/api/router"
 	"infini.sh/framework/core/kv"
@@ -53,7 +54,7 @@ func (h APIHandler) uploadAttachment(w http.ResponseWriter, r *http.Request, ps
 			return
 		}
 		// Upload to S3
-		if fileID, err := uploadToBlobStore(ctx, file, fileHeader.Filename); err != nil {
+		if fileID, err := UploadToBlobStore(ctx, "", file, fileHeader.Filename, "", false); err != nil {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
 		} else {
@@ -241,7 +242,18 @@ func getMimeType(file multipart.File) (string, error) {
 	return mimeType, nil
 }
 
-func uploadToBlobStore(ctx *orm.Context, file multipart.File, fileName string) (string, error) {
+// Helper function to upload the attachment specified by [file] to the
+// blob store.
+//
+// Arguments:
+//
+//   - If [fileID] is not an empty string, it will be used as the file ID.
+//     Otherwise, a random ID will be created and used.
+//   - If [ownerID] is not empty, the created attached will set the owner to it.
+//     Otherwise, owner information will be extracted from cotnext [ctx].
+//   - [replaceIfExists]: If this is true and there is already an attachment with
+//     the same file ID eixsts, replace it.
+func UploadToBlobStore(ctx *orm.Context, fileID string, file multipart.File, fileName string, ownerID string, replaceIfExists bool) (string, error) {
 	defer func() {
 		_ = file.Close()
 	}()
@@ -252,7 +264,9 @@ func uploadToBlobStore(ctx *orm.Context, file multipart.File, fileName string) (
 		return "", fmt.Errorf("failed to read file %s: %v", fileName, err)
 	}
 
-	fileID := util.GetUUID()
+	if fileID == "" {
+		fileID = util.GetUUID()
+	}
 	fileSize := len(data)
 	mimeType, _ := getMimeType(file)
 
@@ -264,14 +278,24 @@ func uploadToBlobStore(ctx *orm.Context, file multipart.File, fileName string) (
 	attachment.Icon = getFileExtension(fileName)
 	attachment.URL = fmt.Sprintf("/attachment/%v", fileID)
 	//attachment.Owner //TODO
+	if ownerID != "" {
+		attachment.SetOwnerID(ownerID)
+	}
 
 	//save attachment metadata
-	err = orm.Create(ctx, &attachment)
+	if replaceIfExists {
+		err = orm.Upsert(ctx, &attachment)
+	} else {
+		err = orm.Create(ctx, &attachment)
+	}
 	if err != nil {
 		panic(err)
 	}
 
 	//save attachment payload
+	//
+	// kv.AddValue will replace the previous value if it already exists so we
+	// don't need to check [replaceIfExists] here.
 	err = kv.AddValue(AttachmentKVBucket, []byte(fileID), data)
 	if err != nil {
 		panic(err)

plugins/connectors/local_fs/plugin.go

@@ -154,9 +154,10 @@ func (p *Plugin) saveDocument(ctx *pipeline.Context, currentPath, basePath strin
 		Source:   core.DataSourceReference{ID: datasource.ID, Type: "connector", Name: datasource.Name},
 		Type:     connectors.TypeFile,
 		Category: filepath.Dir(currentPath),
-		Content:  "", // skip content
-		URL:      currentPath,
-		Size:     int(fileInfo.Size()),
+		// skip content here, which will be popluated by the `read_file_content` processor
+		Content: "",
+		URL:     currentPath,
+		Size:    int(fileInfo.Size()),
 	}
 	doc.System = datasource.System
 	if doc.System == nil {