@dooor-ai/toolkit/src/observability/cortexdb-client.ts

Guards, Evals & Observability for AI applications - works seamlessly with LangChain/LangGraph

/**
 * CortexDB Client for AI Provider proxy
 *
 * Used by guards and evals to call LLMs without exposing API keys.
 * CortexDB acts as a proxy to configured AI providers.
 */

import type { RAGContext } from "../rag/context";
import type { RAGMetadata } from "../rag/types";

export interface CortexDBConfig {
  /** CortexDB Gateway URL (e.g., "http://localhost:8000" or "https://your-cortex.com") */
  baseUrl: string;
  /** API key for CortexDB authentication */
  apiKey: string;
  /** Database name for observability (optional) */
  database?: string;
  /** Project name for tracing (optional) */
  project?: string;
}

export interface AIInvokeRequest {
  /** Prompt to send to the AI */
  prompt: string;
  /** Usage type: evaluation or guard */
  usage: "evaluation" | "guard" | "chat";
  /** Max tokens (optional) */
  maxTokens?: number;
  /** Temperature (optional, 0-1) */
  temperature?: number;
  /** AI Provider name (configured in CortexDB Studio) */
  providerName?: string;
  /** RAG context for retrieval-augmented generation (optional) */
  ragContext?: RAGContext;
}

export interface RAGChunk {
  text: string;
  score: number;
  metadata?: any;
}

export interface AIInvokeResponse {
  /** Generated text (or RAG context if using RAG) */
  text: string;
  /** Token usage */
  usage: {
    totalTokens: number;
    promptTokens: number;
    completionTokens: number;
  };
  /** RAG metadata (if RAG context was provided) */
  ragMetadata?: RAGMetadata;
  /** RAG chunks (if RAG context was provided) */
  ragChunks?: RAGChunk[];
}

/**
 * CortexDB Client - Proxy to AI providers configured in CortexDB Studio
 */
export class CortexDBClient {
  private config: CortexDBConfig;

  constructor(config: CortexDBConfig) {
    this.config = config;
  }

  /**
   * Get client configuration (for internal use by observability backend)
   */
  getConfig(): CortexDBConfig {
    return this.config;
  }

  /**
   * Invoke AI via CortexDB proxy
   *
   * CortexDB will use the configured AI provider (Gemini, OpenAI, etc)
   * based on the usage type and provider name.
   *
   * Optionally supports RAG (Retrieval-Augmented Generation) by passing ragContext.
   */
  async invokeAI(request: AIInvokeRequest): Promise<AIInvokeResponse> {
    // Use RAG endpoint if RAG context is provided
    const isRAG = !!request.ragContext;
    const url = isRAG
      ? `${this.config.baseUrl}/api/ai/rag/invoke`
      : `${this.config.baseUrl}/api/ai/invoke`;

    const payload: any = isRAG
      ? {
          // RAG endpoint payload
          prompt: request.prompt,
          rag_context: request.ragContext!.toJSON(),
          database: this.config.database,
          provider_name: request.providerName,
          max_tokens: request.maxTokens,
          temperature: request.temperature,
        }
      : {
          // Normal AI invoke payload
          prompt: request.prompt,
          usage: request.usage,
          max_tokens: request.maxTokens,
          temperature: request.temperature,
          provider_name: request.providerName,
        };

    const response = await fetch(url, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "Authorization": `Bearer ${this.config.apiKey}`,
      },
      body: JSON.stringify(payload),
    });

    if (!response.ok) {
      const error = await response.text();
      throw new Error(`CortexDB AI invoke failed: ${error}`);
    }

    const data = await response.json() as any;

    // Handle RAG response (new format: context + chunks + metadata)
    if (isRAG && data.context !== undefined) {
      return {
        text: data.context, // RAG returns context, not LLM response
        usage: {
          totalTokens: 0,
          promptTokens: 0,
          completionTokens: 0,
        },
        ragMetadata: data.metadata, // New format uses "metadata" not "rag_metadata"
        ragChunks: data.chunks, // Individual chunks with scores
      };
    }

    // Handle normal AI invoke response (old format)
    return {
      text: data.text,
      usage: {
        totalTokens: data.usage?.total_tokens || 0,
        promptTokens: data.usage?.prompt_tokens || 0,
        completionTokens: data.usage?.completion_tokens || 0,
      },
      ragMetadata: data.rag_metadata,
    };
  }

  /**
   * Save trace to CortexDB (for observability)
   */
  async saveTrace(trace: any): Promise<void> {
    if (!this.config.database) {
      throw new Error("Database not configured for observability");
    }

    const url = `${this.config.baseUrl}/databases/${this.config.database}/traces`;

    const payload = {
      timestamp: Date.now(),
      ...trace,
    };

    console.log("[CortexDBClient] POST", url);
    const response = await fetch(url, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "Authorization": `Bearer ${this.config.apiKey}`,
      },
      body: JSON.stringify(payload),
    });

    if (!response.ok) {
      const errorText = await response.text();
      console.error("[CortexDBClient] Failed to save trace:", response.status, errorText);
      throw new Error(`Failed to save trace: ${response.status} ${errorText}`);
    }

    console.log("[CortexDBClient] Trace saved, status:", response.status);
  }

  /**
   * Save eval result to CortexDB
   */
  async saveEval(evalResult: any): Promise<void> {
    if (!this.config.database) {
      throw new Error("Database not configured for observability");
    }

    const url = `${this.config.baseUrl}/databases/${this.config.database}/evals`;

    console.log("[CortexDBClient] POST", url);
    const response = await fetch(url, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "Authorization": `Bearer ${this.config.apiKey}`,
      },
      body: JSON.stringify(evalResult),
    });

    if (!response.ok) {
      const errorText = await response.text();
      console.error("[CortexDBClient] Failed to save eval:", response.status, errorText);
      throw new Error(`Failed to save eval: ${response.status} ${errorText}`);
    }

    console.log("[CortexDBClient] Eval saved, status:", response.status);
  }

  /**
   * Update tool calls for a trace
   */
  async updateToolCalls(traceId: string, toolCalls: any[]): Promise<void> {
    if (!this.config.database) {
      throw new Error("Database not configured for observability");
    }

    const url = `${this.config.baseUrl}/databases/${this.config.database}/traces/${traceId}/tool_calls`;

    // Retry logic: sometimes the trace might not be committed yet
    const maxRetries = 3;
    const retryDelay = 500; // ms
    
    for (let attempt = 1; attempt <= maxRetries; attempt++) {
      console.log(`[CortexDBClient] POST ${url} (attempt ${attempt}/${maxRetries})`);
      const response = await fetch(url, {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          "Authorization": `Bearer ${this.config.apiKey}`,
        },
        body: JSON.stringify({ tool_calls: toolCalls }),
      });

      if (response.ok) {
        console.log("[CortexDBClient] Tool calls updated, status:", response.status);
        return;
      }

      const errorText = await response.text();
      
      // If 404 and not last attempt, retry after delay
      if (response.status === 404 && attempt < maxRetries) {
        console.log(`[CortexDBClient] Trace not found (404), retrying in ${retryDelay}ms...`);
        await new Promise(resolve => setTimeout(resolve, retryDelay));
        continue;
      }

      // For other errors or last attempt, throw
      console.error("[CortexDBClient] Failed to update tool calls:", response.status, errorText);
      throw new Error(`Failed to update tool calls: ${response.status} ${errorText}`);
    }
  }
}

/**
 * Parse CortexDB connection string
 * 
 * Format: cortexdb://api_key@host:port/database
 * Example: cortexdb://cortexdb_adm123@35.223.201.25:8000/my_evals
 */
export function parseCortexDBConnectionString(connectionString: string): CortexDBConfig {
  const match = connectionString.match(/^cortexdb:\/\/([^@]+)@([^:]+):(\d+)\/(.+)$/);
  
  if (!match) {
    throw new Error(
      'Invalid CortexDB connection string format. Expected: cortexdb://api_key@host:port/database'
    );
  }

  const [, apiKey, host, port, database] = match;

  return {
    baseUrl: `http://${host}:${port}`,
    apiKey,
    database,
  };
}

/**
 * Global CortexDB client instance (configured via configureCortexDB or toolkitConfig)
 */
let globalCortexDBClient: CortexDBClient | null = null;
let globalProviderName: string | null = null;

/**
 * Configure global CortexDB client
 */
export function configureCortexDB(config: CortexDBConfig): void {
  globalCortexDBClient = new CortexDBClient(config);
}

/**
 * Configure CortexDB from connection string (used by toolkitConfig)
 */
export function configureCortexDBFromConnectionString(
  connectionString: string,
  providerName?: string,
  project?: string
): void {
  const config = parseCortexDBConnectionString(connectionString);
  
  if (project) {
    config.project = project;
  }
  
  globalCortexDBClient = new CortexDBClient(config);
  globalProviderName = providerName || null;
}

/**
 * Get global CortexDB client
 */
export function getCortexDBClient(): CortexDBClient {
  if (!globalCortexDBClient) {
    throw new Error(
      "CortexDB client not configured. Provide toolkitConfig.apiKey in ChatDooorGenerativeAI constructor."
    );
  }
  return globalCortexDBClient;
}

/**
 * Get global provider name (set via toolkitConfig)
 */
export function getGlobalProviderName(): string | null {
  return globalProviderName;
}