2024-12-19 11:38:10 +07:00

158 lines
4.0 KiB
TypeScript

import {
fs,
Assistant,
events,
joinPath,
AssistantExtension,
AssistantEvent,
ToolManager,
} from '@janhq/core'
import { RetrievalTool } from './tools/retrieval'
export default class JanAssistantExtension extends AssistantExtension {
private static readonly _homeDir = 'file://assistants'
async onLoad() {
// Register the retrieval tool
ToolManager.instance().register(new RetrievalTool())
// making the assistant directory
const assistantDirExist = await fs.existsSync(
JanAssistantExtension._homeDir
)
if (
localStorage.getItem(`${this.name}-version`) !== VERSION ||
!assistantDirExist
) {
if (!assistantDirExist) await fs.mkdir(JanAssistantExtension._homeDir)
// Write assistant metadata
await this.createJanAssistant()
// Finished migration
localStorage.setItem(`${this.name}-version`, VERSION)
// Update the assistant list
events.emit(AssistantEvent.OnAssistantsUpdate, {})
}
}
/**
* Called when the extension is unloaded.
*/
onUnload(): void {}
async createAssistant(assistant: Assistant): Promise<void> {
const assistantDir = await joinPath([
JanAssistantExtension._homeDir,
assistant.id,
])
if (!(await fs.existsSync(assistantDir))) await fs.mkdir(assistantDir)
// store the assistant metadata json
const assistantMetadataPath = await joinPath([
assistantDir,
'assistant.json',
])
try {
await fs.writeFileSync(
assistantMetadataPath,
JSON.stringify(assistant, null, 2)
)
} catch (err) {
console.error(err)
}
}
async getAssistants(): Promise<Assistant[]> {
try {
// get all the assistant directories
// get all the assistant metadata json
const results: Assistant[] = []
const allFileName: string[] = await fs.readdirSync(
JanAssistantExtension._homeDir
)
for (const fileName of allFileName) {
const filePath = await joinPath([
JanAssistantExtension._homeDir,
fileName,
])
if (!(await fs.fileStat(filePath))?.isDirectory) continue
const jsonFiles: string[] = (await fs.readdirSync(filePath)).filter(
(file: string) => file === 'assistant.json'
)
if (jsonFiles.length !== 1) {
// has more than one assistant file -> ignore
continue
}
const content = await fs.readFileSync(
await joinPath([filePath, jsonFiles[0]]),
'utf-8'
)
const assistant: Assistant =
typeof content === 'object' ? content : JSON.parse(content)
results.push(assistant)
}
return results
} catch (err) {
console.debug(err)
return [this.defaultAssistant]
}
}
async deleteAssistant(assistant: Assistant): Promise<void> {
if (assistant.id === 'jan') {
return Promise.reject('Cannot delete Jan Assistant')
}
// remove the directory
const assistantDir = await joinPath([
JanAssistantExtension._homeDir,
assistant.id,
])
return fs.rm(assistantDir)
}
private async createJanAssistant(): Promise<void> {
await this.createAssistant(this.defaultAssistant)
}
private defaultAssistant: Assistant = {
avatar: '',
thread_location: undefined,
id: 'jan',
object: 'assistant',
created_at: Date.now() / 1000,
name: 'Jan',
description: 'A default assistant that can use all downloaded models',
model: '*',
instructions: '',
tools: [
{
type: 'retrieval',
enabled: false,
useTimeWeightedRetriever: false,
settings: {
top_k: 2,
chunk_size: 1024,
chunk_overlap: 64,
retrieval_template: `Use the following pieces of context to answer the question at the end.
----------------
CONTEXT: {CONTEXT}
----------------
QUESTION: {QUESTION}
----------------
Helpful Answer:`,
},
},
],
file_ids: [],
metadata: undefined,
}
}