@ekodb/ekodb-client 0.12.0 → 0.13.0

package/src/client.ts

@@ -4,7 +4,7 @@
 
 import { encode, decode } from "@msgpack/msgpack";
 import { QueryBuilder, Query as QueryBuilderQuery } from "./query-builder";
-import { SearchQuery, SearchQueryBuilder, SearchResponse } from "./search";
+import { SearchQuery, SearchResponse } from "./search";
 import { Schema, SchemaBuilder, CollectionMetadata } from "./schema";
 import { Script, FunctionResult } from "./functions";
 
@@ -46,8 +46,6 @@ export interface ClientConfig {
   shouldRetry?: boolean;
   /** Maximum number of retry attempts (default: 3) */
   maxRetries?: number;
-  /** Request timeout in milliseconds (default: 30000) */
-  timeout?: number;
   /** Serialization format (default: MessagePack for best performance, use Json for debugging) */
   format?: SerializationFormat;
 }
@@ -134,6 +132,26 @@ export interface BatchDeleteOptions {
   transactionId?: string;
 }
 
+export interface DistinctValuesOptions {
+  /** Optional filter expression (same format as find() filter). */
+  filter?: any;
+  /** Bypass ripple propagation for this query. */
+  bypassRipple?: boolean;
+  /** Bypass cache for this query. */
+  bypassCache?: boolean;
+}
+
+export interface DistinctValuesResponse {
+  /** Collection that was queried. */
+  collection: string;
+  /** Field whose distinct values were returned. */
+  field: string;
+  /** Unique values, sorted alphabetically. */
+  values: any[];
+  /** Number of distinct values. */
+  count: number;
+}
+
 // ========== Chat Interfaces ==========
 
 export interface CollectionConfig {
@@ -298,6 +316,24 @@ export interface EmbedResponse {
   dimensions: number;
 }
 
+/**
+ * Request for stateless raw LLM completion — no session, no history, no RAG.
+ */
+export interface RawCompletionRequest {
+  system_prompt: string;
+  message: string;
+  provider?: string;
+  model?: string;
+  max_tokens?: number;
+}
+
+/**
+ * Response from a raw LLM completion request.
+ */
+export interface RawCompletionResponse {
+  content: string;
+}
+
 /**
  * User function definition - reusable sequence of Functions that can be called by Scripts
  */
@@ -337,7 +373,6 @@ export class EkoDBClient {
   private token: string | null = null;
   private shouldRetry: boolean;
   private maxRetries: number;
-  private timeout: number;
   private format: SerializationFormat;
   private rateLimitInfo: RateLimitInfo | null = null;
 
@@ -348,14 +383,12 @@ export class EkoDBClient {
       this.apiKey = apiKey!;
       this.shouldRetry = true;
       this.maxRetries = 3;
-      this.timeout = 30000;
       this.format = SerializationFormat.MessagePack; // Default to MessagePack for 2-3x performance
     } else {
       this.baseURL = config.baseURL;
       this.apiKey = config.apiKey;
       this.shouldRetry = config.shouldRetry ?? true;
       this.maxRetries = config.maxRetries ?? 3;
-      this.timeout = config.timeout ?? 30000;
       this.format = config.format ?? SerializationFormat.MessagePack; // Default to MessagePack for 2-3x performance
     }
   }
@@ -386,7 +419,7 @@ export class EkoDBClient {
   /**
    * Refresh the authentication token
    */
-  private async refreshToken(): Promise<void> {
+  async refreshToken(): Promise<void> {
     const response = await fetch(`${this.baseURL}/api/auth/token`, {
       method: "POST",
       headers: { "Content-Type": "application/json" },
@@ -410,6 +443,22 @@ export class EkoDBClient {
     this.token = result.token;
   }
 
+  /**
+   * Get the current authentication token.
+   * Returns null if not yet authenticated. Call refreshToken() first.
+   */
+  getToken(): string | null {
+    return this.token;
+  }
+
+  /**
+   * Clear the cached authentication token.
+   * The next request will trigger a fresh token exchange.
+   */
+  clearTokenCache(): void {
+    this.token = null;
+  }
+
   /**
    * Extract rate limit information from response headers
    */
@@ -660,6 +709,32 @@ export class EkoDBClient {
     return this.makeRequest<Record>("GET", `/api/find/${collection}/${id}`);
   }
 
+  /**
+   * Find a document by ID with field projection
+   * @param collection - Collection name
+   * @param id - Document ID
+   * @param selectFields - Fields to include in the result
+   * @param excludeFields - Fields to exclude from the result
+   */
+  async findByIdWithProjection(
+    collection: string,
+    id: string,
+    selectFields?: string[],
+    excludeFields?: string[],
+  ): Promise<Record> {
+    const params = new URLSearchParams();
+    if (selectFields?.length) {
+      params.append("select_fields", selectFields.join(","));
+    }
+    if (excludeFields?.length) {
+      params.append("exclude_fields", excludeFields.join(","));
+    }
+    const url = params.toString()
+      ? `/api/find/${collection}/${id}?${params.toString()}`
+      : `/api/find/${collection}/${id}`;
+    return this.makeRequest<Record>("GET", url);
+  }
+
   /**
    * Update a document
    * @param collection - Collection name
@@ -688,6 +763,52 @@ export class EkoDBClient {
     return this.makeRequest<Record>("PUT", url, record);
   }
 
+  /**
+   * Apply an atomic field action to a single field of a record.
+   *
+   * Use this instead of `update()` for safe concurrent modifications like
+   * incrementing counters, pushing to arrays, or arithmetic operations.
+   *
+   * @param collection - Collection name
+   * @param id - Record ID
+   * @param action - The atomic action: increment, decrement, multiply, divide, modulo,
+   *                 push, pop, shift, unshift, remove, append, clear
+   * @param field - The field name to apply the action to
+   * @param value - The value for the action (omit for pop/shift/clear)
+   */
+  async updateWithAction(
+    collection: string,
+    id: string,
+    action: string,
+    field: string,
+    value?: any,
+  ): Promise<Record> {
+    const url = `/api/update/${collection}/${id}/action/${action}`;
+    return this.makeRequest<Record>("PUT", url, {
+      field,
+      value: value ?? null,
+    });
+  }
+
+  /**
+   * Apply a sequence of atomic field actions to a record in a single request.
+   *
+   * All actions are applied atomically — the record is fetched once, all actions
+   * run in order, and the result is persisted in a single update.
+   *
+   * @param collection - Collection name
+   * @param id - Record ID
+   * @param actions - Array of [action, field, value] tuples
+   */
+  async updateWithActionSequence(
+    collection: string,
+    id: string,
+    actions: [string, string, any][],
+  ): Promise<Record> {
+    const url = `/api/update/sequence/${collection}/${id}`;
+    return this.makeRequest<Record>("PUT", url, actions);
+  }
+
   /**
    * Delete a document
    * @param collection - Collection name
@@ -1313,6 +1434,51 @@ export class EkoDBClient {
     );
   }
 
+  /**
+   * Get distinct (unique) values for a field across all records in a collection.
+   *
+   * Results are deduplicated and sorted alphabetically. Supports an optional filter
+   * to restrict which records are examined.
+   *
+   * @param collection - Collection name
+   * @param field - Field to get distinct values for
+   * @param options - Optional filter and bypass flags
+   *
+   * @example
+   * // All distinct statuses
+   * const resp = await client.distinctValues("orders", "status");
+   * console.log(resp.values); // ["active", "cancelled", "shipped"]
+   *
+   * // Only statuses for US orders
+   * const resp = await client.distinctValues("orders", "status", {
+   *   filter: { type: "Condition", content: { field: "region", operator: "Eq", value: "us" } }
+   * });
+   */
+  async distinctValues(
+    collection: string,
+    field: string,
+    options: DistinctValuesOptions = {},
+  ): Promise<DistinctValuesResponse> {
+    const body: {
+      filter?: any;
+      bypass_ripple?: boolean;
+      bypass_cache?: boolean;
+    } = {};
+    if (options.filter !== undefined) body.filter = options.filter;
+    if (options.bypassRipple !== undefined)
+      body.bypass_ripple = options.bypassRipple;
+    if (options.bypassCache !== undefined)
+      body.bypass_cache = options.bypassCache;
+
+    return this.makeRequest<DistinctValuesResponse>(
+      "POST",
+      `/api/distinct/${collection}/${field}`,
+      body,
+      0,
+      true, // Force JSON
+    );
+  }
+
   /**
    * Health check - verify the ekoDB server is responding
    */
@@ -1348,6 +1514,34 @@ export class EkoDBClient {
     );
   }
 
+  /**
+   * Stateless raw LLM completion — no session, no history, no RAG.
+   *
+   * Sends a system prompt and user message directly to the LLM via ekoDB
+   * and returns the raw text response without any context injection or
+   * conversation management. Use this for structured-output tasks such as
+   * planning where the response must be parsed programmatically.
+   *
+   * @example
+   * const resp = await client.rawCompletion({
+   *   system_prompt: "You are a helpful assistant.",
+   *   message: "Summarize this in JSON.",
+   *   max_tokens: 2048,
+   * });
+   * console.log(resp.content);
+   */
+  async rawCompletion(
+    request: RawCompletionRequest,
+  ): Promise<RawCompletionResponse> {
+    return this.makeRequest<RawCompletionResponse>(
+      "POST",
+      "/api/chat/complete",
+      request,
+      0,
+      true, // Force JSON
+    );
+  }
+
   /**
    * Send a message in an existing chat session
    */
@@ -1560,6 +1754,21 @@ export class EkoDBClient {
     );
   }
 
+  /**
+   * Get all built-in server-side chat tool definitions.
+   * Returns a list of tool objects with name, description, and parameters fields.
+   * Used by planning agents to discover available tools dynamically.
+   */
+  async getChatTools(): Promise<object[]> {
+    return this.makeRequest<object[]>(
+      "GET",
+      "/api/chat/tools",
+      undefined,
+      0,
+      true, // Force JSON
+    );
+  }
+
   /**
    * Get available models for a specific provider
    * @param provider - Provider name (e.g., "openai", "anthropic", "perplexity")
@@ -1916,26 +2125,130 @@ export class EkoDBClient {
   }
 }
 
+/** Mutation notification from a subscription. */
+export interface MutationNotification {
+  collection: string;
+  event: string;
+  recordIds: string[];
+  records?: any;
+  timestamp: string;
+}
+
+/** A chunk/event from a streaming chat response. */
+export type ChatStreamEvent =
+  | { type: "chunk"; content: string }
+  | {
+      type: "end";
+      messageId: string;
+      tokenUsage?: any;
+      toolCallHistory?: any;
+      executionTimeMs: number;
+    }
+  | {
+      type: "toolCall";
+      chatId: string;
+      callId: string;
+      toolName: string;
+      arguments: any;
+    }
+  | { type: "error"; error: string };
+
+/** Definition for a client-side tool the LLM can call. */
+export interface ClientToolDefinition {
+  name: string;
+  description: string;
+  parameters: any;
+}
+
+/** Options for chatSend. */
+export interface ChatSendOptions {
+  bypassRipple?: boolean;
+  clientTools?: ClientToolDefinition[];
+  maxIterations?: number;
+  confirmTools?: string[];
+  excludeTools?: string[];
+}
+
+/** Options for subscribe. */
+export interface SubscribeOptions {
+  filterField?: string;
+  filterValue?: string;
+}
+
+/** EventEmitter-like interface for subscriptions and chat streams. */
+export class EventStream<_T = unknown> {
+  private listeners: Map<string, Array<(data: any) => void>> = new Map();
+  private _closed = false;
+
+  on(event: string, listener: (data: any) => void): this {
+    if (!this.listeners.has(event)) {
+      this.listeners.set(event, []);
+    }
+    this.listeners.get(event)!.push(listener);
+    return this;
+  }
+
+  /** @internal */
+  emit(event: string, data?: any): void {
+    const handlers = this.listeners.get(event);
+    if (handlers) {
+      for (const handler of handlers) {
+        handler(data);
+      }
+    }
+  }
+
+  get closed(): boolean {
+    return this._closed;
+  }
+
+  /** @internal */
+  close(): void {
+    this._closed = true;
+    this.emit("close");
+  }
+}
+
 /**
- * WebSocket client for real-time queries
+ * WebSocket client for real-time queries, subscriptions, and chat streaming.
  */
 export class WebSocketClient {
   private wsURL: string;
   private token: string;
   private ws: any = null;
+  private dispatcherRunning = false;
+
+  // Dispatcher state
+  private pendingRequests: Map<
+    string,
+    { resolve: (value: any) => void; reject: (reason: any) => void }
+  > = new Map();
+  private subscriptions: Map<string, EventStream<MutationNotification>> =
+    new Map();
+  private chatStreams: Map<string, EventStream<ChatStreamEvent>> = new Map();
+  private registerToolsAck: {
+    resolve: (value: any) => void;
+    reject: (reason: any) => void;
+  } | null = null;
 
   constructor(wsURL: string, token: string) {
     this.wsURL = wsURL;
     this.token = token;
   }
 
+  private messageCounter = 0;
+
+  private genMessageId(): string {
+    const counter = this.messageCounter++;
+    return `${Date.now()}-${counter}-${Math.random().toString(36).slice(2, 8)}`;
+  }
+
   /**
-   * Connect to WebSocket
+   * Connect and start the dispatcher.
    */
-  private async connect(): Promise<void> {
-    if (this.ws) return;
+  private async ensureConnected(): Promise<void> {
+    if (this.ws && this.dispatcherRunning) return;
 
-    // Dynamic import for Node.js WebSocket
     const WebSocket = (await import("ws")).default;
 
     let url = this.wsURL;
@@ -1949,49 +2262,318 @@ export class WebSocketClient {
       },
     });
 
-    return new Promise((resolve, reject) => {
+    await new Promise<void>((resolve, reject) => {
       this.ws.on("open", () => resolve());
       this.ws.on("error", (err: Error) => reject(err));
     });
+
+    this.spawnDispatcher();
+  }
+
+  private spawnDispatcher(): void {
+    if (this.dispatcherRunning) return;
+    this.dispatcherRunning = true;
+
+    this.ws.on("message", (data: Buffer) => {
+      try {
+        const msg = JSON.parse(data.toString());
+        this.routeMessage(msg);
+      } catch {
+        // Ignore malformed messages
+      }
+    });
+
+    this.ws.on("close", () => {
+      this.dispatcherRunning = false;
+      // Notify all pending requests
+      for (const [, pending] of this.pendingRequests) {
+        pending.reject(new Error("WebSocket connection closed"));
+      }
+      this.pendingRequests.clear();
+      // Close all chat streams
+      for (const [, stream] of this.chatStreams) {
+        stream.emit("event", { type: "error", error: "Connection closed" });
+        stream.close();
+      }
+      this.chatStreams.clear();
+      // Close all subscriptions
+      for (const [, stream] of this.subscriptions) {
+        stream.close();
+      }
+      this.subscriptions.clear();
+      this.ws = null;
+    });
+  }
+
+  private routeMessage(msg: any): void {
+    switch (msg.type) {
+      case "Success":
+      case "Error": {
+        const messageId = msg.payload?.message_id || msg.payload?.messageId;
+        if (messageId && this.pendingRequests.has(messageId)) {
+          const pending = this.pendingRequests.get(messageId)!;
+          this.pendingRequests.delete(messageId);
+          if (msg.type === "Error") {
+            pending.reject(new Error(msg.message || "Unknown error"));
+          } else {
+            pending.resolve(msg.payload);
+          }
+        } else if (this.registerToolsAck) {
+          const ack = this.registerToolsAck;
+          this.registerToolsAck = null;
+          if (msg.type === "Error") {
+            ack.reject(new Error(msg.message || "Tool registration failed"));
+          } else {
+            ack.resolve(msg.payload);
+          }
+        }
+        break;
+      }
+
+      case "MutationNotification": {
+        const payload = msg.payload;
+        const notification: MutationNotification = {
+          collection: payload.collection,
+          event: payload.event,
+          recordIds: payload.record_ids || payload.recordIds || [],
+          records: payload.records,
+          timestamp: payload.timestamp,
+        };
+        for (const [collection, stream] of this.subscriptions) {
+          if (collection === notification.collection) {
+            stream.emit("mutation", notification);
+          }
+        }
+        break;
+      }
+
+      case "ChatStreamChunk": {
+        const chatId = msg.payload?.chat_id || msg.payload?.chatId;
+        const stream = this.chatStreams.get(chatId);
+        if (stream) {
+          stream.emit("event", {
+            type: "chunk",
+            content: msg.payload.content,
+          } as ChatStreamEvent);
+        }
+        break;
+      }
+
+      case "ChatStreamEnd": {
+        const chatId = msg.payload?.chat_id || msg.payload?.chatId;
+        const stream = this.chatStreams.get(chatId);
+        if (stream) {
+          stream.emit("event", {
+            type: "end",
+            messageId: msg.payload.message_id || msg.payload.messageId || "",
+            tokenUsage: msg.payload.token_usage || msg.payload.tokenUsage,
+            toolCallHistory:
+              msg.payload.tool_call_history || msg.payload.toolCallHistory,
+            executionTimeMs:
+              msg.payload.execution_time_ms || msg.payload.executionTimeMs || 0,
+          } as ChatStreamEvent);
+          this.chatStreams.delete(chatId);
+          stream.close();
+        }
+        break;
+      }
+
+      case "ChatStreamError": {
+        const chatId = msg.payload?.chat_id || msg.payload?.chatId;
+        const stream = this.chatStreams.get(chatId);
+        if (stream) {
+          stream.emit("event", {
+            type: "error",
+            error: msg.payload.error || msg.payload.message || "Unknown error",
+          } as ChatStreamEvent);
+          this.chatStreams.delete(chatId);
+          stream.close();
+        }
+        break;
+      }
+
+      case "ClientToolCall": {
+        const chatId = msg.payload?.chat_id || msg.payload?.chatId;
+        const stream = this.chatStreams.get(chatId);
+        if (stream) {
+          stream.emit("event", {
+            type: "toolCall",
+            chatId,
+            callId: msg.payload.call_id || msg.payload.callId,
+            toolName: msg.payload.tool_name || msg.payload.toolName,
+            arguments: msg.payload.arguments,
+          } as ChatStreamEvent);
+        }
+        break;
+      }
+    }
+  }
+
+  private async sendRequest(request: any): Promise<any> {
+    await this.ensureConnected();
+    const messageId = request.messageId || request.message_id;
+
+    return new Promise((resolve, reject) => {
+      this.pendingRequests.set(messageId, { resolve, reject });
+      try {
+        this.ws.send(JSON.stringify(request));
+      } catch (err) {
+        this.pendingRequests.delete(messageId);
+        reject(err);
+      }
+    });
   }
 
   /**
-   * Find all records in a collection via WebSocket
+   * Find all records in a collection via WebSocket.
    */
   async findAll(collection: string): Promise<Record[]> {
-    await this.connect();
-
-    const messageId = Date.now().toString();
-    const request = {
+    const messageId = this.genMessageId();
+    const payload = await this.sendRequest({
       type: "FindAll",
       messageId,
       payload: { collection },
+    });
+    return payload?.data || [];
+  }
+
+  /**
+   * Subscribe to mutation notifications on a collection.
+   * Returns an EventStream that emits "mutation" events.
+   */
+  async subscribe(
+    collection: string,
+    options?: SubscribeOptions,
+  ): Promise<EventStream<MutationNotification>> {
+    await this.ensureConnected();
+
+    if (this.subscriptions.has(collection)) {
+      throw new Error(`Already subscribed to collection "${collection}"`);
+    }
+
+    const messageId = this.genMessageId();
+    const stream = new EventStream<MutationNotification>();
+    this.subscriptions.set(collection, stream);
+
+    const request: any = {
+      type: "Subscribe",
+      messageId,
+      payload: {
+        collection,
+        ...(options?.filterField && { filter_field: options.filterField }),
+        ...(options?.filterValue && { filter_value: options.filterValue }),
+      },
     };
 
-    return new Promise((resolve, reject) => {
-      this.ws.send(JSON.stringify(request));
+    // Send subscribe request and wait for ack
+    try {
+      await this.sendRequest(request);
+    } catch (err) {
+      this.subscriptions.delete(collection);
+      throw err;
+    }
+    return stream;
+  }
 
-      this.ws.once("message", (data: Buffer) => {
-        const response = JSON.parse(data.toString());
+  /**
+   * Send a chat message and receive a streaming response.
+   * Returns an EventStream that emits "event" with ChatStreamEvent objects.
+   */
+  async chatSend(
+    chatId: string,
+    message: string,
+    options?: ChatSendOptions,
+  ): Promise<EventStream<ChatStreamEvent>> {
+    await this.ensureConnected();
 
-        if (response.type === "Error") {
-          reject(new Error(response.message));
-        } else {
-          resolve(response.payload?.data || []);
-        }
-      });
+    if (this.chatStreams.has(chatId)) {
+      throw new Error(`Chat stream already active for chatId "${chatId}"`);
+    }
+
+    const stream = new EventStream<ChatStreamEvent>();
+    this.chatStreams.set(chatId, stream);
+
+    const request: any = {
+      type: "ChatSend",
+      payload: {
+        chat_id: chatId,
+        message,
+        ...(options?.bypassRipple != null && {
+          bypass_ripple: options.bypassRipple,
+        }),
+        ...(options?.clientTools && { client_tools: options.clientTools }),
+        ...(options?.maxIterations != null && {
+          max_iterations: options.maxIterations,
+        }),
+        ...(options?.confirmTools && { confirm_tools: options.confirmTools }),
+        ...(options?.excludeTools && { exclude_tools: options.excludeTools }),
+      },
+    };
+
+    this.ws.send(JSON.stringify(request));
+    return stream;
+  }
 
-      this.ws.once("error", reject);
+  /**
+   * Register client-side tools for a chat session.
+   */
+  async registerClientTools(
+    chatId: string,
+    tools: ClientToolDefinition[],
+  ): Promise<void> {
+    await this.ensureConnected();
+
+    const request = {
+      type: "RegisterClientTools",
+      payload: {
+        chat_id: chatId,
+        tools,
+      },
+    };
+
+    await new Promise<void>((resolve, reject) => {
+      this.registerToolsAck = {
+        resolve: () => resolve(),
+        reject: (err) => reject(err),
+      };
+      this.ws.send(JSON.stringify(request));
     });
   }
 
   /**
-   * Close the WebSocket connection
+   * Send a tool result back to the server during a chat stream.
+   */
+  async sendToolResult(
+    chatId: string,
+    callId: string,
+    success: boolean,
+    result?: any,
+    error?: string,
+  ): Promise<void> {
+    await this.ensureConnected();
+
+    const request = {
+      type: "ClientToolResult",
+      payload: {
+        chat_id: chatId,
+        call_id: callId,
+        success,
+        ...(result !== undefined && { result }),
+        ...(error !== undefined && { error }),
+      },
+    };
+
+    this.ws.send(JSON.stringify(request));
+  }
+
+  /**
+   * Close the WebSocket connection.
    */
   close(): void {
     if (this.ws) {
       this.ws.close();
       this.ws = null;
+      this.dispatcherRunning = false;
     }
   }
 }