pi-obsidian-memory 0.0.1 → 0.1.1

package/src/services/decision-model/types.ts

@@ -0,0 +1,182 @@
+import type { Static } from "typebox"
+
+import type {
+  CandidateContent,
+  CandidateDecisionRequest,
+  DecisionDiagnostics,
+  DecisionFailureCode,
+  DecisionRequest,
+  DecisionResult,
+} from "#src/domain"
+import type { LayaResponseSchema, LayaResultCandidateSchema } from "./schemas.ts"
+
+export type LayaResponse = Static<typeof LayaResponseSchema>
+export type LayaResultCandidate = Static<typeof LayaResultCandidateSchema>
+
+export type {
+  CandidateContent,
+  CandidateDecisionRequest,
+  CandidateIdentity,
+  CandidateScore,
+  DecisionDiagnostics,
+  DecisionFailure,
+  DecisionFailureCode,
+  DecisionRequest,
+  DecisionResult,
+  DecisionUsage,
+  RelevanceDecision,
+} from "#src/domain"
+
+export interface DecisionLimits {
+  signal: AbortSignal
+}
+
+export interface LayaContext extends DecisionLimits {
+  provider: "laya"
+  endpoint: string
+  model: string
+  maxTokens: number
+  headMaxTokens: number
+  apiKey?: string
+}
+
+export interface JevContext extends DecisionLimits {
+  provider: "typesafe"
+  model: string
+  classifier: JevClassifierPort
+}
+
+export type DecisionContext = LayaContext | JevContext
+
+export interface MemoryDecisionCall {
+  context: DecisionContext
+  request: DecisionRequest
+}
+
+export interface CandidateDecisionCall {
+  context: DecisionContext
+  request: CandidateDecisionRequest
+}
+
+export interface BoolCriteria {
+  true: string
+  false: string
+}
+
+export interface BoolQuestion {
+  type: "bool"
+  instructions: string
+  criteria: BoolCriteria
+}
+
+export interface NeededQuestions {
+  needed: BoolQuestion
+}
+
+export interface RelevantQuestions {
+  relevant: BoolQuestion
+}
+
+export interface CandidateState {
+  message: string
+  candidate: CandidateContent
+}
+
+export interface DecisionPrompt {
+  state: DecisionRequest | CandidateState
+  questions: NeededQuestions | RelevantQuestions
+}
+
+export interface JevModelReference {
+  provider: "typesafe"
+  id: string
+}
+
+export interface JevClassifyOptions {
+  signal: AbortSignal
+  maxRetries: 0
+}
+
+export interface JevClassifyRequest {
+  model: JevModelReference
+  context: DecisionPrompt
+  options: JevClassifyOptions
+}
+
+export interface JevAnswer {
+  type: string
+  probability?: number
+}
+
+export interface JevAnswers {
+  needed?: JevAnswer
+  relevant?: JevAnswer
+}
+
+export interface JevCost {
+  input: number
+  output: number
+  cacheRead: number
+  cacheWrite: number
+  total: number
+}
+
+export interface JevUsage {
+  input: number
+  output: number
+  cacheRead: number
+  cacheWrite: number
+  totalTokens: number
+  cost: JevCost
+}
+
+export interface JevClassifierResult {
+  api: string
+  provider: string
+  model: string
+  answers: JevAnswers
+  usage?: JevUsage
+  stopReason: "stop" | "error" | "aborted"
+  errorMessage?: string
+  timestamp: number
+}
+
+export interface JevClassifierPort {
+  classify(request: JevClassifyRequest): Promise<JevClassifierResult>
+}
+
+export interface ProviderCall {
+  prompt: DecisionPrompt
+  question: "needed" | "relevant"
+  signal: AbortSignal
+  diagnostics: DecisionDiagnostics
+}
+
+export interface DecisionOperation {
+  run(signal: AbortSignal): Promise<DecisionResult>
+}
+
+export interface LayaValueRequest {
+  question: "needed" | "relevant"
+  diagnostics: DecisionDiagnostics
+}
+
+export interface LayaReadRequest extends LayaValueRequest {
+  response: Response
+}
+
+export interface DecisionBatchItem {
+  prompt: DecisionPrompt
+  diagnostics: DecisionDiagnostics
+}
+
+export interface LayaBatchCall {
+  items: DecisionBatchItem[]
+  signal: AbortSignal
+}
+
+export type DecisionInterruption<Outcome> = (code: DecisionFailureCode) => Outcome
+
+export interface InterruptibleOperation<Outcome> {
+  run(signal: AbortSignal): Promise<Outcome>
+}

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

@@ -0,0 +1,45 @@
+import { createHash } from "node:crypto"
+import type { FileHandle } from "node:fs/promises"
+
+import type { ModelAssetRequest } from "./types.ts"
+
+export async function downloadModelAsset(request: ModelAssetRequest, file: FileHandle): Promise<boolean> {
+  request.signal.throwIfAborted()
+  const response = await request.transport({ path: request.asset.path, signal: request.signal })
+  if (!response.ok || !response.body) {
+    await response.body?.cancel()
+    return false
+  }
+  const reader = response.body.getReader()
+  try {
+    return await writeModelAsset(request, file, reader)
+  } finally {
+    await reader.cancel()
+    reader.releaseLock()
+  }
+}
+
+async function writeModelAsset(
+  request: ModelAssetRequest,
+  file: FileHandle,
+  reader: ReadableStreamDefaultReader<Uint8Array>,
+): Promise<boolean> {
+  const hash = createHash("sha256")
+  let loaded = 0
+  while (true) {
+    request.signal.throwIfAborted()
+    const chunk = await reader.read()
+    if (chunk.done) {
+      return loaded === request.asset.bytes && hash.digest("hex") === request.asset.sha256
+    }
+    loaded += chunk.value.byteLength
+    if (loaded > request.asset.bytes) {
+      return false
+    }
+    await file.writeFile(chunk.value)
+    hash.update(chunk.value)
+    if (request.onProgress) {
+      request.onProgress({ loaded, total: request.asset.bytes })
+    }
+  }
+}

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

@@ -0,0 +1,39 @@
+import { localEmbeddingConfiguration, localEmbeddingPolicy } from "#src/domain"
+import { EmbeddingResourceError } from "./errors.ts"
+import type { ModelAssetFetchRequest } from "./types.ts"
+
+export async function fetchPublicModelAsset(request: ModelAssetFetchRequest): Promise<Response> {
+  const model = localEmbeddingConfiguration.model
+  const revision = localEmbeddingPolicy.revision
+  let url = new URL(`https://huggingface.co/${model}/resolve/${revision}/${request.path}`)
+  for (let redirects = 0; redirects <= 5; redirects += 1) {
+    request.signal.throwIfAborted()
+    if (!allowedModelAssetUrl(url)) {
+      break
+    }
+    const response = await fetch(url, { signal: request.signal, credentials: "omit", redirect: "manual" })
+    if (![301, 302, 303, 307, 308].includes(response.status)) {
+      return response
+    }
+    await response.body?.cancel()
+    const location = response.headers.get("location")
+    if (!location) {
+      break
+    }
+    url = new URL(location, url)
+  }
+  throw new EmbeddingResourceError({ succeeded: false, status: "unavailable" })
+}
+
+export function allowedModelAssetUrl(url: URL): boolean {
+  if (url.protocol !== "https:" || url.port !== "" || url.username !== "" || url.password !== "") {
+    return false
+  }
+  return [
+    "huggingface.co",
+    "cas-bridge.xethub.hf.co",
+    "cdn-lfs.huggingface.co",
+    "cdn-lfs-us-1.hf.co",
+    "us.aws.cdn.hf.co",
+  ].includes(url.hostname)
+}

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

@@ -0,0 +1,36 @@
+import { createHash } from "node:crypto"
+import { createReadStream } from "node:fs"
+import { stat } from "node:fs/promises"
+
+import type { AssetInspection, ModelAsset } from "./types.ts"
+
+export async function inspectModelAsset(
+  path: string,
+  asset: ModelAsset,
+  signal: AbortSignal,
+): Promise<AssetInspection> {
+  try {
+    signal.throwIfAborted()
+    const info = await stat(path)
+    if (!info.isFile() || info.size !== asset.bytes) {
+      return "invalid"
+    }
+    const hash = createHash("sha256")
+    for await (const chunk of createReadStream(path, { signal })) {
+      if (!Buffer.isBuffer(chunk)) {
+        return "invalid"
+      }
+      hash.update(chunk)
+    }
+    if (hash.digest("hex") !== asset.sha256) {
+      return "invalid"
+    }
+    return "valid"
+  } catch (error) {
+    signal.throwIfAborted()
+    if (error instanceof Error && "code" in error && error.code === "ENOENT") {
+      return "missing"
+    }
+    return "invalid"
+  }
+}

package/src/services/embeddings/batch.ts

@@ -0,0 +1,32 @@
+import {
+  aggregateVectors,
+  type EmbeddedInput,
+  type EmbeddingBatchRequest,
+  type EmbeddingConfiguration,
+  type EmbeddingResult,
+  embeddingFingerprint,
+  embeddingPolicy,
+} from "#src/domain"
+import type { PreparedChunk } from "./types.ts"
+
+export function finishEmbeddingBatch(
+  configuration: EmbeddingConfiguration,
+  request: EmbeddingBatchRequest,
+  chunks: readonly PreparedChunk[],
+  vectors: readonly (readonly number[])[],
+): EmbeddingResult {
+  if (chunks.length !== vectors.length) {
+    return { succeeded: false, status: "invalid-response" }
+  }
+  const dimensions = embeddingPolicy(configuration).dimensions
+  const items: EmbeddedInput[] = []
+  for (const [index, input] of request.inputs.entries()) {
+    const matching = vectors.filter((_vector, position) => chunks[position]?.inputIndex === index)
+    const vector = aggregateVectors(matching, dimensions)
+    if (!vector) {
+      return { succeeded: false, status: "invalid-response" }
+    }
+    items.push({ key: input.key, vector })
+  }
+  return { succeeded: true, output: { fingerprint: embeddingFingerprint(configuration), dimensions, items } }
+}

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()
+    },
+  }
+}