@dooor-ai/toolkit
Version:
Guards, Evals & Observability for AI applications - works seamlessly with LangChain/LangGraph
243 lines (242 loc) • 9.2 kB
JavaScript
;
/**
* 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.
*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.CortexDBClient = void 0;
exports.parseCortexDBConnectionString = parseCortexDBConnectionString;
exports.configureCortexDB = configureCortexDB;
exports.configureCortexDBFromConnectionString = configureCortexDBFromConnectionString;
exports.getCortexDBClient = getCortexDBClient;
exports.getGlobalProviderName = getGlobalProviderName;
/**
* CortexDB Client - Proxy to AI providers configured in CortexDB Studio
*/
class CortexDBClient {
constructor(config) {
this.config = config;
}
/**
* Get client configuration (for internal use by observability backend)
*/
getConfig() {
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) {
// 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 = 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();
// 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) {
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) {
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, toolCalls) {
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}`);
}
}
}
exports.CortexDBClient = CortexDBClient;
/**
* Parse CortexDB connection string
*
* Format: cortexdb://api_key@host:port/database
* Example: cortexdb://cortexdb_adm123@35.223.201.25:8000/my_evals
*/
function parseCortexDBConnectionString(connectionString) {
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 = null;
let globalProviderName = null;
/**
* Configure global CortexDB client
*/
function configureCortexDB(config) {
globalCortexDBClient = new CortexDBClient(config);
}
/**
* Configure CortexDB from connection string (used by toolkitConfig)
*/
function configureCortexDBFromConnectionString(connectionString, providerName, project) {
const config = parseCortexDBConnectionString(connectionString);
if (project) {
config.project = project;
}
globalCortexDBClient = new CortexDBClient(config);
globalProviderName = providerName || null;
}
/**
* Get global CortexDB client
*/
function getCortexDBClient() {
if (!globalCortexDBClient) {
throw new Error("CortexDB client not configured. Provide toolkitConfig.apiKey in ChatDooorGenerativeAI constructor.");
}
return globalCortexDBClient;
}
/**
* Get global provider name (set via toolkitConfig)
*/
function getGlobalProviderName() {
return globalProviderName;
}