jazz-tools 0.18.25 → 0.18.26

package/src/tools/coValues/coVector.ts

@@ -0,0 +1,432 @@
+import type { RawBinaryCoStream } from "cojson";
+import { cojsonInternals } from "cojson";
+import {
+  AnonymousJazzAgent,
+  CoValue,
+  CoValueClass,
+  getCoValueOwner,
+  Group,
+  ID,
+  Resolved,
+  SubscribeListenerOptions,
+  SubscribeRestArgs,
+  TypeSym,
+} from "../internal.js";
+import {
+  Account,
+  CoValueJazzApi,
+  inspect,
+  loadCoValueWithoutMe,
+  parseCoValueCreateOptions,
+  parseSubscribeRestArgs,
+  subscribeToCoValueWithoutMe,
+  subscribeToExistingCoValue,
+} from "../internal.js";
+
+/**
+ * CoVectors are collaborative storages of vectors (floating point arrays).
+ *
+ * @category CoValues
+ */
+export class CoVector
+  extends Float32Array
+  implements Readonly<Float32Array>, CoValue
+{
+  declare $jazz: CoVectorJazzApi<this>;
+
+  /** @category Type Helpers */
+  declare [TypeSym]: "BinaryCoStream";
+
+  static get [Symbol.species]() {
+    return Float32Array;
+  }
+
+  protected static requiredDimensionsCount: number | undefined = undefined;
+  private declare _isVectorLoaded: boolean;
+  private declare _requiredDimensionsCount: number;
+
+  constructor(
+    options:
+      | {
+          owner: Account | Group;
+        }
+      | {
+          fromRaw: RawBinaryCoStream;
+        },
+  ) {
+    const dimensionsCount = (new.target as typeof CoVector)
+      .requiredDimensionsCount;
+
+    if (dimensionsCount === undefined) {
+      throw new Error(
+        "Instantiating CoVector without a dimensions count is not allowed. Use co.vector(...).create() instead.",
+      );
+    }
+
+    // Initialize empty Float32Array buffer with the expected vector length
+    // to be filled with the vector data later
+    super(dimensionsCount);
+
+    const isFromRaw = "fromRaw" in options;
+
+    const raw: RawBinaryCoStream = isFromRaw
+      ? options.fromRaw
+      : options.owner.$jazz.raw.createBinaryStream();
+
+    Object.defineProperties(this, {
+      [TypeSym]: { value: "BinaryCoStream", enumerable: false },
+      $jazz: {
+        value: new CoVectorJazzApi(this, raw),
+        enumerable: false,
+      },
+      _isVectorLoaded: { value: false, enumerable: false, writable: true },
+      _requiredDimensionsCount: {
+        value: dimensionsCount,
+        enumerable: false,
+        writable: false,
+      },
+    });
+
+    if (isFromRaw) {
+      this.loadVectorData();
+    }
+  }
+
+  /** @category Internals */
+  static fromRaw<V extends CoVector>(
+    this: CoValueClass<V> & typeof CoVector,
+    raw: RawBinaryCoStream,
+  ) {
+    return new this({ fromRaw: raw });
+  }
+
+  /**
+   * Create a new `CoVector` instance with the given vector.
+   *
+   * @category Creation
+   * @deprecated Use `co.vector(...).create` instead.
+   */
+  static create<S extends CoVector>(
+    this: CoValueClass<S> & typeof CoVector,
+    vector: number[] | Float32Array,
+    options?: { owner?: Account | Group } | Account | Group,
+  ) {
+    const vectorAsFloat32Array =
+      vector instanceof Float32Array ? vector : new Float32Array(vector);
+
+    const givenVectorDimensions =
+      vectorAsFloat32Array.byteLength / vectorAsFloat32Array.BYTES_PER_ELEMENT;
+
+    if (
+      this.requiredDimensionsCount !== undefined &&
+      givenVectorDimensions !== this.requiredDimensionsCount
+    ) {
+      throw new Error(
+        `Vector dimension mismatch! Expected ${this.requiredDimensionsCount} dimensions, got ${
+          givenVectorDimensions
+        }`,
+      );
+    }
+
+    const coVector = new this(parseCoValueCreateOptions(options));
+    coVector.setVectorData(vectorAsFloat32Array);
+
+    const byteArray = CoVector.toByteArray(vectorAsFloat32Array);
+
+    coVector.$jazz.raw.startBinaryStream({
+      mimeType: "application/vector+octet-stream",
+      totalSizeBytes: byteArray.byteLength,
+    });
+
+    const chunkSize =
+      cojsonInternals.TRANSACTION_CONFIG.MAX_RECOMMENDED_TX_SIZE;
+
+    // Although most embedding vectors are small
+    // (3072-dimensional vector is only 12,288 bytes),
+    // we should still chunk the data to avoid transaction size limits
+    for (let idx = 0; idx < byteArray.length; idx += chunkSize) {
+      coVector.$jazz.raw.pushBinaryStreamChunk(
+        byteArray.slice(idx, idx + chunkSize),
+      );
+    }
+    coVector.$jazz.raw.endBinaryStream();
+
+    return coVector;
+  }
+
+  private static toByteArray(vector: Float32Array): Uint8Array {
+    // zero copy view of the vector bytes
+    return new Uint8Array(vector.buffer, vector.byteOffset, vector.byteLength);
+  }
+
+  private static fromByteArray(bytesChunks: Uint8Array[]): Float32Array {
+    const total = bytesChunks.reduce((acc, c) => acc + c.byteLength, 0);
+
+    if (total % 4 !== 0)
+      throw new Error("[INTERNAL] Total byte length must be multiple of 4");
+
+    const u8 = new Uint8Array(total);
+    let off = 0;
+
+    for (const c of bytesChunks) {
+      u8.set(c, off);
+      off += c.byteLength;
+    }
+
+    return new Float32Array(u8.buffer, u8.byteOffset, total / 4);
+  }
+
+  private loadVectorData(): void {
+    if (this._isVectorLoaded === true) {
+      return;
+    }
+
+    const chunks = this.$jazz.raw.getBinaryChunks();
+
+    if (!chunks) {
+      // This should never happen
+      throw new Error(`CoVector '${this.$jazz.raw.id}' is not loaded`);
+    }
+
+    const vector = CoVector.fromByteArray(chunks.chunks);
+
+    if (vector.length !== this._requiredDimensionsCount) {
+      throw new Error(
+        `Vector dimension mismatch! CoVector '${this.$jazz.raw.id}' loaded with ${vector.length} dimensions, but the schema requires ${this._requiredDimensionsCount} dimensions`,
+      );
+    }
+
+    this.setVectorData(vector);
+
+    return;
+  }
+
+  private setVectorData(vector: Float32Array): void {
+    super.set(vector, 0);
+    this._isVectorLoaded = true;
+  }
+
+  /**
+   * Get a JSON representation of the `CoVector`
+   * @category Content
+   */
+  toJSON(): Array<number> {
+    return Array.from(this);
+  }
+
+  /** @internal */
+  [inspect]() {
+    return this.toJSON();
+  }
+
+  /**
+   * Load a `CoVector`
+   *
+   * @category Subscription & Loading
+   * @deprecated Use `co.vector(...).load` instead.
+   */
+  static async load<C extends CoVector>(
+    this: CoValueClass<C>,
+    id: ID<C>,
+    options?: {
+      loadAs?: Account | AnonymousJazzAgent;
+    },
+  ): Promise<C | null> {
+    const coVector = await loadCoValueWithoutMe(this, id, options);
+
+    /**
+     * We are only interested in the entire vector. Since most vectors are small (<15kB),
+     * we can wait for the stream to be complete before returning the vector
+     */
+    if (!coVector?.$jazz.raw.isBinaryStreamEnded()) {
+      return new Promise<C | null>((resolve) => {
+        subscribeToCoValueWithoutMe(
+          this,
+          id,
+          options || {},
+          (value, unsubscribe) => {
+            if (value.$jazz.raw.isBinaryStreamEnded()) {
+              unsubscribe();
+              resolve(value);
+            }
+          },
+        );
+      });
+    }
+
+    coVector.loadVectorData();
+    return coVector;
+  }
+
+  /**
+   * Subscribe to a `CoVector`, when you have an ID but don't have a `CoVector` instance yet
+   * @category Subscription & Loading
+   * @deprecated Use `co.vector(...).subscribe` instead.
+   */
+  static subscribe<V extends CoVector>(
+    this: CoValueClass<V>,
+    id: ID<V>,
+    listener: (value: Resolved<V, true>, unsubscribe: () => void) => void,
+  ): () => void;
+  static subscribe<V extends CoVector>(
+    this: CoValueClass<V>,
+    id: ID<V>,
+    options: SubscribeListenerOptions<V, true>,
+    listener: (value: Resolved<V, true>, unsubscribe: () => void) => void,
+  ): () => void;
+  static subscribe<V extends CoVector>(
+    this: CoValueClass<V>,
+    id: ID<V>,
+    ...args: SubscribeRestArgs<V, true>
+  ): () => void {
+    const { options, listener } = parseSubscribeRestArgs(args);
+    return subscribeToCoValueWithoutMe<V, true>(this, id, options, listener);
+  }
+
+  // CoVector mutation method overrides, as CoVectors aren't meant to be mutated
+  /**
+   * Calling `copyWithin` on a CoVector is forbidden. CoVectors are immutable.
+   * @deprecated If you want to change the vector, replace the former instance of CoVector with a new one.
+   */
+  override copyWithin(target: number, start: number, end?: number): never {
+    throw new Error("Cannot mutate a CoVector using `copyWithin`");
+  }
+  /**
+   * Calling `fill` on a CoVector is forbidden. CoVectors are immutable.
+   * @deprecated If you want to change the vector, replace the former instance of CoVector with a new one.
+   */
+  override fill(value: number, start?: number, end?: number): never {
+    throw new Error("Cannot mutate a CoVector using `fill`");
+  }
+  /**
+   * Calling `reverse` on a CoVector is forbidden. CoVectors are immutable.
+   * @deprecated If you want to change the vector, replace the former instance of CoVector with a new one.
+   */
+  override reverse(): never {
+    throw new Error("Cannot mutate a CoVector using `reverse`");
+  }
+  /**
+   * Calling `set` on a CoVector is forbidden. CoVectors are immutable.
+   * @deprecated If you want to change the vector, replace the former instance of CoVector with a new one.
+  //  */
+  override set(array: ArrayLike<number>, offset?: number): never {
+    throw new Error("Cannot mutate a CoVector using `set`");
+  }
+  /**
+   * Calling `sort` on a CoVector is forbidden. CoVectors are immutable.
+   * @deprecated If you want to change the vector, replace the former instance of CoVector with a new one.
+   */
+  override sort(compareFn?: (a: number, b: number) => number): never {
+    throw new Error("Cannot mutate a CoVector using `sort`");
+  }
+}
+
+export class CoVectorJazzApi<V extends CoVector> extends CoValueJazzApi<V> {
+  constructor(
+    private coVector: V,
+    public raw: RawBinaryCoStream,
+  ) {
+    super(coVector);
+  }
+
+  get owner(): Group {
+    return getCoValueOwner(this.coVector);
+  }
+
+  /**
+   * An instance method to subscribe to an existing `CoVector`
+   * @category Subscription & Loading
+   */
+  subscribe<B extends CoVector>(
+    this: CoVectorJazzApi<B>,
+    listener: (value: Resolved<B, true>) => void,
+  ): () => void {
+    return subscribeToExistingCoValue(this.coVector, {}, listener);
+  }
+
+  /**
+   * Wait for the `CoVector` to be uploaded to the other peers.
+   *
+   * @category Subscription & Loading
+   */
+  waitForSync(options?: { timeout?: number }) {
+    return this.raw.core.waitForSync(options);
+  }
+
+  // Vector operations
+  /**
+   * Calculate the magnitude of this vector.
+   */
+  magnitude(): number {
+    return VectorCalculation.magnitude(this.coVector);
+  }
+
+  /**
+   * Normalize this vector.
+   * @returns A new instance of a normalized vector.
+   */
+  normalize(): Float32Array {
+    return VectorCalculation.normalize(this.coVector);
+  }
+
+  /**
+   * Calculate the dot product of this vector and another vector.
+   */
+  dotProduct(otherVector: CoVector | Float32Array | number[]): number {
+    return VectorCalculation.dotProduct(this.coVector, otherVector);
+  }
+
+  /**
+   * Calculate the cosine similarity between this vector and another vector.
+   *
+   * @returns A value between `-1` and `1`:
+   * - `1` means the vectors are identical
+   * - `0` means the vectors are orthogonal (i.e. no similarity)
+   * - `-1` means the vectors are opposite direction (perfectly dissimilar)
+   */
+  cosineSimilarity(otherVector: CoVector | Float32Array | number[]): number {
+    return VectorCalculation.cosineSimilarity(this.coVector, otherVector);
+  }
+}
+
+const VectorCalculation = {
+  magnitude: (vector: Float32Array | number[]) => {
+    let sum = 0;
+    for (const v of vector) {
+      sum += v * v;
+    }
+    return Math.sqrt(sum);
+  },
+  normalize: (vector: Float32Array) => {
+    const mag = VectorCalculation.magnitude(vector);
+
+    if (mag === 0) {
+      return new Float32Array(vector.length).fill(0);
+    }
+
+    return vector.map((v) => v / mag);
+  },
+  dotProduct: (vectorA: Float32Array, vectorB: Float32Array | number[]) => {
+    if (vectorA.length !== vectorB.length) {
+      throw new Error(
+        `Vector dimensions don't match: ${vectorA.length} vs ${vectorB.length}`,
+      );
+    }
+
+    return vectorA.reduce((sum, a, i) => sum + a * vectorB[i]!, 0);
+  },
+  cosineSimilarity: (
+    vectorA: Float32Array,
+    vectorB: Float32Array | number[],
+  ) => {
+    const magnitudeA = VectorCalculation.magnitude(vectorA);
+    const magnitudeB = VectorCalculation.magnitude(vectorB);
+
+    if (magnitudeA === 0 || magnitudeB === 0) {
+      return 0;
+    }
+
+    const dotProductAB = VectorCalculation.dotProduct(vectorA, vectorB);
+    return dotProductAB / (magnitudeA * magnitudeB);
+  },
+};

package/src/tools/implementation/zodSchema/coExport.ts

@@ -7,6 +7,7 @@ export { CoFeedSchema as Feed } from "./schemaTypes/CoFeedSchema.js";
 export { PlainTextSchema as PlainText } from "./schemaTypes/PlainTextSchema.js";
 export { RichTextSchema as RichText } from "./schemaTypes/RichTextSchema.js";
 export { FileStreamSchema as FileStream } from "./schemaTypes/FileStreamSchema.js";
+export { CoVectorSchema as Vector } from "./schemaTypes/CoVectorSchema.js";
 export { CoInput as input } from "./typeConverters/CoFieldSchemaInit.js";
 export {
   AccountSchema as Account,
@@ -23,6 +24,7 @@ export {
   coPlainTextDefiner as plainText,
   coRichTextDefiner as richText,
   coFileStreamDefiner as fileStream,
+  coVectorDefiner as vector,
   coImageDefiner as image,
   coAccountDefiner as account,
   coGroupDefiner as group,

package/src/tools/implementation/zodSchema/runtimeConverters/coValueSchemaTransformation.ts

@@ -12,12 +12,14 @@ import {
   CoValueClass,
   FileStream,
   FileStreamSchema,
+  CoVectorSchema,
   PlainTextSchema,
   SchemaUnion,
   enrichAccountSchema,
   enrichCoMapSchema,
   isCoValueClass,
   Group,
+  CoVector,
 } from "../../../internal.js";
 import { coField } from "../../schema.js";
 
@@ -123,6 +125,17 @@ export function hydrateCoreCoValueSchema<S extends AnyCoreCoValueSchema>(
   } else if (schema.builtin === "FileStream") {
     const coValueClass = FileStream;
     return new FileStreamSchema(coValueClass) as CoValueSchemaFromCoreSchema<S>;
+  } else if (schema.builtin === "CoVector") {
+    const dimensions = schema.dimensions;
+
+    const coValueClass = class CoVectorWithDimensions extends CoVector {
+      protected static requiredDimensionsCount = dimensions;
+    };
+
+    return new CoVectorSchema(
+      dimensions,
+      coValueClass,
+    ) as CoValueSchemaFromCoreSchema<S>;
   } else if (schema.builtin === "CoPlainText") {
     const coValueClass = CoPlainText;
     return new PlainTextSchema(coValueClass) as CoValueSchemaFromCoreSchema<S>;

package/src/tools/implementation/zodSchema/schemaTypes/CoVectorSchema.ts

@@ -0,0 +1,105 @@
+import {
+  Account,
+  AnonymousJazzAgent,
+  CoVector,
+  Group,
+  InstanceOrPrimitiveOfSchema,
+  InstanceOrPrimitiveOfSchemaCoValuesNullable,
+  coOptionalDefiner,
+} from "../../../internal.js";
+import { CoOptionalSchema } from "./CoOptionalSchema.js";
+import { CoreCoValueSchema } from "./CoValueSchema.js";
+
+export interface CoreCoVectorSchema extends CoreCoValueSchema {
+  builtin: "CoVector";
+  dimensions: number;
+}
+
+export function createCoreCoVectorSchema(
+  dimensions: number,
+): CoreCoVectorSchema {
+  return {
+    collaborative: true as const,
+    builtin: "CoVector" as const,
+    dimensions,
+  };
+}
+
+export class CoVectorSchema implements CoreCoVectorSchema {
+  readonly collaborative = true as const;
+  readonly builtin = "CoVector" as const;
+
+  constructor(
+    public dimensions: number,
+    private coValueClass: typeof CoVector,
+  ) {}
+
+  /**
+   * Create a `CoVector` from a given vector.
+   */
+  create(
+    vector: number[] | Float32Array,
+    options?: { owner: Group } | Group,
+  ): CoVectorInstance;
+  /**
+   * Create a `CoVector` from a given vector.
+   *
+   * @deprecated Creating CoValues with an Account as owner is deprecated. Use a Group instead.
+   */
+  create(
+    vector: number[] | Float32Array,
+    options?: { owner: Account | Group } | Account | Group,
+  ): CoVectorInstance;
+  create(
+    vector: number[] | Float32Array,
+    options?: { owner: Account | Group } | Account | Group,
+  ): CoVectorInstance {
+    return this.coValueClass.create(vector, options);
+  }
+
+  /**
+   * Load a `CoVector` with a given ID.
+   */
+  load(
+    id: string,
+    options?: { loadAs: Account | AnonymousJazzAgent },
+  ): Promise<CoVectorInstanceNullable> {
+    return this.coValueClass.load(id, options);
+  }
+
+  /**
+   * Subscribe to a `CoVector`, when you have an ID but don't have a `CoVector` instance yet
+   */
+  subscribe(
+    id: string,
+    options: { loadAs: Account | AnonymousJazzAgent },
+    listener: (
+      value: CoVectorInstanceNullable,
+      unsubscribe: () => void,
+    ) => void,
+  ): () => void;
+  subscribe(
+    id: string,
+    listener: (
+      value: CoVectorInstanceNullable,
+      unsubscribe: () => void,
+    ) => void,
+  ): () => void;
+  subscribe(...args: [any, ...any[]]) {
+    // @ts-expect-error
+    return this.coValueClass.subscribe(...args);
+  }
+
+  getCoValueClass(): typeof CoVector {
+    return this.coValueClass;
+  }
+
+  optional(): CoOptionalSchema<this> {
+    return coOptionalDefiner(this);
+  }
+}
+
+export type CoVectorInstance = InstanceOrPrimitiveOfSchema<CoVectorSchema>;
+
+export type CoVectorInstanceNullable =
+  InstanceOrPrimitiveOfSchemaCoValuesNullable<CoVectorSchema>;

package/src/tools/implementation/zodSchema/typeConverters/InstanceOfSchema.ts

@@ -7,9 +7,11 @@ import {
   CoMap,
   CoPlainText,
   CoRichText,
+  CoVector,
   CoValueClass,
   CoreAccountSchema,
   CoreCoRecordSchema,
+  CoreCoVectorSchema,
   FileStream,
   Group,
 } from "../../../internal.js";
@@ -63,11 +65,13 @@ export type InstanceOfSchema<S extends CoValueClass | AnyZodOrCoValueSchema> =
                     ? CoRichText
                     : S extends CoreFileStreamSchema
                       ? FileStream
-                      : S extends CoreCoOptionalSchema<infer T>
-                        ? InstanceOrPrimitiveOfSchema<T> | undefined
-                        : S extends CoDiscriminatedUnionSchema<infer Members>
-                          ? InstanceOrPrimitiveOfSchema<Members[number]>
-                          : never
+                      : S extends CoreCoVectorSchema
+                        ? Readonly<CoVector>
+                        : S extends CoreCoOptionalSchema<infer T>
+                          ? InstanceOrPrimitiveOfSchema<T> | undefined
+                          : S extends CoDiscriminatedUnionSchema<infer Members>
+                            ? InstanceOrPrimitiveOfSchema<Members[number]>
+                            : never
     : S extends CoValueClass
       ? InstanceType<S>
       : never;

package/src/tools/implementation/zodSchema/typeConverters/InstanceOfSchemaCoValuesNullable.ts

@@ -13,6 +13,8 @@ import {
   CoreCoListSchema,
   CoreCoMapSchema,
   CoreCoRecordSchema,
+  CoreCoVectorSchema,
+  CoVector,
   FileStream,
   Group,
 } from "../../../internal.js";
@@ -70,15 +72,19 @@ export type InstanceOfSchemaCoValuesNullable<
                   ? CoRichText | null
                   : S extends CoreFileStreamSchema
                     ? FileStream | null
-                    : S extends CoreCoOptionalSchema<infer Inner>
-                      ?
-                          | InstanceOrPrimitiveOfSchemaCoValuesNullable<Inner>
-                          | undefined
-                      : S extends CoreCoDiscriminatedUnionSchema<infer Members>
-                        ? InstanceOrPrimitiveOfSchemaCoValuesNullable<
-                            Members[number]
-                          >
-                        : never
+                    : S extends CoreCoVectorSchema
+                      ? Readonly<CoVector> | null
+                      : S extends CoreCoOptionalSchema<infer Inner>
+                        ?
+                            | InstanceOrPrimitiveOfSchemaCoValuesNullable<Inner>
+                            | undefined
+                        : S extends CoreCoDiscriminatedUnionSchema<
+                              infer Members
+                            >
+                          ? InstanceOrPrimitiveOfSchemaCoValuesNullable<
+                              Members[number]
+                            >
+                          : never
   : S extends CoValueClass
     ? InstanceType<S> | null
     : never;

package/src/tools/implementation/zodSchema/zodCo.ts

@@ -9,6 +9,7 @@ import {
   type CoRecordSchema,
   type DefaultProfileShape,
   type FileStreamSchema,
+  type CoVectorSchema,
   ImageDefinition,
   type PlainTextSchema,
   type Simplify,
@@ -18,6 +19,7 @@ import {
   createCoreCoMapSchema,
   createCoreCoPlainTextSchema,
   createCoreFileStreamSchema,
+  createCoreCoVectorSchema,
   hydrateCoreCoValueSchema,
   isAnyCoValueSchema,
   isCoValueClass,
@@ -193,6 +195,19 @@ export const coFileStreamDefiner = (): FileStreamSchema => {
   return hydrateCoreCoValueSchema(coreSchema);
 };
 
+export const coVectorDefiner = (dimensions: number): CoVectorSchema => {
+  const isPositiveInteger = Number.isInteger(dimensions) && dimensions > 0;
+
+  if (!isPositiveInteger) {
+    throw new Error(
+      "co.vector() expects the vector dimensions count to be a positive integer",
+    );
+  }
+
+  const coreSchema = createCoreCoVectorSchema(dimensions);
+  return hydrateCoreCoValueSchema(coreSchema);
+};
+
 export const coPlainTextDefiner = (): PlainTextSchema => {
   const coreSchema = createCoreCoPlainTextSchema();
   return hydrateCoreCoValueSchema(coreSchema);