pi-obsidian-memory 0.0.1 → 0.1.1

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

package/src/services/embeddings/responses.ts

@@ -0,0 +1,41 @@
+import { Value } from "typebox/value"
+
+import { normalizeVector } from "#src/domain"
+import { GoogleEmbeddingResponseSchema, OpenAIEmbeddingResponseSchema } from "./schemas.ts"
+import type { ParsedVectors } from "./types.ts"
+
+export function parseOpenAIEmbeddings(text: string, count: number): ParsedVectors {
+  try {
+    const response = JSON.parse(text)
+    if (!Value.Check(OpenAIEmbeddingResponseSchema, response) || response.data.length !== count) {
+      return { succeeded: false, status: "invalid-response" }
+    }
+    const vectors = new Map<number, number[]>()
+    for (const item of response.data) {
+      const vector = normalizeVector(item.embedding, 1536)
+      if (item.index >= count || vectors.has(item.index) || !vector) {
+        return { succeeded: false, status: "invalid-response" }
+      }
+      vectors.set(item.index, vector)
+    }
+    return { succeeded: true, output: [...vectors.entries()].sort(([a], [b]) => a - b).map(([, vector]) => vector) }
+  } catch {
+    return { succeeded: false, status: "invalid-response" }
+  }
+}
+
+export function parseGoogleEmbedding(text: string): ParsedVectors {
+  try {
+    const response = JSON.parse(text)
+    if (!Value.Check(GoogleEmbeddingResponseSchema, response)) {
+      return { succeeded: false, status: "invalid-response" }
+    }
+    const vector = normalizeVector(response.embedding.values, 768)
+    if (!vector) {
+      return { succeeded: false, status: "invalid-response" }
+    }
+    return { succeeded: true, output: [vector] }
+  } catch {
+    return { succeeded: false, status: "invalid-response" }
+  }
+}

package/src/services/embeddings/runtime.ts

@@ -0,0 +1,13 @@
+import { EmbeddingResourceError } from "./errors.ts"
+import { loadCachedLocalModel } from "./inference.ts"
+import { provisionLocalModel } from "./model-cache.ts"
+import type { LocalModel, LocalModelOptions } from "./types.ts"
+
+export async function loadLocalModel(options: LocalModelOptions): Promise<LocalModel> {
+  const cache = await provisionLocalModel(options)
+  if (!cache.succeeded) {
+    throw new EmbeddingResourceError(cache)
+  }
+  options.signal.throwIfAborted()
+  return loadCachedLocalModel(cache.output)
+}

package/src/services/embeddings/schemas.ts

@@ -0,0 +1,53 @@
+import Type from "typebox"
+
+export const OpenAIEmbeddingResponseSchema = Type.Object(
+  {
+    object: Type.Literal("list"),
+    model: Type.Literal("text-embedding-3-small"),
+    data: Type.Array(
+      Type.Object(
+        {
+          object: Type.Literal("embedding"),
+          index: Type.Integer({ minimum: 0 }),
+          embedding: Type.Array(Type.Number(), { minItems: 1536, maxItems: 1536 }),
+        },
+        { additionalProperties: false },
+      ),
+      { maxItems: 32 },
+    ),
+    usage: Type.Object(
+      { prompt_tokens: Type.Integer({ minimum: 0 }), total_tokens: Type.Integer({ minimum: 0 }) },
+      { additionalProperties: false },
+    ),
+  },
+  { additionalProperties: false },
+)
+
+const GoogleEmbeddingUsageSchema = Type.Object(
+  {
+    promptTokenCount: Type.Optional(Type.Integer({ minimum: 0 })),
+    promptTokenDetails: Type.Optional(
+      Type.Array(
+        Type.Object(
+          { modality: Type.Literal("TEXT"), tokenCount: Type.Integer({ minimum: 0 }) },
+          { additionalProperties: false },
+        ),
+      ),
+    ),
+  },
+  { additionalProperties: false },
+)
+
+export const GoogleEmbeddingResponseSchema = Type.Object(
+  {
+    embedding: Type.Object(
+      {
+        values: Type.Array(Type.Number(), { minItems: 768, maxItems: 768 }),
+        shape: Type.Optional(Type.Array(Type.Integer({ minimum: 1 }), { maxItems: 4 })),
+      },
+      { additionalProperties: false },
+    ),
+    usageMetadata: Type.Optional(GoogleEmbeddingUsageSchema),
+  },
+  { additionalProperties: false },
+)

package/src/services/embeddings/types.ts

@@ -0,0 +1,116 @@
+import type { EmbeddingBatchRequest, EmbeddingFailure, EmbeddingResult } from "#src/domain"
+import type { Succeeded } from "#src/shared"
+
+export type {
+  EmbeddedInput,
+  EmbeddingBatch,
+  EmbeddingBatchRequest,
+  EmbeddingConfiguration,
+  EmbeddingFailure,
+  EmbeddingFailureStatus,
+  EmbeddingInput,
+  EmbeddingResult,
+} from "#src/domain"
+
+export interface EmbeddingOperationContext {
+  signal: AbortSignal
+}
+
+export interface RemoteEmbeddingContext extends EmbeddingOperationContext {
+  apiKey: string
+  endpoint?: string
+}
+
+export interface LocalRuntimeOptions {
+  cachePath: string
+  onProgress?: DownloadProgressListener
+}
+
+export interface DownloadProgress {
+  loaded: number
+  total: number
+}
+
+export type DownloadProgressListener = (progress: DownloadProgress) => void
+
+export interface LocalModel {
+  countTokens(text: string): number
+  extract(text: string): Promise<readonly number[] | null>
+  dispose(): Promise<void>
+}
+
+export type LocalModelLoader = (options: LocalModelOptions) => Promise<LocalModel>
+
+export interface LocalModelOptions extends LocalRuntimeOptions {
+  signal: AbortSignal
+}
+
+export interface LocalEmbeddingOptions extends LocalRuntimeOptions {
+  loadModel?: LocalModelLoader
+}
+
+export interface LocalEmbedder {
+  embed(context: EmbeddingOperationContext, request: EmbeddingBatchRequest): Promise<EmbeddingResult>
+  dispose(): Promise<void>
+}
+
+export interface LocalEmbeddingState {
+  options: LocalEmbeddingOptions
+  model: Promise<LocalModel> | null
+  queue: Promise<void>
+  disposed: boolean
+}
+
+export type EmbeddingOperation = (signal: AbortSignal) => Promise<EmbeddingResult>
+export type TextSize = (text: string) => number
+
+export interface ChunkPolicy {
+  maximum: number
+  overlap: number
+  measure: TextSize
+  prefix: string
+}
+
+export interface PreparedChunk {
+  inputIndex: number
+  text: string
+}
+
+export type PreparedChunks = Succeeded<PreparedChunk[]> | EmbeddingFailure
+
+export interface EmbeddingHttpRequest {
+  endpoint: string
+  apiKey: string
+  provider: "openai" | "google"
+  body: string
+  signal: AbortSignal
+}
+
+export type EmbeddingHttpResult = Succeeded<string> | EmbeddingFailure
+export type ParsedVectors = Succeeded<number[][]> | EmbeddingFailure
+
+export type ModelAssetPath = "config.json" | "tokenizer_config.json" | "tokenizer.json" | "onnx/model_quantized.onnx"
+
+export interface ModelAsset {
+  path: ModelAssetPath
+  bytes: number
+  sha256: string
+}
+
+export interface ModelAssetFetchRequest {
+  path: ModelAssetPath
+  signal: AbortSignal
+}
+
+export type ModelAssetTransport = (request: ModelAssetFetchRequest) => Promise<Response>
+
+export interface ModelAssetRequest {
+  directory: string
+  asset: ModelAsset
+  signal: AbortSignal
+  transport: ModelAssetTransport
+  onProgress?: DownloadProgressListener
+}
+
+export type ModelCacheResult = Succeeded<string> | EmbeddingFailure
+export type AssetInspection = "valid" | "missing" | "invalid"

package/src/services/id/id.ts

@@ -0,0 +1,5 @@
+import { ulid } from "ulid"
+
+export function newNoteId(): string {
+  return ulid()
+}

package/src/services/id/index.ts

@@ -0,0 +1 @@
+export * from "./id.ts"

package/src/services/index.ts

@@ -0,0 +1,11 @@
+export * from "./capture-cadence/index.ts"
+export * from "./clock/index.ts"
+export * from "./config/index.ts"
+export * from "./decision-model/index.ts"
+export * from "./embeddings/index.ts"
+export * from "./id/index.ts"
+export * from "./llm/index.ts"
+export * from "./search-index/index.ts"
+export * from "./shell-edit/index.ts"
+export * from "./shell-scripts/index.ts"
+export * from "./vault/index.ts"

package/src/services/llm/capture.ts

@@ -0,0 +1,19 @@
+import { CaptureActionSchema, captureInputProblems, parseCaptureAction } from "#src/domain"
+import { checkStep, completeStep, completionRequest, modelCallFailure } from "./complete.ts"
+import { CAPTURE_POLICY, DATA_POLICY } from "./constants.ts"
+import type { CaptureRequest, CaptureStepResult, CompletionPort } from "./types.ts"
+
+export async function captureStep(request: CaptureRequest, complete: CompletionPort): Promise<CaptureStepResult> {
+  const { context, transcriptDelta, observations, explicitSaves, limits } = request
+  const input = { context, transcriptDelta, observations, explicitSaves, limits }
+  const problems = captureInputProblems(input)
+  if (problems.length > 0) return modelCallFailure(`Invalid capture input: ${problems.join("; ")}`)
+  const snapshot = structuredClone(input)
+  const data = JSON.stringify(snapshot)
+  const refused = checkStep(request)
+  if (refused) return refused
+  const system = `${DATA_POLICY}\n${CAPTURE_POLICY}\nAction schema: ${JSON.stringify(CaptureActionSchema)}`
+  const result = await completeStep(completionRequest(request, { system, data }), complete)
+  if (!result.succeeded) return result
+  return parseCaptureAction(result.output, snapshot)
+}

package/src/services/llm/check-chat-models.ts

@@ -0,0 +1,38 @@
+import { NO_API_KEY } from "./constants.ts"
+import type { ChatModelAuth, ChatModelCheck, ChatModelPort, ChatModelReference } from "./types.ts"
+
+export function checkChatModels<TModel>(
+  references: ChatModelReference[],
+  port: ChatModelPort<TModel>,
+): Promise<ChatModelCheck[]> {
+  return Promise.all(references.map((reference) => checkChatModel(reference, port)))
+}
+
+async function checkChatModel<TModel>(
+  reference: ChatModelReference,
+  port: ChatModelPort<TModel>,
+): Promise<ChatModelCheck> {
+  const label = `${reference.provider}/${reference.id}`
+  const model = port.find(reference.provider, reference.id)
+  if (!model) {
+    return { model: reference, status: "model-not-found", problem: `Model ${label} is not registered in pi` }
+  }
+
+  if (!port.configured(model)) {
+    return { model: reference, status: "no-credentials", problem: `${NO_API_KEY} "${reference.provider}"` }
+  }
+
+  const auth = await resolveAuth(port, model, label)
+  if (!auth.ok) {
+    return { model: reference, status: "no-credentials", problem: auth.error }
+  }
+  return { model: reference, status: "ok" }
+}
+
+async function resolveAuth<TModel>(port: ChatModelPort<TModel>, model: TModel, label: string): Promise<ChatModelAuth> {
+  try {
+    return await port.auth(model)
+  } catch (error) {
+    return { ok: false, error: `Credentials for ${label} could not be resolved: ${String(error)}` }
+  }
+}