refactor: introduce inference tools (#2493)

core/src/browser/extensions/engines/AIEngine.ts

@@ -36,7 +36,7 @@ export abstract class AIEngine extends BaseExtension {
  * Registers AI Engines
  */
  registerEngine() {
- EngineManager.instance()?.register(this)
+ EngineManager.instance().register(this)
  }
 
  /**

core/src/browser/extensions/engines/EngineManager.ts

@@ -23,7 +23,10 @@ export class EngineManager {
  return this.engines.get(provider) as T | undefined
  }
 
- static instance(): EngineManager | undefined {
- return window.core?.engineManager as EngineManager
+ /**
+ * The instance of the engine manager.
+ */
+ static instance(): EngineManager {
+ return window.core?.engineManager as EngineManager ?? new EngineManager()
  }
 }

core/src/browser/extensions/engines/OAIEngine.ts

@@ -106,6 +106,7 @@ export abstract class OAIEngine extends AIEngine {
  return
  }
  message.status = MessageStatus.Error
+ message.error_code = err.code
  events.emit(MessageEvent.OnMessageUpdate, message)
  },
  })

core/src/browser/extensions/engines/helpers/sse.ts

@@ -1,5 +1,5 @@
 import { Observable } from 'rxjs'
-import { ModelRuntimeParams } from '../../../../types'
+import { ErrorCode, ModelRuntimeParams } from '../../../../types'
 /**
  * Sends a request to the inference server to generate a response based on the recent messages.
  * @param recentMessages - An array of recent messages to use as context for the inference.
@@ -34,6 +34,16 @@ export function requestInference(
  signal: controller?.signal,
  })
  .then(async (response) => {
+ if (!response.ok) {
+ const data = await response.json()
+ const error = {
+ message: data.error?.message ?? 'Error occurred.',
+ code: data.error?.code ?? ErrorCode.Unknown,
+ }
+ subscriber.error(error)
+ subscriber.complete()
+ return
+ }
  if (model.parameters.stream === false) {
  const data = await response.json()
  subscriber.next(data.choices[0]?.message?.content ?? '')

core/src/browser/index.ts

@@ -27,3 +27,9 @@ export * from './extension'
  * @module
  */
 export * from './extensions'
+
+/**
+ * Export all base tools.
+ * @module
+ */
+export * from './tools'

core/src/browser/tools/index.ts

@@ -0,0 +1,2 @@
+export * from './manager'
+export * from './tool'

core/src/browser/tools/manager.ts

@@ -0,0 +1,47 @@
+import { AssistantTool, MessageRequest } from '../../types'
+import { InferenceTool } from './tool'
+
+/**
+ * Manages the registration and retrieval of inference tools.
+ */
+export class ToolManager {
+ public tools = new Map<string, InferenceTool>()
+
+ /**
+ * Registers a tool.
+ * @param tool - The tool to register.
+ */
+ register<T extends InferenceTool>(tool: T) {
+ this.tools.set(tool.name, tool)
+ }
+
+ /**
+ * Retrieves a tool by it's name.
+ * @param name - The name of the tool to retrieve.
+ * @returns The tool, if found.
+ */
+ get<T extends InferenceTool>(name: string): T | undefined {
+ return this.tools.get(name) as T | undefined
+ }
+
+ /*
+ ** Process the message request with the tools.
+ */
+ process(request: MessageRequest, tools: AssistantTool[]): Promise<MessageRequest> {
+ return tools.reduce((prevPromise, currentTool) => {
+ return prevPromise.then((prevResult) => {
+ return currentTool.enabled
+ ? this.get(currentTool.type)?.process(prevResult, currentTool) ??
+ Promise.resolve(prevResult)
+ : Promise.resolve(prevResult)
+ })
+ }, Promise.resolve(request))
+ }
+
+ /**
+ * The instance of the tool manager.
+ */
+ static instance(): ToolManager {
+ return (window.core?.toolManager as ToolManager) ?? new ToolManager()
+ }
+}

core/src/browser/tools/tool.ts

@@ -0,0 +1,12 @@
+import { AssistantTool, MessageRequest } from '../../types'
+
+/**
+ * Represents a base inference tool.
+ */
+export abstract class InferenceTool {
+ abstract name: string
+ /*
+ ** Process a message request and return the processed message request.
+ */
+ abstract process(request: MessageRequest, tool?: AssistantTool): Promise<MessageRequest>
+}

core/src/types/model/modelEntity.ts

@@ -7,7 +7,6 @@ export type ModelInfo = {
  settings: ModelSettingParams
  parameters: ModelRuntimeParams
  engine?: InferenceEngine
- proxy_model?: InferenceEngine
 }
 
 /**
@@ -21,8 +20,6 @@ export enum InferenceEngine {
  groq = 'groq',
  triton_trtllm = 'triton_trtllm',
  nitro_tensorrt_llm = 'nitro-tensorrt-llm',
-
- tool_retrieval_enabled = 'tool_retrieval_enabled',
 }
 
 export type ModelArtifact = {
@@ -94,8 +91,6 @@ export type Model = {
  * The model engine.
  */
  engine: InferenceEngine
-
- proxy_model?: InferenceEngine
 }
 
 export type ModelMetadata = {

extensions/assistant-extension/src/index.ts

@@ -1,26 +1,21 @@
 import {
  fs,
  Assistant,
- MessageRequest,
  events,
- InferenceEngine,
- MessageEvent,
- InferenceEvent,
  joinPath,
- executeOnMain,
  AssistantExtension,
  AssistantEvent,
+ ToolManager,
 } from '@janhq/core'
+import { RetrievalTool } from './tools/retrieval'
 
 export default class JanAssistantExtension extends AssistantExtension {
  private static readonly _homeDir = 'file://assistants'
- private static readonly _threadDir = 'file://threads'
-
- controller = new AbortController()
- isCancelled = false
- retrievalThreadId: string | undefined = undefined
 
  async onLoad() {
+ // Register the retrieval tool
+ ToolManager.instance().register(new RetrievalTool())
+
  // making the assistant directory
  const assistantDirExist = await fs.existsSync(
  JanAssistantExtension._homeDir
@@ -38,140 +33,6 @@ export default class JanAssistantExtension extends AssistantExtension {
  // Update the assistant list
  events.emit(AssistantEvent.OnAssistantsUpdate, {})
  }
-
- // Events subscription
- events.on(MessageEvent.OnMessageSent, (data: MessageRequest) =>
- JanAssistantExtension.handleMessageRequest(data, this)
- )
-
- events.on(InferenceEvent.OnInferenceStopped, () => {
- JanAssistantExtension.handleInferenceStopped(this)
- })
- }
-
- private static async handleInferenceStopped(instance: JanAssistantExtension) {
- instance.isCancelled = true
- instance.controller?.abort()
- }
-
- private static async handleMessageRequest(
- data: MessageRequest,
- instance: JanAssistantExtension
- ) {
- instance.isCancelled = false
- instance.controller = new AbortController()
-
- if (
- data.model?.engine !== InferenceEngine.tool_retrieval_enabled ||
- !data.messages ||
- // TODO: Since the engine is defined, its unsafe to assume that assistant tools are defined
- // That could lead to an issue where thread stuck at generating response
- !data.thread?.assistants[0]?.tools
- ) {
- return
- }
-
- const latestMessage = data.messages[data.messages.length - 1]
-
- // 1. Ingest the document if needed
- if (
- latestMessage &&
- latestMessage.content &&
- typeof latestMessage.content !== 'string' &&
- latestMessage.content.length > 1
- ) {
- const docFile = latestMessage.content[1]?.doc_url?.url
- if (docFile) {
- await executeOnMain(
- NODE,
- 'toolRetrievalIngestNewDocument',
- docFile,
- data.model?.proxy_model
- )
- }
- } else if (
- // Check whether we need to ingest document or not
- // Otherwise wrong context will be sent
- !(await fs.existsSync(
- await joinPath([
- JanAssistantExtension._threadDir,
- data.threadId,
- 'memory',
- ])
- ))
- ) {
- // No document ingested, reroute the result to inference engine
- const output = {
- ...data,
- model: {
- ...data.model,
- engine: data.model.proxy_model,
- },
- }
- events.emit(MessageEvent.OnMessageSent, output)
- return
- }
- // 2. Load agent on thread changed
- if (instance.retrievalThreadId !== data.threadId) {
- await executeOnMain(NODE, 'toolRetrievalLoadThreadMemory', data.threadId)
-
- instance.retrievalThreadId = data.threadId
-
- // Update the text splitter
- await executeOnMain(
- NODE,
- 'toolRetrievalUpdateTextSplitter',
- data.thread.assistants[0].tools[0]?.settings?.chunk_size ?? 4000,
- data.thread.assistants[0].tools[0]?.settings?.chunk_overlap ?? 200
- )
- }
-
- // 3. Using the retrieval template with the result and query
- if (latestMessage.content) {
- const prompt =
- typeof latestMessage.content === 'string'
- ? latestMessage.content
- : latestMessage.content[0].text
- // Retrieve the result
- const retrievalResult = await executeOnMain(
- NODE,
- 'toolRetrievalQueryResult',
- prompt
- )
- console.debug('toolRetrievalQueryResult', retrievalResult)
-
- // Update message content
- if (data.thread?.assistants[0]?.tools && retrievalResult)
- data.messages[data.messages.length - 1].content =
- data.thread.assistants[0].tools[0].settings?.retrieval_template
- ?.replace('{CONTEXT}', retrievalResult)
- .replace('{QUESTION}', prompt)
- }
-
- // Filter out all the messages that are not text
- data.messages = data.messages.map((message) => {
- if (
- message.content &&
- typeof message.content !== 'string' &&
- (message.content.length ?? 0) > 0
- ) {
- return {
- ...message,
- content: [message.content[0]],
- }
- }
- return message
- })
-
- // 4. Reroute the result to inference engine
- const output = {
- ...data,
- model: {
- ...data.model,
- engine: data.model.proxy_model,
- },
- }
- events.emit(MessageEvent.OnMessageSent, output)
  }
 
  /**