import { Assistant, AssistantExtension, ToolManager } from '@janhq/core'
import { RetrievalTool } from './tools/retrieval'

export default class JanAssistantExtension extends AssistantExtension {

  async onLoad() {
    // Register the retrieval tool
    ToolManager.instance().register(new RetrievalTool())
  }

  /**
   * Called when the extension is unloaded.
   */
  onUnload(): void {}

  async getAssistants(): Promise<Assistant[]> {
    return [this.defaultAssistant]
  }

  /** DEPRECATED */
  async createAssistant(assistant: Assistant): Promise<void> {}
  async deleteAssistant(assistant: Assistant): Promise<void> {}

  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,
  }
}