fix: Improve ChromeAI model parameters

This commit is contained in:
n4ze3m 2024-08-31 20:17:10 +05:30
parent 26e799b246
commit 9bd3193ad9
2 changed files with 10 additions and 14 deletions

View File

@ -11,11 +11,6 @@ import { BaseMessage, AIMessageChunk } from "@langchain/core/messages"
import { ChatGenerationChunk } from "@langchain/core/outputs" import { ChatGenerationChunk } from "@langchain/core/outputs"
import { IterableReadableStream } from "@langchain/core/utils/stream" import { IterableReadableStream } from "@langchain/core/utils/stream"
export interface AI {
canCreateTextSession(): Promise<AIModelAvailability>
createTextSession(options?: AITextSessionOptions): Promise<AITextSession>
defaultTextSessionOptions(): Promise<AITextSessionOptions>
}
export interface AITextSession { export interface AITextSession {
prompt(input: string): Promise<string> prompt(input: string): Promise<string>
@ -44,7 +39,7 @@ export interface ChromeAIInputs extends BaseChatModelParams {
promptFormatter?: (messages: BaseMessage[]) => string promptFormatter?: (messages: BaseMessage[]) => string
} }
export interface ChromeAICallOptions extends BaseLanguageModelCallOptions {} export interface ChromeAICallOptions extends BaseLanguageModelCallOptions { }
function formatPrompt(messages: BaseMessage[]): string { function formatPrompt(messages: BaseMessage[]): string {
return messages return messages
@ -88,9 +83,9 @@ function formatPrompt(messages: BaseMessage[]): string {
export class ChatChromeAI extends SimpleChatModel<ChromeAICallOptions> { export class ChatChromeAI extends SimpleChatModel<ChromeAICallOptions> {
session?: AITextSession session?: AITextSession
temperature = 0.5 temperature = 0.8
topK = 40 topK = 120
promptFormatter: (messages: BaseMessage[]) => string promptFormatter: (messages: BaseMessage[]) => string
@ -121,16 +116,17 @@ export class ChatChromeAI extends SimpleChatModel<ChromeAICallOptions> {
} }
const { ai } = window as any const { ai } = window as any
const canCreateTextSession = await ai.canCreateTextSession() const capabilities = await ai?.assistant?.capabilities()
const canCreateTextSession = capabilities?.available
if (canCreateTextSession === AIModelAvailability.No) { if (canCreateTextSession === AIModelAvailability.No) {
throw new Error("The AI model is not available.") throw new Error("The AI model is not available.")
} else if (canCreateTextSession === AIModelAvailability.AfterDownload) { } else if (canCreateTextSession === AIModelAvailability.AfterDownload) {
throw new Error("The AI model is not yet downloaded.") throw new Error("The AI model is not yet downloaded.")
} }
this.session = await ai.createTextSession({ this.session = await ai?.assistant?.create({
topK: this.topK, temperature: this.temperature,
temperature: this.temperature topK: this.topK
}) })
} }

View File

@ -12,8 +12,8 @@ export const getChromeAISupported = async () => {
} }
//@ts-ignore //@ts-ignore
const createSession = await ai?.canCreateGenericSession() const capabilities = await ai?.assistant?.capabilities()
if (createSession !== "readily") { if (capabilities?.available !== "readily") {
return "ai_not_ready" return "ai_not_ready"
} }