From 6272ab107f13c1e1b96ace971b9df7a264ac8d72 Mon Sep 17 00:00:00 2001
From: Rotimi Best <best@classroomio.com>
Date: Tue, 19 May 2026 04:58:47 +0100
Subject: [PATCH] refactor: streamline text and title generation model creation
 by consolidating provider logic

---
 .../api/src/services/agent/text-generation.ts | 32 ++-----------------
 .../src/services/agent/title-generation.ts    | 28 ++--------------
 2 files changed, 4 insertions(+), 56 deletions(-)

diff --git a/apps/api/src/services/agent/text-generation.ts b/apps/api/src/services/agent/text-generation.ts
index fcdb8d1e..97fabcc8 100644
--- a/apps/api/src/services/agent/text-generation.ts
+++ b/apps/api/src/services/agent/text-generation.ts
@@ -1,8 +1,5 @@
 import { generateText } from 'ai';
-import { AIProvider, type AIProviderConfig } from '@cio/ai-assistant';
-import { createOpenAI } from '@ai-sdk/openai';
-import { createAnthropic } from '@ai-sdk/anthropic';
-import { createGoogleGenerativeAI } from '@ai-sdk/google';
+import { AIProvider, type AIProviderConfig, createModel } from '@cio/ai-assistant';
 
 const TEXT_GEN_MODELS: Record<AIProvider, string> = {
   [AIProvider.OPENAI]: 'gpt-4o-mini',
@@ -11,31 +8,6 @@ const TEXT_GEN_MODELS: Record<AIProvider, string> = {
   [AIProvider.MOONSHOT]: 'kimi-k2.6'
 };
 
-function createTextGenModel(config: AIProviderConfig) {
-  const modelName = TEXT_GEN_MODELS[config.provider];
-
-  switch (config.provider) {
-    case AIProvider.OPENAI: {
-      const openai = createOpenAI({ apiKey: config.apiKey });
-      return openai(modelName);
-    }
-    case AIProvider.ANTHROPIC: {
-      const anthropic = createAnthropic({ apiKey: config.apiKey });
-      return anthropic(modelName);
-    }
-    case AIProvider.GOOGLE: {
-      const google = createGoogleGenerativeAI({ apiKey: config.apiKey });
-      return google(modelName);
-    }
-    case AIProvider.MOONSHOT: {
-      const moonshot = createOpenAI({ apiKey: config.apiKey, baseURL: 'https://api.moonshot.cn/v1' });
-      return moonshot(modelName);
-    }
-    default:
-      throw new Error(`Unsupported AI provider: ${config.provider}`);
-  }
-}
-
 export interface GenerateFieldTextResult {
   text: string;
   usage: { promptTokens: number; completionTokens: number; totalTokens: number };
@@ -49,8 +21,8 @@ export async function generateFieldText(
   context: string | undefined,
   providerConfig: AIProviderConfig
 ): Promise<GenerateFieldTextResult> {
-  const model = createTextGenModel(providerConfig);
   const modelName = TEXT_GEN_MODELS[providerConfig.provider];
+  const model = createModel({ ...providerConfig, model: modelName });
 
   const contextLine = context ? `Context: ${context}.` : '';
   const formatLine =
diff --git a/apps/api/src/services/agent/title-generation.ts b/apps/api/src/services/agent/title-generation.ts
index b2f8d465..4798da00 100644
--- a/apps/api/src/services/agent/title-generation.ts
+++ b/apps/api/src/services/agent/title-generation.ts
@@ -1,9 +1,6 @@
 import { generateText, generateObject } from 'ai';
 import { z } from 'zod';
-import { AIProvider, type AIProviderConfig } from '@cio/ai-assistant';
-import { createOpenAI } from '@ai-sdk/openai';
-import { createAnthropic } from '@ai-sdk/anthropic';
-import { createGoogleGenerativeAI } from '@ai-sdk/google';
+import { AIProvider, type AIProviderConfig, createModel } from '@cio/ai-assistant';
 
 /**
  * Cheap models used specifically for title generation to minimize token cost.
@@ -16,28 +13,7 @@ const TITLE_MODELS: Record<AIProvider, string> = {
 };
 
 function createTitleModel(config: AIProviderConfig) {
-  const modelName = TITLE_MODELS[config.provider];
-
-  switch (config.provider) {
-    case AIProvider.OPENAI: {
-      const openai = createOpenAI({ apiKey: config.apiKey });
-      return openai(modelName);
-    }
-    case AIProvider.ANTHROPIC: {
-      const anthropic = createAnthropic({ apiKey: config.apiKey });
-      return anthropic(modelName);
-    }
-    case AIProvider.GOOGLE: {
-      const google = createGoogleGenerativeAI({ apiKey: config.apiKey });
-      return google(modelName);
-    }
-    case AIProvider.MOONSHOT: {
-      const moonshot = createOpenAI({ apiKey: config.apiKey, baseURL: 'https://api.moonshot.cn/v1' });
-      return moonshot(modelName);
-    }
-    default:
-      throw new Error(`Unsupported AI provider: ${config.provider}`);
-  }
+  return createModel({ ...config, model: TITLE_MODELS[config.provider] });
 }
 
 /**