chore(embeddings): use framework embeddings, refactor ai providers

Signed-off-by: Radek Ježek <[email protected]>

.env.example

@@ -45,10 +45,7 @@ RUN_BULLMQ_WORKERS=runs,runs-cleanup,vectorStores-cleanup,vectorStores-fileProce
 # --- BACKEND SECTION ---
 
 # LLM backend, possible values are: ollama, openai, watsonx, bam, ibm-vllm
-LLM_BACKEND=watsonx
-
-# Embedding backend, possible values are: ollama, openai, watsonx, bam, caikit
-EMBEDDING_BACKEND=watsonx
+AI_BACKEND=watsonx
 
 # Extraction backend, possible values are: docling, unstructured-opensource, unstructured-api
 EXTRACTION_BACKEND=docling

README.md

@@ -90,7 +90,7 @@ pnpm mikro-orm seeder:run
 cp .env.example .env
 ```
 
-- Add values the env vars: CRYPTO_CIPHER_KEY, LLM_BACKEND, EMBEDDING_BACKEND and API key for which ever provider you have chosen.
+- Add values the env vars: CRYPTO_CIPHER_KEY, AI_BACKEND and API key for which ever provider you have chosen.
 
 - Run the bee-api:
 

seeders/DatabaseSeeder.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { ref, type EntityManager } from '@mikro-orm/core';
+import { type EntityManager, ref } from '@mikro-orm/core';
 import { Seeder } from '@mikro-orm/seeder';
 
 import { OrganizationUserRole, ProjectRole } from '@/administration/entities/constants';
@@ -30,7 +30,8 @@ import { ProjectApiKey } from '@/administration/entities/project-api-key.entity'
 import { API_KEY_PREFIX, scryptSecret } from '@/auth/utils';
 import { IBM_ORGANIZATION_OWNER_ID } from '@/config';
 import { redactProjectKeyValue } from '@/administration/helpers';
-import { Agent, getDefaultModel } from '@/runs/execution/constants';
+import { Agent } from '@/runs/execution/constants';
+import { defaultAIProvider } from '@/runs/execution/provider';
 
 const USER_EXTERNAL_ID = 'test';
 const PROJECT_API_KEY = `${API_KEY_PREFIX}testkey`;
@@ -91,7 +92,7 @@ export class DatabaseSeeder extends Seeder {
       redactedValue: redactProjectKeyValue(PROJECT_API_KEY)
     });
     const beeAssistant = new Assistant({
-      model: getDefaultModel(),
+      model: defaultAIProvider.createChatLLM().modelId,
       agent: Agent.BEE,
       tools: [
         {
@@ -120,7 +121,7 @@ export class DatabaseSeeder extends Seeder {
       }
     });
     const streamlitAssistant = new Assistant({
-      model: getDefaultModel(),
+      model: defaultAIProvider.createChatLLM().modelId,
       agent: Agent.STREAMLIT,
       tools: [],
       name: 'Builder Assistant',

src/assistants/assistants.service.ts

@@ -41,7 +41,8 @@ import { APIError, APIErrorCode } from '@/errors/error.entity.js';
 import { Tool, ToolType } from '@/tools/entities/tool/tool.entity.js';
 import { getUpdatedValue } from '@/utils/update.js';
 import { createDeleteResponse } from '@/utils/delete.js';
-import { Agent, getDefaultModel } from '@/runs/execution/constants.js';
+import { Agent } from '@/runs/execution/constants.js';
+import { defaultAIProvider } from '@/runs/execution/provider';
 
 export function toDto(assistant: Loaded<Assistant>): AssistantDto {
   return {
@@ -113,7 +114,7 @@ export async function createAssistant({
     description: description ?? undefined,
     metadata,
     topP: top_p ?? undefined,
-    model: model ?? getDefaultModel(),
+    model: model ?? defaultAIProvider.createChatLLM().modelId,
     agent,
     temperature: temperature ?? undefined,
     systemPromptOverwrite: system_prompt_overwrite ?? undefined

src/chat/chat.service.ts

@@ -26,11 +26,10 @@ import {
 import { ChatMessageRole } from './constants';
 import { Chat } from './entities/chat.entity';
 
-import { createChatLLM } from '@/runs/execution/factory';
 import { getLogger } from '@/logger';
 import { APIError, APIErrorCode } from '@/errors/error.entity';
-import { getDefaultModel } from '@/runs/execution/constants';
 import { ORM } from '@/database';
+import { defaultAIProvider } from '@/runs/execution/provider';
 
 const getChatLogger = () => getLogger();
 
@@ -52,12 +51,12 @@ export function toDto(chat: Loaded<Chat>): ChatCompletionCreateResponse {
 }
 
 export async function createChatCompletion({
-  model = getDefaultModel(),
+  model,
   messages,
   response_format
 }: ChatCompletionCreateBody): Promise<ChatCompletionCreateResponse> {
-  const llm = createChatLLM({ model });
-  const chat = new Chat({ model, messages, responseFormat: response_format });
+  const llm = defaultAIProvider.createChatLLM({ model });
+  const chat = new Chat({ model: llm.modelId, messages, responseFormat: response_format });
   await ORM.em.persistAndFlush(chat);
   try {
     const schema = response_format?.json_schema.schema;

src/config.ts

@@ -18,11 +18,10 @@ import 'dotenv/config';
 import { difference } from 'remeda';
 
 import {
+  AIBackend,
   CodeInterpreterStorageBackend,
-  LLMBackend,
   SearchToolBackend
 } from './runs/execution/constants';
-import { EmbeddingBackend } from './embedding/constants';
 import { ExtractionBackend } from './files/extraction/constants';
 
 import { QueueName } from '@/jobs/constants.js';
@@ -81,8 +80,7 @@ export const AUTH_CLIENT_SECRET = getEnv('AUTH_CLIENT_SECRET');
 export const AUTH_AUDIENCE = getEnv('AUTH_AUDIENCE');
 
 // Backends
-export const LLM_BACKEND = getEnum('LLM_BACKEND', Object.values(LLMBackend));
-export const EMBEDDING_BACKEND = getEnum('EMBEDDING_BACKEND', Object.values(EmbeddingBackend));
+export const AI_BACKEND = getEnum('AI_BACKEND', Object.values(AIBackend));
 export const EXTRACTION_BACKEND = getEnum('EXTRACTION_BACKEND', Object.values(ExtractionBackend));
 
 export const OLLAMA_URL = getEnv('OLLAMA_URL', null);