Merge pull request #165 from FlowiseAI/feature/Cohere

Bugfix/Cohere API key not found
This commit is contained in:
Henry Heng
2023-05-24 13:40:01 +01:00
committed by GitHub
2 changed files with 79 additions and 1 deletions
@@ -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
@@ -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`')
}
}
}