takos-control 1.0.0

package/src/adapters/pgvector-store.ts

@@ -0,0 +1,236 @@
+/**
+ * PostgreSQL + pgvector implementation of the Cloudflare VectorizeIndex
+ * binding interface.
+ *
+ * Requires the `vector` extension to be installed on the database.
+ * See the companion migration file for the schema definition.
+ */
+
+// Minimal Pool interface to avoid depending on @types/pg.
+interface PgPool {
+  query(text: string, values?: unknown[]): Promise<{ rows: Record<string, unknown>[]; rowCount: number | null }>;
+}
+
+// ---------------------------------------------------------------------------
+// Types — mirror the subset of VectorizeIndex used by EmbeddingsService
+// ---------------------------------------------------------------------------
+
+interface VectorizeVector {
+  id: string;
+  values: number[];
+  namespace?: string;
+  metadata?: Record<string, unknown>;
+}
+
+interface VectorizeMatch {
+  id: string;
+  score: number;
+  values?: number[];
+  namespace?: string;
+  metadata?: Record<string, unknown>;
+}
+
+interface VectorizeMatches {
+  matches: VectorizeMatch[];
+  count: number;
+}
+
+interface VectorizeQueryOptions {
+  topK?: number;
+  namespace?: string;
+  returnValues?: boolean;
+  returnMetadata?: boolean | string;
+  filter?: Record<string, unknown>;
+}
+
+interface VectorizeVectorMutation {
+  ids: string[];
+  count: number;
+}
+
+// ---------------------------------------------------------------------------
+// Config
+// ---------------------------------------------------------------------------
+
+export type PgVectorStoreConfig = {
+  pool: PgPool;
+  tableName?: string;
+};
+
+const DEFAULT_TABLE = 'vector_embeddings';
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function toSqlVector(values: number[]): string {
+  return `[${values.join(',')}]`;
+}
+
+/**
+ * Build a WHERE clause fragment from a Vectorize-style metadata filter.
+ *
+ * The filter is a flat record `{ key: value }` which we translate to
+ * `metadata @> '{"key": value}'::jsonb` conditions.
+ */
+function buildFilterClause(
+  filter: Record<string, unknown> | undefined,
+  paramOffset: number,
+): { clause: string; params: unknown[] } {
+  if (!filter || Object.keys(filter).length === 0) {
+    return { clause: '', params: [] };
+  }
+
+  const conditions: string[] = [];
+  const params: unknown[] = [];
+
+  for (const [key, value] of Object.entries(filter)) {
+    params.push(JSON.stringify({ [key]: value }));
+    conditions.push(`metadata @> $${paramOffset + params.length}::jsonb`);
+  }
+
+  return {
+    clause: `AND ${conditions.join(' AND ')}`,
+    params,
+  };
+}
+
+// ---------------------------------------------------------------------------
+// Public API
+// ---------------------------------------------------------------------------
+
+export function createPgVectorStore(config: PgVectorStoreConfig) {
+  const { pool } = config;
+  const table = config.tableName ?? DEFAULT_TABLE;
+
+  return {
+    // -- query() -----------------------------------------------------------
+    async query(
+      vector: number[],
+      options?: VectorizeQueryOptions,
+    ): Promise<VectorizeMatches> {
+      const topK = options?.topK ?? 10;
+      const returnMetadata = options?.returnMetadata !== false && options?.returnMetadata !== 'none';
+      const returnValues = options?.returnValues === true;
+
+      const { clause: filterClause, params: filterParams } = buildFilterClause(
+        options?.filter as Record<string, unknown> | undefined,
+        1, // $1 is the vector
+      );
+
+      const selectCols = [
+        'id',
+        `1 - (embedding <=> $1::vector) AS score`,
+        ...(returnMetadata ? ['metadata'] : []),
+        ...(returnValues ? ['embedding'] : []),
+      ].join(', ');
+
+      const sql = `
+        SELECT ${selectCols}
+        FROM ${table}
+        WHERE 1=1 ${filterClause}
+        ORDER BY embedding <=> $1::vector
+        LIMIT ${topK}
+      `;
+
+      const result = await pool.query(sql, [toSqlVector(vector), ...filterParams]);
+
+      const matches: VectorizeMatch[] = result.rows.map((row: Record<string, unknown>) => ({
+        id: row.id as string,
+        score: row.score as number,
+        ...(returnMetadata && row.metadata ? { metadata: row.metadata as Record<string, unknown> } : {}),
+        ...(returnValues && row.embedding ? { values: parsePgVector(row.embedding as string) } : {}),
+      }));
+
+      return { matches, count: matches.length };
+    },
+
+    // -- upsert() ----------------------------------------------------------
+    async upsert(vectors: VectorizeVector[]): Promise<VectorizeVectorMutation> {
+      if (vectors.length === 0) return { ids: [], count: 0 };
+
+      const ids: string[] = [];
+
+      // Batch upsert using a single multi-row INSERT ... ON CONFLICT.
+      const valuePlaceholders: string[] = [];
+      const params: unknown[] = [];
+
+      for (let i = 0; i < vectors.length; i++) {
+        const v = vectors[i];
+        const offset = i * 3;
+        valuePlaceholders.push(`($${offset + 1}, $${offset + 2}::vector, $${offset + 3}::jsonb)`);
+        params.push(v.id, toSqlVector(v.values), JSON.stringify(v.metadata ?? {}));
+        ids.push(v.id);
+      }
+
+      const sql = `
+        INSERT INTO ${table} (id, embedding, metadata)
+        VALUES ${valuePlaceholders.join(', ')}
+        ON CONFLICT (id) DO UPDATE SET
+          embedding = EXCLUDED.embedding,
+          metadata = EXCLUDED.metadata
+      `;
+
+      await pool.query(sql, params);
+
+      return { ids, count: ids.length };
+    },
+
+    // -- deleteByIds() -----------------------------------------------------
+    async deleteByIds(ids: string[]): Promise<VectorizeVectorMutation> {
+      if (ids.length === 0) return { ids: [], count: 0 };
+
+      const placeholders = ids.map((_, i) => `$${i + 1}`).join(', ');
+      const sql = `DELETE FROM ${table} WHERE id IN (${placeholders})`;
+
+      const result = await pool.query(sql, ids);
+      const deletedCount = result.rowCount ?? 0;
+
+      return { ids: ids.slice(0, deletedCount), count: deletedCount };
+    },
+
+    // -- getByIds() --------------------------------------------------------
+    async getByIds(ids: string[]): Promise<VectorizeVector[]> {
+      if (ids.length === 0) return [];
+
+      const placeholders = ids.map((_, i) => `$${i + 1}`).join(', ');
+      const sql = `SELECT id, embedding, metadata FROM ${table} WHERE id IN (${placeholders})`;
+
+      const result = await pool.query(sql, ids);
+
+      return result.rows.map((row: Record<string, unknown>) => ({
+        id: row.id as string,
+        values: parsePgVector(row.embedding as string),
+        metadata: (row.metadata as Record<string, unknown>) ?? {},
+      }));
+    },
+
+    // -- describe() --------------------------------------------------------
+    async describe() {
+      const sql = `SELECT COUNT(*) as total FROM ${table}`;
+      const result = await pool.query(sql);
+      const count = Number(result.rows[0]?.total ?? 0);
+      return {
+        name: table,
+        config: { dimensions: 0, metric: 'cosine' as const },
+        vectorsCount: count,
+        processedUpToDatetime: new Date().toISOString(),
+        processedUpToMutation: '',
+      };
+    },
+  };
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+/**
+ * Parse pgvector text representation `[0.1,0.2,0.3]` into number[].
+ */
+function parsePgVector(raw: string): number[] {
+  if (typeof raw !== 'string') return [];
+  const inner = raw.replace(/^\[/, '').replace(/\]$/, '');
+  if (!inner) return [];
+  return inner.split(',').map(Number);
+}

package/src/adapters/pubsub-queue.ts

@@ -0,0 +1,127 @@
+import { PubSub, type Message } from '@google-cloud/pubsub';
+import type { Queue } from '../shared/types/bindings.ts';
+import type { ConsumableQueue, LocalQueueName, LocalQueueRecord } from '../local-platform/queue-runtime.ts';
+
+// ---------------------------------------------------------------------------
+// Config
+// ---------------------------------------------------------------------------
+
+export type PubSubQueueConfig = {
+  projectId?: string;
+  topicName: string;
+  keyFilePath?: string;
+  /** Subscription name for pulling messages. Required for receive(). */
+  subscriptionName?: string;
+  /** Logical queue name used by the worker loop for DLQ / retry-limit lookup. */
+  queueName?: LocalQueueName;
+};
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function lazyPubSub(config: PubSubQueueConfig): () => PubSub {
+  let pubsub: PubSub | undefined;
+  return () => {
+    if (!pubsub) {
+      pubsub = new PubSub({
+        ...(config.projectId ? { projectId: config.projectId } : {}),
+        ...(config.keyFilePath ? { keyFilename: config.keyFilePath } : {}),
+      });
+    }
+    return pubsub;
+  };
+}
+
+const RECEIVE_TIMEOUT_MS = 20_000;
+
+// ---------------------------------------------------------------------------
+// Public API
+// ---------------------------------------------------------------------------
+
+/**
+ * Create a Pub/Sub-backed queue that implements both the standard Queue (send)
+ * interface and the ConsumableQueue (receive) interface required by the
+ * worker poll loop.
+ *
+ * `receive()` attaches a temporary message listener on the subscription and
+ * waits up to 20 s.  The message is immediately acknowledged on receipt,
+ * matching the pop semantics of local/Redis queues.
+ * A `subscriptionName` must be provided in config for receive() to work.
+ */
+export function createPubSubQueue<T = unknown>(config: PubSubQueueConfig): Queue<T> & Partial<ConsumableQueue<T>> {
+  const getPubSub = lazyPubSub(config);
+
+  function getTopic() {
+    return getPubSub().topic(config.topicName);
+  }
+
+  const queue = {
+    // -- ConsumableQueue metadata ------------------------------------------
+    ...(config.queueName ? { queueName: config.queueName } : {}),
+
+    // -- Send --------------------------------------------------------------
+    async send(
+      message: T,
+      _options?: unknown,
+    ): Promise<void> {
+      await getTopic().publishMessage({ json: message as Record<string, unknown> });
+    },
+
+    async sendBatch(
+      messages: Iterable<{ body: T }>,
+    ): Promise<void> {
+      const topic = getTopic();
+      for (const msg of messages) {
+        await topic.publishMessage({ json: msg.body as Record<string, unknown> });
+      }
+    },
+
+    // -- Receive (ConsumableQueue) -----------------------------------------
+    async receive(): Promise<LocalQueueRecord<T> | null> {
+      if (!config.subscriptionName) {
+        throw new Error(
+          'PubSub receive() requires a subscriptionName in config. ' +
+          'Set GCP_PUBSUB_{NAME}_SUBSCRIPTION to enable consuming.',
+        );
+      }
+
+      const subscription = getPubSub().subscription(config.subscriptionName);
+
+      // Use a temporary message listener with a timeout.
+      // The subscription's streaming pull will deliver messages to the handler.
+      return new Promise<LocalQueueRecord<T> | null>((resolve) => {
+        let resolved = false;
+
+        const timer = setTimeout(() => {
+          if (resolved) return;
+          resolved = true;
+          subscription.removeListener('message', onMessage);
+          resolve(null);
+        }, RECEIVE_TIMEOUT_MS);
+
+        function onMessage(message: Message) {
+          if (resolved) return;
+          resolved = true;
+          clearTimeout(timer);
+          subscription.removeListener('message', onMessage);
+
+          const body = JSON.parse(message.data?.toString('utf-8') ?? '{}') as T;
+          const deliveryAttempt = message.deliveryAttempt ?? 1;
+
+          // Immediately acknowledge — matches the pop semantics of local queues.
+          message.ack();
+
+          resolve({
+            body,
+            attempts: deliveryAttempt,
+          });
+        }
+
+        subscription.on('message', onMessage);
+      });
+    },
+  };
+
+  return queue as unknown as Queue<T> & Partial<ConsumableQueue<T>>;
+}

package/src/adapters/r2-compat-types.ts

@@ -0,0 +1,54 @@
+import type {
+  R2Object,
+} from '../../shared/types/bindings.ts';
+
+// ---------------------------------------------------------------------------
+// Cloudflare R2 compatibility types
+//
+// These interfaces mirror the Cloudflare Workers R2 shapes so that non-R2
+// object-store adapters (S3, GCS, …) can produce values that satisfy the same
+// contracts without pulling in @cloudflare/workers-types as a dependency.
+// ---------------------------------------------------------------------------
+
+// R2Objects is not re-exported from bindings, define inline to match the Cloudflare shape.
+export interface R2Objects {
+  objects: R2Object[];
+  truncated: boolean;
+  cursor?: string;
+  delimitedPrefixes: string[];
+}
+
+/** Matches the shape of Cloudflare R2Checksums without requiring the abstract class import. */
+export interface R2ChecksumsLike {
+  readonly md5?: ArrayBuffer;
+  readonly sha1?: ArrayBuffer;
+  readonly sha256?: ArrayBuffer;
+  readonly sha384?: ArrayBuffer;
+  readonly sha512?: ArrayBuffer;
+  toJSON(): Record<string, string | undefined>;
+}
+
+/** Matches the shape of Cloudflare R2HTTPMetadata. */
+export interface R2HTTPMetadataLike {
+  contentType?: string;
+  contentLanguage?: string;
+  contentDisposition?: string;
+  contentEncoding?: string;
+  cacheControl?: string;
+  cacheExpiry?: Date;
+}
+
+/** Matches the Cloudflare R2Range type. */
+export interface R2RangeLike {
+  offset: number;
+  length?: number;
+}
+
+/** Build an empty R2Checksums-compatible object. */
+export function emptyChecksums(): R2ChecksumsLike {
+  return {
+    toJSON() {
+      return {};
+    },
+  };
+}