Skip to content

Commit 35a32dc

Browse files
eric2788eric2788
committed
added web-llm to llm provider
1 parent 48d015a commit 35a32dc

14 files changed

+311
-39
lines changed

‎package.json

+2
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
"@ffmpeg/ffmpeg": "^0.12.10",
2828
"@ffmpeg/util": "^0.12.1",
2929
"@material-tailwind/react": "^2.1.9",
30+
"@mlc-ai/web-llm": "^0.2.73",
3031
"@plasmohq/messaging": "^0.6.2",
3132
"@plasmohq/storage": "^1.9.3",
3233
"@react-hooks-library/core": "^0.5.2",
@@ -64,6 +65,7 @@
6465
"@types/react": "18.2.37",
6566
"@types/react-dom": "18.2.15",
6667
"@types/semver": "^7.5.8",
68+
"@webgpu/types": "^0.1.49",
6769
"dotenv": "^16.4.5",
6870
"esbuild": "^0.20.2",
6971
"gify-parse": "^1.0.7",

‎pnpm-lock.yaml

+24
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎src/contents/index/mounter.tsx

+1-1
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,7 @@ function createApp(roomId: string, plasmo: PlasmoSpec, info: StreamInfo): App {
173173
onClick: () => this.start()
174174
},
175175
position: 'top-left',
176-
duration: 6000000,
176+
duration: Infinity,
177177
dismissible: false
178178
})
179179
return

‎src/hooks/life-cycle.ts

+30
Original file line numberDiff line numberDiff line change
@@ -74,4 +74,34 @@ export function useTimeoutElement(before: JSX.Element, after: JSX.Element, timeo
7474
}, [after, before, timeout]);
7575

7676
return element;
77+
}
78+
79+
/**
80+
* A React hook that ensures the provided disposable resource is properly disposed of
81+
* when the component is unmounted.
82+
*
83+
* @param disposable - The resource to be disposed of, which can be either a `Disposable`
84+
* or an `AsyncDisposable`. The resource must implement either the `Symbol.dispose` or
85+
* `Symbol.asyncDispose` method.
86+
*
87+
* @example
88+
* ```typescript
89+
* const MyComponent = () => {
90+
* const disposableResource = useMemo(() => createDisposableResource(), []);
91+
* withDisposable(disposableResource);
92+
*
93+
* return <div>My Component</div>;
94+
* };
95+
* ```
96+
*/
97+
export function useDisposable(disposable: Disposable | AsyncDisposable) {
98+
useEffect(() => {
99+
return () => {
100+
if (!!disposable[Symbol.dispose]) {
101+
disposable[Symbol.dispose]();
102+
} else if (!!disposable[Symbol.asyncDispose]) {
103+
disposable[Symbol.asyncDispose]();
104+
}
105+
}
106+
}, [])
77107
}

‎src/llms/cloudflare-ai.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ export default class CloudFlareAI implements LLMProviders {
4141
console.warn('Cloudflare AI session is not supported')
4242
return {
4343
...this,
44-
[Symbol.dispose]: () => { }
44+
[Symbol.asyncDispose]: async () => { }
4545
}
4646
}
4747

‎src/llms/gemini-nano.ts

+4-4
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ export default class GeminiNano implements LLMProviders {
1919
return session.prompt(chat)
2020
} finally {
2121
console.debug('[gemini nano] done')
22-
session[Symbol.dispose]()
22+
await session[Symbol.asyncDispose]()
2323
}
2424
}
2525

@@ -33,7 +33,7 @@ export default class GeminiNano implements LLMProviders {
3333
}
3434
} finally {
3535
console.debug('[gemini nano] done')
36-
session[Symbol.dispose]()
36+
await session[Symbol.asyncDispose]()
3737
}
3838
}
3939

@@ -82,7 +82,7 @@ class GeminiAssistant implements Session<LLMProviders> {
8282
}
8383
}
8484

85-
[Symbol.dispose](): void {
85+
async [Symbol.asyncDispose]() {
8686
this.assistant.destroy()
8787
}
8888
}
@@ -105,7 +105,7 @@ class GeminiSummarizer implements Session<LLMProviders> {
105105
}
106106
}
107107

108-
[Symbol.dispose](): void {
108+
async [Symbol.asyncDispose]() {
109109
this.summarizer.destroy()
110110
}
111111

‎src/llms/index.ts

+5-3
Original file line numberDiff line numberDiff line change
@@ -3,21 +3,23 @@ import type { SettingSchema as LLMSchema } from '~options/fragments/llm'
33
import cloudflare from './cloudflare-ai'
44
import nano from './gemini-nano'
55
import worker from './remote-worker'
6+
import webllm from './web-llm'
67

78
export interface LLMProviders {
89
cumulative: boolean
9-
validate(): Promise<void>
10+
validate(progress?: (p: number, t: string) => void): Promise<void>
1011
prompt(chat: string): Promise<string>
1112
promptStream(chat: string): AsyncGenerator<string>
1213
asSession(): Promise<Session<LLMProviders>>
1314
}
1415

15-
export type Session<T> = Disposable & Omit<T, 'asSession' | 'validate' | 'cumulative'>
16+
export type Session<T> = AsyncDisposable & Omit<T, 'asSession' | 'validate' | 'cumulative'>
1617

1718
const llms = {
1819
cloudflare,
1920
nano,
20-
worker
21+
worker,
22+
webllm
2123
}
2224

2325
export type LLMs = typeof llms

‎src/llms/models.ts

+33
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
import type { LLMTypes } from "~llms"
2+
3+
export type ModelList = {
4+
providers: LLMTypes[]
5+
models: string[]
6+
}
7+
8+
const models: ModelList[] = [
9+
{
10+
providers: ['worker', 'cloudflare'],
11+
models: [
12+
'@cf/qwen/qwen1.5-14b-chat-awq',
13+
'@cf/qwen/qwen1.5-7b-chat-awq',
14+
'@cf/qwen/qwen1.5-1.8b-chat',
15+
'@hf/google/gemma-7b-it',
16+
'@hf/nousresearch/hermes-2-pro-mistral-7b'
17+
]
18+
},
19+
{
20+
providers: [ 'webllm' ],
21+
models: [
22+
'Qwen2-7B-Instruct-q4f32_1-MLC',
23+
'Qwen2.5-14B-Instruct-q4f16_1-MLC',
24+
'gemma-2-9b-it-q4f16_1-MLC',
25+
'Qwen2.5-3B-Instruct-q0f16-MLC',
26+
'Phi-3-mini-128k-instruct-q0f16-MLC',
27+
'Phi-3.5-mini-instruct-q4f16_1-MLC-1k'
28+
]
29+
}
30+
]
31+
32+
33+
export default models

‎src/llms/remote-worker.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ export default class RemoteWorker implements LLMProviders {
5454
console.warn('Remote worker session is not supported')
5555
return {
5656
...this,
57-
[Symbol.dispose]: () => { }
57+
[Symbol.asyncDispose]: async () => { }
5858
}
5959
}
6060

‎src/llms/web-llm.ts

+87
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
import type { MLCEngine } from "@mlc-ai/web-llm";
2+
import type { LLMProviders, Session } from "~llms";
3+
import type { SettingSchema } from "~options/fragments/llm";
4+
5+
export default class WebLLM implements LLMProviders {
6+
7+
private static readonly DEFAULT_MODEL: string = 'Qwen2-7B-Instruct-q4f32_1-MLC'
8+
9+
private readonly model: string
10+
11+
constructor(settings: SettingSchema) {
12+
this.model = settings.model || WebLLM.DEFAULT_MODEL
13+
}
14+
15+
cumulative: boolean = true
16+
17+
async validate(progresser?: (p: number, t: string) => void): Promise<void> {
18+
await this.initializeEngine(progresser)
19+
}
20+
21+
async prompt(chat: string): Promise<string> {
22+
const session = await this.asSession()
23+
try {
24+
console.debug('[web-llm] prompting: ', chat)
25+
return session.prompt(chat)
26+
} finally {
27+
console.debug('[web-llm] done')
28+
await session[Symbol.asyncDispose]()
29+
}
30+
}
31+
32+
async *promptStream(chat: string): AsyncGenerator<string> {
33+
const session = await this.asSession()
34+
try {
35+
console.debug('[web-llm] prompting stream: ', chat)
36+
const res = session.promptStream(chat)
37+
for await (const chunk of res) {
38+
yield chunk
39+
}
40+
} finally {
41+
console.debug('[web-llm] done')
42+
await session[Symbol.asyncDispose]()
43+
}
44+
}
45+
46+
async asSession(): Promise<Session<LLMProviders>> {
47+
const engine = await this.initializeEngine()
48+
return {
49+
async prompt(chat: string) {
50+
await engine.interruptGenerate()
51+
const c = await engine.completions.create({
52+
prompt: chat,
53+
max_tokens: 512,
54+
temperature: 0.2,
55+
})
56+
return c.choices[0]?.text ?? engine.getMessage()
57+
},
58+
async *promptStream(chat: string): AsyncGenerator<string> {
59+
await engine.interruptGenerate()
60+
const chunks = await engine.completions.create({
61+
prompt: chat,
62+
max_tokens: 512,
63+
temperature: 0.2,
64+
stream: true
65+
})
66+
for await (const chunk of chunks) {
67+
yield chunk.choices[0]?.text || "";
68+
if (chunk.usage) {
69+
console.debug('Usage:', chunk.usage)
70+
}
71+
}
72+
},
73+
[Symbol.asyncDispose]: engine.unload
74+
}
75+
}
76+
77+
private async initializeEngine(progresser?: (p: number, t: string) => void): Promise<MLCEngine> {
78+
const { CreateMLCEngine } = await import('@mlc-ai/web-llm')
79+
return CreateMLCEngine(this.model, {
80+
initProgressCallback: (progress) => {
81+
progresser?.(progress.progress, "正在下载AI模型到本地")
82+
console.log('初始化进度:', progress)
83+
}
84+
})
85+
}
86+
87+
}

‎src/options/components/Selector.tsx

+2-1
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ export type SelectorProps<T> = {
1515
disabled?: boolean
1616
options: SelectorOption<T>[]
1717
className?: string
18+
emptyValue?: string
1819
}
1920

2021

@@ -37,7 +38,7 @@ function Selector<T = any>(props: SelectorProps<T>): JSX.Element {
3738
<label className="text-sm ml-1 font-medium text-gray-900 dark:text-white">{props.label}</label>
3839
<div ref={dropdownRef} className={`mt-2 ${props.disabled ? 'cursor-not-allowed' : 'cursor-pointer'}`} onClick={() => !props.disabled && setOpen(!isOpen)}>
3940
<div className={`inline-flex justify-between h-full w-full rounded-md border border-gray-300 dark:border-gray-600 shadow-sm px-4 py-2 text-sm font-medium text-gray-700 dark:text-white ${props.disabled ? 'opacity-50 bg-transparent' : 'bg-white dark:bg-gray-800 hover:bg-gray-50 dark:hover:bg-gray-900'} focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-offset-gray-100 dark:focus:ring-gray-500`}>
40-
{props.options.find((option) => option.value === props.value)?.label ?? String(props.value)}
41+
{props.options.find((option) => option.value === props.value)?.label || String(props.value || (props.emptyValue || '请选择'))}
4142
<svg className="-mr-1 ml-2 h-5 w-5" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 20 20" fill="currentColor" aria-hidden="true">
4243
<path fillRule="evenodd" d="M5.293 7.293a1 1 0 011.414 0L10 10.586l3.293-3.293a1 1 0 111.414 1.414l-4 4a1 1 0 01-1.414 0l-4-4a1 1 0 010-1.414z" clipRule="evenodd" />
4344
</svg>

0 commit comments

Comments
 (0)