zclaw_openfang/desktop/src/lib/embedding-client.ts

/**
 * Embedding Client - Vector Embedding Operations
 *
 * Client for interacting with embedding APIs via Tauri backend.
 * Supports multiple providers: OpenAI, Zhipu, Doubao, Qwen, DeepSeek.
 */

import { invoke } from '@tauri-apps/api/core';
import { secureStorage } from './secure-storage';
import { createLogger } from './logger';

const logger = createLogger('embedding-client');

export interface EmbeddingConfig {
  provider: string;
  model: string;
  apiKey: string;
  endpoint: string;
  enabled: boolean;
}

export interface EmbeddingResponse {
  embedding: number[];
  model: string;
  usage?: {
    prompt_tokens: number;
    total_tokens: number;
  };
}

export interface EmbeddingProvider {
  id: string;
  name: string;
  defaultModel: string;
  dimensions: number;
}

const EMBEDDING_STORAGE_KEY = 'zclaw-embedding-config';
const EMBEDDING_KEY_SECURE = 'zclaw-secure-embedding-apikey';

/**
 * Load embedding config from localStorage. apiKey is NOT included;
 * use loadEmbeddingApiKey() to retrieve it from secure storage.
 */
export function loadEmbeddingConfig(): EmbeddingConfig {
  try {
    const stored = localStorage.getItem(EMBEDDING_STORAGE_KEY);
    if (stored) {
      const parsed = JSON.parse(stored);
      return { ...parsed, apiKey: '' };
    }
  } catch (e) {
    logger.debug('Failed to load embedding config', { error: e });
  }
  return {
    provider: 'local',
    model: 'tfidf',
    apiKey: '',
    endpoint: '',
    enabled: false,
  };
}

/**
 * Save embedding config to localStorage. API key is NOT saved here;
 * use saveEmbeddingApiKey() separately.
 */
export function saveEmbeddingConfig(config: EmbeddingConfig): void {
  try {
    const { apiKey: _, ...rest } = config;
    localStorage.setItem(EMBEDDING_STORAGE_KEY, JSON.stringify(rest));
  } catch (e) {
    logger.debug('Failed to save embedding config', { error: e });
  }
}

/**
 * Load embedding API key from secure storage.
 */
export async function loadEmbeddingApiKey(): Promise<string | null> {
  return secureStorage.get(EMBEDDING_KEY_SECURE);
}

/**
 * Save embedding API key to secure storage.
 */
export async function saveEmbeddingApiKey(apiKey: string): Promise<void> {
  if (!apiKey.trim()) {
    await secureStorage.delete(EMBEDDING_KEY_SECURE);
    return;
  }
  await secureStorage.set(EMBEDDING_KEY_SECURE, apiKey.trim());
}

export async function getEmbeddingProviders(): Promise<EmbeddingProvider[]> {
  const result = await invoke<[string, string, string, number][]>('embedding_providers');
  return result.map(([id, name, defaultModel, dimensions]) => ({
    id,
    name,
    defaultModel,
    dimensions,
  }));
}

export async function createEmbedding(
  text: string,
  config?: Partial<EmbeddingConfig>
): Promise<EmbeddingResponse> {
  const savedConfig = loadEmbeddingConfig();
  const provider = config?.provider ?? savedConfig.provider;
  // Resolve apiKey: use explicit config value, then secure storage, then empty
  const explicitKey = config?.apiKey?.trim();
  const apiKey = explicitKey || await loadEmbeddingApiKey() || '';
  const model = config?.model ?? savedConfig.model;
  const endpoint = config?.endpoint ?? savedConfig.endpoint;

  if (provider === 'local') {
    throw new Error('Local TF-IDF mode does not support API embedding');
  }

  if (!apiKey) {
    throw new Error('API Key is required for embedding');
  }

  return invoke<EmbeddingResponse>('embedding_create', {
    provider,
    apiKey,
    text,
    model: model || undefined,
    endpoint: endpoint || undefined,
  });
}

export async function createEmbeddings(
  texts: string[],
  config?: Partial<EmbeddingConfig>
): Promise<EmbeddingResponse[]> {
  const results: EmbeddingResponse[] = [];
  for (const text of texts) {
    const result = await createEmbedding(text, config);
    results.push(result);
  }
  return results;
}

export function cosineSimilarity(a: number[], b: number[]): number {
  if (a.length !== b.length) {
    throw new Error('Vectors must have the same length');
  }

  let dotProduct = 0;
  let normA = 0;
  let normB = 0;

  for (let i = 0; i < a.length; i++) {
    dotProduct += a[i] * b[i];
    normA += a[i] * a[i];
    normB += b[i] * b[i];
  }

  const denom = Math.sqrt(normA * normB);
  if (denom === 0) {
    return 0;
  }

  return dotProduct / denom;
}

export class EmbeddingClient {
  private config: EmbeddingConfig;

  constructor(config?: EmbeddingConfig) {
    this.config = config ?? loadEmbeddingConfig();
    // If no explicit apiKey was provided and config was loaded from localStorage,
    // the apiKey will be empty. It will be resolved from secure storage lazily.
  }

  get isApiMode(): boolean {
    return this.config.provider !== 'local' && this.config.enabled && !!this.config.apiKey;
  }

  async embed(text: string): Promise<number[]> {
    // Resolve apiKey from secure storage if not in config
    const effectiveConfig = this.config.apiKey
      ? this.config
      : { ...this.config, apiKey: await loadEmbeddingApiKey() ?? '' };
    const response = await createEmbedding(text, effectiveConfig);
    return response.embedding;
  }

  async embedBatch(texts: string[]): Promise<number[][]> {
    const responses = await createEmbeddings(texts, this.config);
    return responses.map(r => r.embedding);
  }

  similarity(vec1: number[], vec2: number[]): number {
    return cosineSimilarity(vec1, vec2);
  }

  updateConfig(config: Partial<EmbeddingConfig>): void {
    this.config = { ...this.config, ...config };
    if (config.provider !== undefined || config.apiKey !== undefined) {
      this.config.enabled = this.config.provider !== 'local' && !!this.config.apiKey;
    }
    // Save non-key fields to localStorage
    saveEmbeddingConfig(this.config);
    // Save apiKey to secure storage (fire-and-forget)
    if (config.apiKey !== undefined) {
      saveEmbeddingApiKey(config.apiKey).catch(e => logger.debug('Failed to save embedding API key', { error: e }));
    }
  }

  getConfig(): EmbeddingConfig {
    return { ...this.config };
  }
}

let embeddingClientInstance: EmbeddingClient | null = null;

export function getEmbeddingClient(): EmbeddingClient {
  if (!embeddingClientInstance) {
    embeddingClientInstance = new EmbeddingClient();
  }
  return embeddingClientInstance;
}

export function resetEmbeddingClient(): void {
  embeddingClientInstance = null;
}