From 0ed84312523d0c4d95a41b93f231d57911138664 Mon Sep 17 00:00:00 2001
From: Henry <hzj94@hotmail.com>
Date: Wed, 24 May 2023 13:27:38 +0100
Subject: [PATCH] fix Cohere API key not found

---
 .../components/nodes/llms/Cohere/Cohere.ts    |  2 +-
 packages/components/nodes/llms/Cohere/core.ts | 78 +++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 packages/components/nodes/llms/Cohere/core.ts

diff --git a/packages/components/nodes/llms/Cohere/Cohere.ts b/packages/components/nodes/llms/Cohere/Cohere.ts
index dc632ec31e8..a7e9c696ce2 100644
--- a/packages/components/nodes/llms/Cohere/Cohere.ts
+++ b/packages/components/nodes/llms/Cohere/Cohere.ts
@@ -1,6 +1,6 @@
 import { INode, INodeData, INodeParams } from '../../../src/Interface'
 import { getBaseClasses } from '../../../src/utils'
-import { Cohere, CohereInput } from 'langchain/llms/cohere'
+import { Cohere, CohereInput } from './core'
 
 class Cohere_LLMs implements INode {
     label: string
diff --git a/packages/components/nodes/llms/Cohere/core.ts b/packages/components/nodes/llms/Cohere/core.ts
new file mode 100644
index 00000000000..97c81571028
--- /dev/null
+++ b/packages/components/nodes/llms/Cohere/core.ts
@@ -0,0 +1,78 @@
+import { LLM, BaseLLMParams } from 'langchain/llms/base'
+
+export interface CohereInput extends BaseLLMParams {
+    /** Sampling temperature to use */
+    temperature?: number
+
+    /**
+     * Maximum number of tokens to generate in the completion.
+     */
+    maxTokens?: number
+
+    /** Model to use */
+    model?: string
+
+    apiKey?: string
+}
+
+export class Cohere extends LLM implements CohereInput {
+    temperature = 0
+
+    maxTokens = 250
+
+    model: string
+
+    apiKey: string
+
+    constructor(fields?: CohereInput) {
+        super(fields ?? {})
+
+        const apiKey = fields?.apiKey ?? undefined
+
+        if (!apiKey) {
+            throw new Error('Please set the COHERE_API_KEY environment variable or pass it to the constructor as the apiKey field.')
+        }
+
+        this.apiKey = apiKey
+        this.maxTokens = fields?.maxTokens ?? this.maxTokens
+        this.temperature = fields?.temperature ?? this.temperature
+        this.model = fields?.model ?? this.model
+    }
+
+    _llmType() {
+        return 'cohere'
+    }
+
+    /** @ignore */
+    async _call(prompt: string, options: this['ParsedCallOptions']): Promise<string> {
+        const { cohere } = await Cohere.imports()
+
+        cohere.init(this.apiKey)
+
+        // Hit the `generate` endpoint on the `large` model
+        const generateResponse = await this.caller.callWithOptions({ signal: options.signal }, cohere.generate.bind(cohere), {
+            prompt,
+            model: this.model,
+            max_tokens: this.maxTokens,
+            temperature: this.temperature,
+            end_sequences: options.stop
+        })
+        try {
+            return generateResponse.body.generations[0].text
+        } catch {
+            throw new Error('Could not parse response.')
+        }
+    }
+
+    /** @ignore */
+    static async imports(): Promise<{
+        cohere: typeof import('cohere-ai')
+    }> {
+        try {
+            const { default: cohere } = await import('cohere-ai')
+            return { cohere }
+        } catch (e) {
+            throw new Error('Please install cohere-ai as a dependency with, e.g. `yarn add cohere-ai`')
+        }
+    }
+}