diff --git a/libs/langchain-community/src/experimental/llms/chrome_ai.ts b/libs/langchain-community/src/experimental/llms/chrome_ai.ts index 618a8b6ecfbc..c41ba13121a6 100644 --- a/libs/langchain-community/src/experimental/llms/chrome_ai.ts +++ b/libs/langchain-community/src/experimental/llms/chrome_ai.ts @@ -5,25 +5,94 @@ import { GenerationChunk } from "@langchain/core/outputs"; import { IterableReadableStream } from "@langchain/core/utils/stream"; import { BaseLLMParams, LLM } from "@langchain/core/language_models/llms"; -export interface AI { - canCreateTextSession(): Promise; - createTextSession(options?: AITextSessionOptions): Promise; - defaultTextSessionOptions(): Promise; +export interface AILanguageModelFactory { + create(options?: AILanguageModelCreateOptions): Promise; + capabilities(): Promise; } -export interface AITextSession { - prompt(input: string): Promise; - promptStreaming(input: string): ReadableStream; +export interface AILanguageModel extends EventTarget { + prompt( + input: AILanguageModelPromptInput, + options?: AILanguageModelPromptOptions + ): Promise; + promptStreaming( + input: AILanguageModelPromptInput, + options?: AILanguageModelPromptOptions + ): ReadableStream; + + countPromptTokens( + input: AILanguageModelPromptInput, + options?: AILanguageModelPromptOptions + ): Promise; + + get maxTokens(): number; + get tokensSoFar(): number; + get tokensLeft(): number; + + get topK(): number; + get temperature(): number; + + oncontextoverflow: (event: Event) => void; + + clone(options?: AILanguageModelCloneOptions): Promise; destroy(): void; - clone(): AITextSession; } -export interface AITextSessionOptions { +interface AILanguageModelCapabilities { + readonly available: AICapabilityAvailability; + languageAvailable(languageTag: string): AICapabilityAvailability; + + get defaultTopK(): number | undefined; + get maxTopK(): number | undefined; + get defaultTemperature(): number | undefined; + get maxTemperature(): number | undefined; +} + +interface AILanguageModelCreateOptions { + signal?: AbortSignal; + monitor?: AICreateMonitorCallback; + systemPrompt?: string; + initialPrompts?: AILanguageModelInitialPrompt[]; topK: number; temperature: number; } -export type AIModelAvailability = "readily" | "after-download" | "no"; +export interface AILanguageModelInitialPrompt { + role: AILanguageModelInitialPromptRole; + content: string; +} + +export interface AILanguageModelPrompt { + role: AILanguageModelPromptRole; + content: string; +} + +export interface AILanguageModelPromptOptions { + signal?: AbortSignal; +} + +export interface AILanguageModelCloneOptions { + signal?: AbortSignal; +} + +export type AILanguageModelPromptInput = + | string + | AILanguageModelPrompt + | AILanguageModelPrompt[]; + +enum AILanguageModelInitialPromptRole { + "system", + "user", + "assistant", +} + +enum AILanguageModelPromptRole { + "user", + "assistant", +} + +export type AICapabilityAvailability = "yes" | "no"; +export type AICreateMonitorCallback = () => void; export interface ChromeAIInputs extends BaseLLMParams { topK?: number; @@ -93,14 +162,14 @@ export class ChromeAI extends LLM { `Could not initialize ChromeAI instance. Make sure you are running a version of Chrome with the proper experimental flags enabled.\n\nError message: ${e.message}` ); } - const { available } = await aiInstance.assistant.capabilities(); + const { available } = await aiInstance.languageModel.capabilities(); if (available === "no") { throw new Error("The AI model is not available."); } else if (available === "after-download") { throw new Error("The AI model is not yet downloaded."); } - const session = await aiInstance.assistant.create({ + const session = await aiInstance.languageModel.create({ systemPrompt: this.systemPrompt, topK: this.topK, temperature: this.temperature,