pi-obsidian-memory 0.0.0 → 0.1.0

package/src/services/embeddings/chunks.ts

@@ -0,0 +1,80 @@
+import type { EmbeddingBatchRequest, EmbeddingFailure } from "#src/domain"
+import type { ChunkPolicy, PreparedChunk, PreparedChunks } from "./types.ts"
+
+export function validateEmbeddingRequest(request: EmbeddingBatchRequest): EmbeddingFailure | null {
+  if (request.inputs.length === 0 || new Set(request.inputs.map((input) => input.key)).size !== request.inputs.length) {
+    return { succeeded: false, status: "invalid-input" }
+  }
+  if (request.inputs.some((input) => input.text.trim().length === 0 || input.key.length === 0)) {
+    return { succeeded: false, status: "invalid-input" }
+  }
+  if (request.inputs.length > 32 || request.inputs.reduce((size, input) => size + input.text.length, 0) > 1_000_000) {
+    return { succeeded: false, status: "input-too-long" }
+  }
+  return null
+}
+
+export function prepareChunks(request: EmbeddingBatchRequest, policy: ChunkPolicy): PreparedChunks {
+  const invalid = validateEmbeddingRequest(request)
+  if (invalid) {
+    return invalid
+  }
+  const output: PreparedChunk[] = []
+  for (const [inputIndex, input] of request.inputs.entries()) {
+    if (request.purpose === "query" && policy.measure(policy.prefix + input.text) > policy.maximum) {
+      return { succeeded: false, status: "input-too-long" }
+    }
+    const chunks = chunkText(input.text, policy)
+    if (!chunks || output.length + chunks.length > 256) {
+      return { succeeded: false, status: "input-too-long" }
+    }
+    output.push(...chunks.map((text) => ({ inputIndex, text })))
+  }
+  return { succeeded: true, output }
+}
+
+export function chunkText(text: string, policy: ChunkPolicy): string[] | null {
+  const characters = [...text]
+  const chunks: string[] = []
+  let start = 0
+  while (start < characters.length) {
+    const end = fittingEnd(characters, start, policy)
+    if (end === start || chunks.length >= 256) {
+      return null
+    }
+    chunks.push(policy.prefix + characters.slice(start, end).join(""))
+    if (end === characters.length) {
+      return chunks
+    }
+    start = overlapStart(characters, start, end, policy)
+  }
+  return chunks
+}
+
+function fittingEnd(characters: readonly string[], start: number, policy: ChunkPolicy): number {
+  let low = start
+  let high = characters.length
+  while (low < high) {
+    const middle = Math.ceil((low + high) / 2)
+    if (policy.measure(policy.prefix + characters.slice(start, middle).join("")) <= policy.maximum) {
+      low = middle
+    } else {
+      high = middle - 1
+    }
+  }
+  return low
+}
+
+function overlapStart(characters: readonly string[], start: number, end: number, policy: ChunkPolicy): number {
+  let low = start + 1
+  let high = end
+  while (low < high) {
+    const middle = Math.floor((low + high) / 2)
+    if (policy.measure(characters.slice(middle, end).join("")) > policy.overlap) {
+      low = middle + 1
+    } else {
+      high = middle
+    }
+  }
+  return low
+}

package/src/services/embeddings/configuration.ts

@@ -0,0 +1,8 @@
+export {
+  EmbeddingConfigurationSchema,
+  embeddingFingerprint,
+  embeddingPolicy,
+  googleEmbeddingConfiguration,
+  localEmbeddingConfiguration,
+  openAIEmbeddingConfiguration,
+} from "#src/domain"

package/src/services/embeddings/constants.ts

@@ -0,0 +1,24 @@
+import type { ModelAsset } from "./types.ts"
+
+export const localModelAssets: readonly ModelAsset[] = [
+  {
+    path: "config.json",
+    bytes: 658,
+    sha256: "cb99455288675345e1a4f411438d5d0adbba5fbd3a67ea4fb03c015433b996c1",
+  },
+  {
+    path: "tokenizer_config.json",
+    bytes: 443,
+    sha256: "a1d6bc8734a6f635dc158508bef000f8e2e5a759c7d92f984b2c86e5ff53425b",
+  },
+  {
+    path: "tokenizer.json",
+    bytes: 17_082_730,
+    sha256: "0b44a9d7b51c3c62626640cda0e2c2f70fdacdc25bbbd68038369d14ebdf4c39",
+  },
+  {
+    path: "onnx/model_quantized.onnx",
+    bytes: 118_308_185,
+    sha256: "f80102d3f2a1229f387d3c81909990d8945513e347b0eab049f7de3c6f98c193",
+  },
+]

package/src/services/embeddings/errors.ts

@@ -0,0 +1,11 @@
+import type { EmbeddingFailure } from "#src/domain"
+
+export class EmbeddingResourceError extends Error {
+  readonly failure: EmbeddingFailure
+
+  constructor(failure: EmbeddingFailure) {
+    super("Embedding resource unavailable")
+    this.name = "EmbeddingResourceError"
+    this.failure = failure
+  }
+}

package/src/services/embeddings/google.ts

@@ -0,0 +1,90 @@
+import { type EmbeddingBatchRequest, type EmbeddingResult, googleEmbeddingConfiguration } from "#src/domain"
+import { finishEmbeddingBatch } from "./batch.ts"
+import { prepareChunks } from "./chunks.ts"
+import { postEmbedding } from "./http.ts"
+import { runEmbeddingOperation } from "./operation.ts"
+import { parseGoogleEmbedding } from "./responses.ts"
+import type { ParsedVectors, PreparedChunk, RemoteEmbeddingContext } from "./types.ts"
+
+export function embedGoogle(context: RemoteEmbeddingContext, request: EmbeddingBatchRequest): Promise<EmbeddingResult> {
+  return runEmbeddingOperation(context, (signal) => embedGoogleBatch(context, request, signal))
+}
+
+async function embedGoogleBatch(
+  context: RemoteEmbeddingContext,
+  request: EmbeddingBatchRequest,
+  signal: AbortSignal,
+): Promise<EmbeddingResult> {
+  const chunks = prepareChunks(request, {
+    maximum: 1024,
+    overlap: 128,
+    measure: (text) => Buffer.byteLength(text, "utf8"),
+    prefix: "",
+  })
+  if (!chunks.succeeded) {
+    return chunks
+  }
+
+  const vectors = await requestGoogleChunks(context, request.purpose, chunks.output, signal)
+  if (!vectors.succeeded) {
+    return vectors
+  }
+
+  return finishEmbeddingBatch(googleEmbeddingConfiguration, request, chunks.output, vectors.output)
+}
+
+async function requestGoogleChunks(
+  context: RemoteEmbeddingContext,
+  purpose: EmbeddingBatchRequest["purpose"],
+  chunks: readonly PreparedChunk[],
+  signal: AbortSignal,
+): Promise<ParsedVectors> {
+  const output: number[][] = []
+  for (let start = 0; start < chunks.length; start += 2) {
+    signal.throwIfAborted()
+    const group = chunks.slice(start, start + 2)
+    const responses = await Promise.all(group.map((chunk) => requestGoogle(context, chunk.text, purpose, signal)))
+    for (const response of responses) {
+      if (!response.succeeded) {
+        return response
+      }
+      output.push(...response.output)
+    }
+  }
+  return { succeeded: true, output }
+}
+
+async function requestGoogle(
+  context: RemoteEmbeddingContext,
+  text: string,
+  purpose: EmbeddingBatchRequest["purpose"],
+  signal: AbortSignal,
+): Promise<ParsedVectors> {
+  let endpoint = "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:embedContent"
+  if (typeof context.endpoint === "string") {
+    endpoint = context.endpoint
+  }
+
+  let taskType = "RETRIEVAL_DOCUMENT"
+  if (purpose === "query") {
+    taskType = "RETRIEVAL_QUERY"
+  }
+
+  const response = await postEmbedding({
+    endpoint,
+    apiKey: context.apiKey,
+    provider: "google",
+    signal,
+    body: JSON.stringify({
+      content: { parts: [{ text }] },
+      taskType,
+      outputDimensionality: 768,
+      embedContentConfig: { autoTruncate: false },
+    }),
+  })
+  if (!response.succeeded) {
+    return response
+  }
+
+  return parseGoogleEmbedding(response.output)
+}

package/src/services/embeddings/http.ts

@@ -0,0 +1,64 @@
+import type { EmbeddingHttpRequest, EmbeddingHttpResult } from "./types.ts"
+
+export async function postEmbedding(request: EmbeddingHttpRequest): Promise<EmbeddingHttpResult> {
+  if (request.apiKey.trim().length === 0 || !allowedEmbeddingEndpoint(request)) {
+    return { succeeded: false, status: "unavailable" }
+  }
+  const headers = new Headers({ "content-type": "application/json" })
+  if (request.provider === "openai") {
+    headers.set("authorization", `Bearer ${request.apiKey}`)
+  } else {
+    headers.set("x-goog-api-key", request.apiKey)
+  }
+  const response = await fetch(request.endpoint, {
+    method: "POST",
+    headers,
+    body: request.body,
+    signal: request.signal,
+    redirect: "error",
+  })
+  if (!response.ok) {
+    await response.body?.cancel()
+    return { succeeded: false, status: "unavailable" }
+  }
+  return readEmbeddingBody(response)
+}
+
+function allowedEmbeddingEndpoint(request: EmbeddingHttpRequest): boolean {
+  const url = new URL(request.endpoint)
+  if (url.username !== "" || url.password !== "" || url.search !== "" || url.hash !== "") {
+    return false
+  }
+  if (url.protocol === "http:" && ["127.0.0.1", "[::1]"].includes(url.hostname)) {
+    return true
+  }
+  if (request.provider === "openai") {
+    return url.href === "https://api.openai.com/v1/embeddings"
+  }
+  return url.href === "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:embedContent"
+}
+
+async function readEmbeddingBody(response: Response): Promise<EmbeddingHttpResult> {
+  if (!response.body) {
+    return { succeeded: false, status: "invalid-response" }
+  }
+  const reader = response.body.getReader()
+  const buffers: Uint8Array[] = []
+  let length = 0
+  try {
+    while (true) {
+      const next = await reader.read()
+      if (next.done) {
+        return { succeeded: true, output: Buffer.concat(buffers).toString("utf8") }
+      }
+      length += next.value.byteLength
+      if (length > 4_000_000) {
+        await reader.cancel()
+        return { succeeded: false, status: "invalid-response" }
+      }
+      buffers.push(next.value)
+    }
+  } finally {
+    reader.releaseLock()
+  }
+}

package/src/services/embeddings/index.ts

@@ -0,0 +1,17 @@
+export * from "./asset-download.ts"
+export * from "./asset-fetch.ts"
+export * from "./asset-inspection.ts"
+export * from "./batch.ts"
+export * from "./chunks.ts"
+export * from "./configuration.ts"
+export * from "./constants.ts"
+export * from "./google.ts"
+export * from "./inference.ts"
+export * from "./local.ts"
+export * from "./model-asset.ts"
+export * from "./model-cache.ts"
+export * from "./openai.ts"
+export * from "./operation.ts"
+export * from "./responses.ts"
+export * from "./runtime.ts"
+export * from "./types.ts"

package/src/services/embeddings/inference.ts

@@ -0,0 +1,35 @@
+import type { FeatureExtractionPipeline } from "@huggingface/transformers"
+
+import type { LocalModel } from "./types.ts"
+
+export async function loadCachedLocalModel(modelPath: string): Promise<LocalModel> {
+  const { pipeline } = await import("@huggingface/transformers")
+  const extractor = await pipeline("feature-extraction", modelPath, {
+    device: "cpu",
+    dtype: "q8",
+    local_files_only: true,
+    cache_dir: modelPath,
+  })
+  return localModelFromPipeline(extractor)
+}
+
+function localModelFromPipeline(extractor: FeatureExtractionPipeline): LocalModel {
+  return {
+    countTokens(text) {
+      return extractor.tokenizer.encode(text, { add_special_tokens: true }).length
+    },
+    async extract(text) {
+      const tensor = await extractor(text, { pooling: "mean", normalize: true })
+      if (tensor.dims.length !== 2 || tensor.dims[0] !== 1 || tensor.dims[1] !== 384) {
+        return null
+      }
+      if (!(tensor.data instanceof Float32Array)) {
+        return null
+      }
+      return Array.from(tensor.data)
+    },
+    async dispose() {
+      await extractor.dispose()
+    },
+  }
+}

package/src/services/embeddings/local.ts

@@ -0,0 +1,138 @@
+import { type EmbeddingBatchRequest, type EmbeddingResult, localEmbeddingConfiguration } from "#src/domain"
+import { finishEmbeddingBatch } from "./batch.ts"
+import { prepareChunks, validateEmbeddingRequest } from "./chunks.ts"
+import { runEmbeddingOperation } from "./operation.ts"
+import { loadLocalModel } from "./runtime.ts"
+import type {
+  LocalEmbedder,
+  LocalEmbeddingOptions,
+  LocalEmbeddingState,
+  LocalModel,
+  ParsedVectors,
+  PreparedChunk,
+} from "./types.ts"
+
+export function createLocalEmbedder(options: LocalEmbeddingOptions): LocalEmbedder {
+  const state: LocalEmbeddingState = { options: { ...options }, model: null, queue: Promise.resolve(), disposed: false }
+  return {
+    embed(context, request) {
+      return runEmbeddingOperation(context, (signal) => enqueueLocalEmbedding(state, request, signal))
+    },
+    async dispose() {
+      state.disposed = true
+      await state.queue
+      if (state.model) {
+        const model = await state.model.catch(() => null)
+        if (model) {
+          await model.dispose()
+        }
+        state.model = null
+      }
+    },
+  }
+}
+
+function enqueueLocalEmbedding(
+  state: LocalEmbeddingState,
+  request: EmbeddingBatchRequest,
+  signal: AbortSignal,
+): Promise<EmbeddingResult> {
+  const pending = state.queue.then(() => embedLocally(state, request, signal))
+  state.queue = pending.then(
+    () => {},
+    () => {},
+  )
+  return pending
+}
+
+async function getLocalModel(state: LocalEmbeddingState, signal: AbortSignal): Promise<LocalModel> {
+  if (state.model) {
+    return state.model
+  }
+
+  let loader = loadLocalModel
+  if (state.options.loadModel) {
+    loader = state.options.loadModel
+  }
+  const options = { cachePath: state.options.cachePath, signal }
+  if (state.options.onProgress) {
+    state.model = loader({ ...options, onProgress: state.options.onProgress })
+  } else {
+    state.model = loader(options)
+  }
+
+  const pending = state.model
+  try {
+    const model = await pending
+    if (signal.aborted || state.disposed) {
+      await model.dispose()
+      signal.throwIfAborted()
+      throw new Error("Local embedder disposed")
+    }
+    return model
+  } catch (error) {
+    if (state.model === pending) {
+      state.model = null
+    }
+    throw error
+  }
+}
+
+async function embedLocally(
+  state: LocalEmbeddingState,
+  request: EmbeddingBatchRequest,
+  signal: AbortSignal,
+): Promise<EmbeddingResult> {
+  signal.throwIfAborted()
+  if (state.disposed || state.options.cachePath.trim().length === 0) {
+    return { succeeded: false, status: "unavailable" }
+  }
+
+  const invalid = validateEmbeddingRequest(request)
+  if (invalid) {
+    return invalid
+  }
+
+  const model = await getLocalModel(state, signal)
+  signal.throwIfAborted()
+
+  let prefix = "passage: "
+  if (request.purpose === "query") {
+    prefix = "query: "
+  }
+
+  const chunks = prepareChunks(request, {
+    maximum: 512,
+    overlap: 32,
+    measure: (text) => model.countTokens(text),
+    prefix,
+  })
+  if (!chunks.succeeded) {
+    return chunks
+  }
+
+  const vectors = await extractLocalChunks(model, chunks.output, signal)
+  if (!vectors.succeeded) {
+    return vectors
+  }
+
+  return finishEmbeddingBatch(localEmbeddingConfiguration, request, chunks.output, vectors.output)
+}
+
+async function extractLocalChunks(
+  model: LocalModel,
+  chunks: readonly PreparedChunk[],
+  signal: AbortSignal,
+): Promise<ParsedVectors> {
+  const output: number[][] = []
+  for (const chunk of chunks) {
+    signal.throwIfAborted()
+    const vector = await model.extract(chunk.text)
+    if (!vector) {
+      return { succeeded: false, status: "invalid-response" }
+    }
+    output.push([...vector])
+  }
+  signal.throwIfAborted()
+  return { succeeded: true, output }
+}

package/src/services/embeddings/model-asset.ts

@@ -0,0 +1,68 @@
+import { randomUUID } from "node:crypto"
+import { link, mkdir, open, unlink } from "node:fs/promises"
+import { basename, dirname, join } from "node:path"
+
+import { downloadModelAsset } from "./asset-download.ts"
+import { inspectModelAsset } from "./asset-inspection.ts"
+import type { ModelAssetRequest, ModelCacheResult } from "./types.ts"
+
+export async function provisionModelAsset(request: ModelAssetRequest): Promise<ModelCacheResult> {
+  try {
+    const path = join(request.directory, request.asset.path)
+    const status = await inspectModelAsset(path, request.asset, request.signal)
+    if (status === "valid") {
+      return { succeeded: true, output: path }
+    }
+    if (status === "invalid") {
+      return { succeeded: false, status: "unavailable" }
+    }
+    await mkdir(dirname(path), { recursive: true })
+    return await publishModelAsset(request, path)
+  } catch {
+    if (request.signal.aborted) {
+      return { succeeded: false, status: "aborted" }
+    }
+    return { succeeded: false, status: "unavailable" }
+  }
+}
+
+async function publishModelAsset(request: ModelAssetRequest, path: string): Promise<ModelCacheResult> {
+  const temporary = join(dirname(path), `.${basename(path)}.${randomUUID()}.tmp`)
+  const file = await open(temporary, "wx")
+  let result: ModelCacheResult = { succeeded: false, status: "unavailable" }
+  try {
+    const valid = await downloadModelAsset(request, file)
+    await file.sync()
+    await file.close()
+    request.signal.throwIfAborted()
+    if (valid) {
+      result = await linkModelAsset(request, temporary, path)
+    }
+  } catch {
+    if (request.signal.aborted) {
+      result = { succeeded: false, status: "aborted" }
+    }
+  } finally {
+    try {
+      await file.close()
+      await unlink(temporary)
+    } catch {
+      result = { succeeded: false, status: "unavailable", cleanupFailed: true }
+    }
+  }
+  return result
+}
+
+async function linkModelAsset(request: ModelAssetRequest, temporary: string, path: string): Promise<ModelCacheResult> {
+  try {
+    await link(temporary, path)
+    return { succeeded: true, output: path }
+  } catch (error) {
+    if (error instanceof Error && "code" in error && error.code === "EEXIST") {
+      if ((await inspectModelAsset(path, request.asset, request.signal)) === "valid") {
+        return { succeeded: true, output: path }
+      }
+    }
+    return { succeeded: false, status: "unavailable" }
+  }
+}

package/src/services/embeddings/model-cache.ts

@@ -0,0 +1,30 @@
+import { isAbsolute, join } from "node:path"
+
+import { localEmbeddingConfiguration, localEmbeddingPolicy } from "#src/domain"
+import { fetchPublicModelAsset } from "./asset-fetch.ts"
+import { localModelAssets } from "./constants.ts"
+import { provisionModelAsset } from "./model-asset.ts"
+import type { LocalModelOptions, ModelAssetRequest, ModelCacheResult } from "./types.ts"
+
+export async function provisionLocalModel(options: LocalModelOptions): Promise<ModelCacheResult> {
+  if (!isAbsolute(options.cachePath)) {
+    return { succeeded: false, status: "unavailable" }
+  }
+  const directory = join(options.cachePath, localEmbeddingConfiguration.model, localEmbeddingPolicy.revision)
+  for (const asset of localModelAssets) {
+    const request: ModelAssetRequest = {
+      directory,
+      asset,
+      signal: options.signal,
+      transport: fetchPublicModelAsset,
+    }
+    if (options.onProgress) {
+      request.onProgress = options.onProgress
+    }
+    const result = await provisionModelAsset(request)
+    if (!result.succeeded) {
+      return result
+    }
+  }
+  return { succeeded: true, output: directory }
+}

package/src/services/embeddings/openai.ts

@@ -0,0 +1,64 @@
+import { type EmbeddingBatchRequest, type EmbeddingResult, openAIEmbeddingConfiguration } from "#src/domain"
+import { finishEmbeddingBatch } from "./batch.ts"
+import { prepareChunks } from "./chunks.ts"
+import { postEmbedding } from "./http.ts"
+import { runEmbeddingOperation } from "./operation.ts"
+import { parseOpenAIEmbeddings } from "./responses.ts"
+import type { ParsedVectors, PreparedChunk, RemoteEmbeddingContext } from "./types.ts"
+
+export function embedOpenAI(context: RemoteEmbeddingContext, request: EmbeddingBatchRequest): Promise<EmbeddingResult> {
+  return runEmbeddingOperation(context, (signal) => embedOpenAIBatch(context, request, signal))
+}
+
+async function embedOpenAIBatch(
+  context: RemoteEmbeddingContext,
+  request: EmbeddingBatchRequest,
+  signal: AbortSignal,
+): Promise<EmbeddingResult> {
+  const chunks = prepareChunks(request, {
+    maximum: 4096,
+    overlap: 128,
+    measure: (text) => Buffer.byteLength(text, "utf8"),
+    prefix: "",
+  })
+  if (!chunks.succeeded) {
+    return chunks
+  }
+  const vectors: number[][] = []
+  for (let start = 0; start < chunks.output.length; start += 32) {
+    signal.throwIfAborted()
+    const response = await requestOpenAI(context, chunks.output.slice(start, start + 32), signal)
+    if (!response.succeeded) {
+      return response
+    }
+    vectors.push(...response.output)
+  }
+  return finishEmbeddingBatch(openAIEmbeddingConfiguration, request, chunks.output, vectors)
+}
+
+async function requestOpenAI(
+  context: RemoteEmbeddingContext,
+  chunks: readonly PreparedChunk[],
+  signal: AbortSignal,
+): Promise<ParsedVectors> {
+  let endpoint = "https://api.openai.com/v1/embeddings"
+  if (typeof context.endpoint === "string") {
+    endpoint = context.endpoint
+  }
+  const response = await postEmbedding({
+    endpoint,
+    apiKey: context.apiKey,
+    provider: "openai",
+    signal,
+    body: JSON.stringify({
+      model: openAIEmbeddingConfiguration.model,
+      input: chunks.map((chunk) => chunk.text),
+      encoding_format: "float",
+      dimensions: 1536,
+    }),
+  })
+  if (!response.succeeded) {
+    return response
+  }
+  return parseOpenAIEmbeddings(response.output, chunks.length)
+}

package/src/services/embeddings/operation.ts

@@ -0,0 +1,31 @@
+import type { EmbeddingFailure, EmbeddingResult } from "#src/domain"
+import { EmbeddingResourceError } from "./errors.ts"
+import type { EmbeddingOperation, EmbeddingOperationContext } from "./types.ts"
+
+export async function runEmbeddingOperation(
+  context: EmbeddingOperationContext,
+  operation: EmbeddingOperation,
+): Promise<EmbeddingResult> {
+  if (context.signal.aborted) {
+    return { succeeded: false, status: "aborted" }
+  }
+  const { signal } = context
+  let cancel = () => {}
+  const aborted = new Promise<EmbeddingFailure>((resolve) => {
+    cancel = () => resolve({ succeeded: false, status: "aborted" })
+    signal.addEventListener("abort", cancel, { once: true })
+  })
+  try {
+    return await Promise.race([operation(signal), aborted])
+  } catch (error) {
+    if (signal.aborted) {
+      return { succeeded: false, status: "aborted" }
+    }
+    if (error instanceof EmbeddingResourceError) {
+      return error.failure
+    }
+    return { succeeded: false, status: "unavailable" }
+  } finally {
+    signal.removeEventListener("abort", cancel)
+  }
+}