import { LLM, type BaseLLMParams } from 'langchain/llms'
import Replicate from 'replicate'

export interface VicunaInput {
  apiKey?: string
}

export class Vicuna extends LLM implements VicunaInput {
  model: `${string}/${string}:${string}` =
    'replicate/vicuna-13b:a68b84083b703ab3d5fbf31b6e25f16be2988e4c3e21fe79c2ff1c18b99e61c1'

  input = { max_length: 4000, temperature: 0.2 }

  apiKey: string

  constructor(fields: VicunaInput & BaseLLMParams) {
    super(fields)

    const apiKey =
      fields?.apiKey ??
      // eslint-disable-next-line no-process-env
      (typeof process !== 'undefined' && process.env?.REPLICATE_API_KEY)

    if (!apiKey) {
      throw new Error('Please set the REPLICATE_API_KEY environment variable')
    }

    this.apiKey = apiKey
  }

  _llmType() {
    return 'vicuna'
  }

  async _call(prompt: string, _stop?: string[]): Promise<string> {
    _stop // unused
    const replicate = new Replicate({
      userAgent: 'langchain',
      auth: this.apiKey
    })

    const output = (await this.caller.call(() =>
      replicate.run(this.model, {
        wait: true,
        input: {
          ...this.input,
          prompt
        }
      })
    )) as string[]

    output[0] = `${output[0]} `
    return output.join('')
  }
}
Vicuna LLM
