Add search and RAG workspace flows

lib/server/ai.ts

@@ -1,4 +1,4 @@
-import { generateText } from 'ai';
+import { embedMany, generateText } from 'ai';
 import { createZhipu } from 'zhipu-ai-provider';
 
 type AiWorkload = 'report' | 'extraction';
@@ -31,13 +31,35 @@ type AiGenerateOutput = {
   text: string;
 };
 
+type AiEmbedOutput = {
+  embeddings: number[][];
+};
+
 type RunAiAnalysisOptions = GetAiConfigOptions & {
   workload?: AiWorkload;
   createModel?: (config: AiConfig) => unknown;
   generate?: (input: AiGenerateInput) => Promise<AiGenerateOutput>;
 };
 
+type EmbeddingConfig = {
+  provider: AiProvider;
+  apiKey?: string;
+  baseUrl: string;
+  model: 'embedding-3';
+  dimensions: 256;
+};
+
+type RunAiEmbeddingsOptions = GetAiConfigOptions & {
+  createModel?: (config: EmbeddingConfig) => unknown;
+  embed?: (input: {
+    model: unknown;
+    values: string[];
+  }) => Promise<AiEmbedOutput>;
+};
+
 const CODING_API_BASE_URL = 'https://api.z.ai/api/coding/paas/v4';
+const SEARCH_EMBEDDING_MODEL = 'embedding-3';
+const SEARCH_EMBEDDING_DIMENSIONS = 256;
 
 let warnedIgnoredZhipuBaseUrl = false;
 
@@ -97,6 +119,30 @@ async function defaultGenerate(input: AiGenerateInput): Promise<AiGenerateOutput
   return { text: result.text };
 }
 
+function defaultCreateEmbeddingModel(config: EmbeddingConfig) {
+  const zhipu = createZhipu({
+    apiKey: config.apiKey,
+    baseURL: config.baseUrl
+  });
+
+  return zhipu.textEmbeddingModel(config.model, {
+    dimensions: config.dimensions
+  });
+}
+
+async function defaultEmbed(input: {
+  model: unknown;
+  values: string[];
+}): Promise<AiEmbedOutput> {
+  const result = await embedMany({
+    model: input.model as never,
+    values: input.values,
+    maxRetries: 0
+  });
+
+  return { embeddings: result.embeddings as number[][] };
+}
+
 export function getAiConfig(options?: GetAiConfigOptions) {
   return getReportAiConfig(options);
 }
@@ -121,6 +167,19 @@ export function getExtractionAiConfig(options?: GetAiConfigOptions) {
   };
 }
 
+export function getEmbeddingAiConfig(options?: GetAiConfigOptions) {
+  const env = options?.env ?? process.env;
+  warnIgnoredZhipuBaseUrl(env, options?.warn ?? console.warn);
+
+  return {
+    provider: 'zhipu',
+    apiKey: envValue('ZHIPU_API_KEY', env),
+    baseUrl: CODING_API_BASE_URL,
+    model: SEARCH_EMBEDDING_MODEL,
+    dimensions: SEARCH_EMBEDDING_DIMENSIONS
+  } satisfies EmbeddingConfig;
+}
+
 export function isAiConfigured(options?: GetAiConfigOptions) {
   const config = getReportAiConfig(options);
   return Boolean(config.apiKey);
@@ -160,6 +219,31 @@ export async function runAiAnalysis(prompt: string, systemPrompt?: string, optio
   };
 }
 
+export async function runAiEmbeddings(values: string[], options?: RunAiEmbeddingsOptions) {
+  const sanitizedValues = values
+    .map((value) => value.trim())
+    .filter((value) => value.length > 0);
+
+  if (sanitizedValues.length === 0) {
+    return [];
+  }
+
+  const config = getEmbeddingAiConfig(options);
+  if (!config.apiKey) {
+    throw new Error('ZHIPU_API_KEY is required for AI workloads');
+  }
+
+  const createModel = options?.createModel ?? defaultCreateEmbeddingModel;
+  const embed = options?.embed ?? defaultEmbed;
+  const model = createModel(config);
+  const result = await embed({
+    model,
+    values: sanitizedValues
+  });
+
+  return result.embeddings.map((embedding) => embedding.map((value) => Number(value)));
+}
+
 export function __resetAiWarningsForTests() {
   warnedIgnoredZhipuBaseUrl = false;
 }