Chore/cohere embed models (FlowiseAI#2035)

* Update of Cohere embedding models * Update of Cohere embedding models * adding official docs link to inputType
menloparklab · Mar 26, 2024 · f116dba · f116dba
1 parent 794818b
commit f116dba
Showing 1 changed file with 60 additions and 4 deletions.
diff --git a/packages/components/nodes/embeddings/CohereEmbedding/CohereEmbedding.ts b/packages/components/nodes/embeddings/CohereEmbedding/CohereEmbedding.ts
@@ -17,7 +17,7 @@ class CohereEmbedding_Embeddings implements INode {
     constructor() {
         this.label = 'Cohere Embeddings'
         this.name = 'cohereEmbeddings'
-        this.version = 1.0
+        this.version = 2.0
         this.type = 'CohereEmbeddings'
         this.icon = 'Cohere.svg'
         this.category = 'Embeddings'
@@ -35,27 +35,82 @@ class CohereEmbedding_Embeddings implements INode {
                 name: 'modelName',
                 type: 'options',
                 options: [
+                    {
+                        label: 'embed-english-v3.0',
+                        name: 'embed-english-v3.0',
+                        description: 'Embedding Dimensions: 1024'
+                    },
+                    {
+                        label: 'embed-english-light-v3.0',
+                        name: 'embed-english-light-v3.0',
+                        description: 'Embedding Dimensions: 384'
+                    },
+                    {
+                        label: 'embed-multilingual-v3.0',
+                        name: 'embed-multilingual-v3.0',
+                        description: 'Embedding Dimensions: 1024'
+                    },
+                    {
+                        label: 'embed-multilingual-light-v3.0',
+                        name: 'embed-multilingual-light-v3.0',
+                        description: 'Embedding Dimensions: 384'
+                    },
                     {
                         label: 'embed-english-v2.0',
-                        name: 'embed-english-v2.0'
+                        name: 'embed-english-v2.0',
+                        description: 'Embedding Dimensions: 4096'
                     },
                     {
                         label: 'embed-english-light-v2.0',
-                        name: 'embed-english-light-v2.0'
+                        name: 'embed-english-light-v2.0',
+                        description: 'Embedding Dimensions: 1024'
                     },
                     {
                         label: 'embed-multilingual-v2.0',
-                        name: 'embed-multilingual-v2.0'
+                        name: 'embed-multilingual-v2.0',
+                        description: 'Embedding Dimensions: 768'
                     }
                 ],
                 default: 'embed-english-v2.0',
                 optional: true
+            },
+            {
+                label: 'Type',
+                name: 'inputType',
+                type: 'options',
+                description:
+                    'Specifies the type of input passed to the model. Required for embedding models v3 and higher. <a target="_blank" href="https://docs.cohere.com/reference/embed">Official Docs</a>',
+                options: [
+                    {
+                        label: 'search_document',
+                        name: 'search_document',
+                        description: 'Use this to encode documents for embeddings that you store in a vector database for search use-cases'
+                    },
+                    {
+                        label: 'search_query',
+                        name: 'search_query',
+                        description: 'Use this when you query your vector DB to find relevant documents.'
+                    },
+                    {
+                        label: 'classification',
+                        name: 'classification',
+                        description: 'Use this when you use the embeddings as an input to a text classifier'
+                    },
+                    {
+                        label: 'clustering',
+                        name: 'clustering',
+                        description: 'Use this when you want to cluster the embeddings.'
+                    }
+                ],
+                default: 'search_query',
+                optional: true
             }
         ]
     }
 
     async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
         const modelName = nodeData.inputs?.modelName as string
+        const inputType = nodeData.inputs?.inputType as string
 
         const credentialData = await getCredentialData(nodeData.credential ?? '', options)
         const cohereApiKey = getCredentialParam('cohereApiKey', credentialData, nodeData)
@@ -65,6 +120,7 @@ class CohereEmbedding_Embeddings implements INode {
         }
 
         if (modelName) obj.model = modelName
+        if (inputType) obj.inputType = inputType
 
         const model = new CohereEmbeddings(obj)
         return model