Use token.js for non-tools calls (#4973)

* deprecate inference()

* fix tool_choice. only startModel for Cortex

* appease linter

* remove sse

* add stopInferencing support. temporarily with OpenAI

* use abortSignal in token.js

* bump token.js version

Author: gau-nernst

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

@@ -6,6 +6,7 @@ import { AIEngine } from './AIEngine'
  */
 export class EngineManager {
   public engines = new Map<string, AIEngine>()
+  public controller: AbortController | null = null
 
   /**
    * Registers an engine.

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

@@ -12,11 +12,7 @@ import {
   ChatCompletionRole,
   ContentType,
 } from '../../../types'
-import { requestInference } from './helpers/sse'
-import { ulid } from 'ulidx'
 
-jest.mock('./helpers/sse')
-jest.mock('ulidx')
 jest.mock('../../events')
 
 class TestOAIEngine extends OAIEngine {
@@ -48,79 +44,6 @@ describe('OAIEngine', () => {
     )
   })
 
-  it('should handle inference request', async () => {
-    const data: MessageRequest = {
-      model: { engine: 'test-provider', id: 'test-model' } as any,
-      threadId: 'test-thread',
-      type: MessageRequestType.Thread,
-      assistantId: 'test-assistant',
-      messages: [{ role: ChatCompletionRole.User, content: 'Hello' }],
-    }
-
-    ;(ulid as jest.Mock).mockReturnValue('test-id')
-    ;(requestInference as jest.Mock).mockReturnValue({
-      subscribe: ({ next, complete }: any) => {
-        next('test response')
-        complete()
-      },
-    })
-
-    await engine.inference(data)
-
-    expect(requestInference).toHaveBeenCalledWith(
-      'http://test-inference-url',
-      expect.objectContaining({ model: 'test-model' }),
-      expect.any(Object),
-      expect.any(AbortController),
-      { Authorization: 'Bearer test-token' },
-      undefined
-    )
-
-    expect(events.emit).toHaveBeenCalledWith(
-      MessageEvent.OnMessageResponse,
-      expect.objectContaining({ id: 'test-id' })
-    )
-    expect(events.emit).toHaveBeenCalledWith(
-      MessageEvent.OnMessageUpdate,
-      expect.objectContaining({
-        content: [
-          {
-            type: ContentType.Text,
-            text: { value: 'test response', annotations: [] },
-          },
-        ],
-        status: MessageStatus.Ready,
-      })
-    )
-  })
-
-  it('should handle inference error', async () => {
-    const data: MessageRequest = {
-      model: { engine: 'test-provider', id: 'test-model' } as any,
-      threadId: 'test-thread',
-      type: MessageRequestType.Thread,
-      assistantId: 'test-assistant',
-      messages: [{ role: ChatCompletionRole.User, content: 'Hello' }],
-    }
-
-    ;(ulid as jest.Mock).mockReturnValue('test-id')
-    ;(requestInference as jest.Mock).mockReturnValue({
-      subscribe: ({ error }: any) => {
-        error({ message: 'test error', code: 500 })
-      },
-    })
-
-    await engine.inference(data)
-
-    expect(events.emit).toHaveBeenLastCalledWith(
-      MessageEvent.OnMessageUpdate,
-      expect.objectContaining({
-        status: 'error',
-        error_code: 500,
-      })
-    )
-  })
-
   it('should stop inference', () => {
     engine.stopInference()
     expect(engine.isCancelled).toBe(true)

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

@@ -1,18 +1,9 @@
-import { requestInference } from './helpers/sse'
-import { ulid } from 'ulidx'
 import { AIEngine } from './AIEngine'
 import {
-  ChatCompletionRole,
-  ContentType,
   InferenceEvent,
   MessageEvent,
   MessageRequest,
-  MessageRequestType,
-  MessageStatus,
   Model,
-  ModelInfo,
-  ThreadContent,
-  ThreadMessage,
 } from '../../../types'
 import { events } from '../../events'
 
@@ -53,112 +44,6 @@ export abstract class OAIEngine extends AIEngine {
    */
   override onUnload(): void {}
 
-  /*
-   * Inference request
-   */
-  override async inference(data: MessageRequest) {
-    if (!data.model?.id) {
-      events.emit(MessageEvent.OnMessageResponse, {
-        status: MessageStatus.Error,
-        content: [
-          {
-            type: ContentType.Text,
-            text: {
-              value: 'No model ID provided',
-              annotations: [],
-            },
-          },
-        ],
-      })
-      return
-    }
-
-    const timestamp = Date.now() / 1000
-    const message: ThreadMessage = {
-      id: ulid(),
-      thread_id: data.thread?.id ?? data.threadId,
-      type: data.type,
-      assistant_id: data.assistantId,
-      role: ChatCompletionRole.Assistant,
-      content: [],
-      status: MessageStatus.Pending,
-      created_at: timestamp,
-      completed_at: timestamp,
-      object: 'thread.message',
-    }
-
-    if (data.type !== MessageRequestType.Summary) {
-      events.emit(MessageEvent.OnMessageResponse, message)
-    }
-
-    this.isCancelled = false
-    this.controller = new AbortController()
-
-    const model: ModelInfo = {
-      ...(this.loadedModel ? this.loadedModel : {}),
-      ...data.model,
-    }
-
-    const header = await this.headers()
-    let requestBody = {
-      messages: data.messages ?? [],
-      model: model.id,
-      stream: true,
-      tools: data.tools,
-      ...model.parameters,
-    }
-    if (this.transformPayload) {
-      requestBody = this.transformPayload(requestBody)
-    }
-
-    requestInference(
-      this.inferenceUrl,
-      requestBody,
-      model,
-      this.controller,
-      header,
-      this.transformResponse
-    ).subscribe({
-      next: (content: any) => {
-        const messageContent: ThreadContent = {
-          type: ContentType.Text,
-          text: {
-            value: content.trim(),
-            annotations: [],
-          },
-        }
-        message.content = [messageContent]
-        events.emit(MessageEvent.OnMessageUpdate, message)
-      },
-      complete: async () => {
-        message.status = message.content.length
-          ? MessageStatus.Ready
-          : MessageStatus.Error
-        events.emit(MessageEvent.OnMessageUpdate, message)
-      },
-      error: async (err: any) => {
-        if (this.isCancelled || message.content.length) {
-          message.status = MessageStatus.Stopped
-          events.emit(MessageEvent.OnMessageUpdate, message)
-          return
-        }
-        message.status = MessageStatus.Error
-        message.content[0] = {
-          type: ContentType.Text,
-          text: {
-            value:
-              typeof message === 'string'
-                ? err.message
-                : (JSON.stringify(err.message) ?? err.detail),
-            annotations: [],
-          },
-        }
-        message.error_code = err.code
-        events.emit(MessageEvent.OnMessageUpdate, message)
-      },
-    })
-  }
-
   /**
    * Stops the inference.
    */