@takara-ai/miru-code 0.1.0

package/src/index/files.ts

@@ -0,0 +1,126 @@
+import { basename, extname } from "node:path";
+import type { ContentType } from "../types.ts";
+
+const MAX_FILE_BYTES = 1_000_000;
+const EMPTY_FILE_BYTES = 128;
+
+const EXTENSION_TO_LANGUAGE: Record<string, string> = {
+  ".py": "python",
+  ".pyi": "python",
+  ".js": "javascript",
+  ".mjs": "javascript",
+  ".cjs": "javascript",
+  ".jsx": "javascript",
+  ".ts": "typescript",
+  ".tsx": "typescript",
+  ".mts": "typescript",
+  ".cts": "typescript",
+  ".go": "go",
+  ".rs": "rust",
+  ".java": "java",
+  ".kt": "kotlin",
+  ".kts": "kotlin",
+  ".c": "c",
+  ".h": "c",
+  ".cpp": "cpp",
+  ".cc": "cpp",
+  ".cxx": "cpp",
+  ".hpp": "cpp",
+  ".cs": "csharp",
+  ".rb": "ruby",
+  ".php": "php",
+  ".swift": "swift",
+  ".scala": "scala",
+  ".clj": "clojure",
+  ".cljs": "clojure",
+  ".ex": "elixir",
+  ".exs": "elixir",
+  ".erl": "erlang",
+  ".hs": "haskell",
+  ".lua": "lua",
+  ".sh": "bash",
+  ".bash": "bash",
+  ".zsh": "bash",
+  ".fish": "fish",
+  ".sql": "sql",
+  ".r": "r",
+  ".R": "r",
+  ".dart": "dart",
+  ".zig": "zig",
+  ".vue": "vue",
+  ".svelte": "svelte",
+  ".md": "markdown",
+  ".mdx": "markdown",
+  ".rst": "rst",
+  ".txt": "text",
+  ".json": "json",
+  ".yaml": "yaml",
+  ".yml": "yaml",
+  ".toml": "toml",
+  ".ini": "ini",
+  ".cfg": "ini",
+  ".xml": "xml",
+  ".html": "html",
+  ".css": "css",
+  ".scss": "scss",
+  ".less": "less",
+  ".dockerfile": "dockerfile",
+};
+
+const CODE_EXTENSIONS = new Set(
+  Object.entries(EXTENSION_TO_LANGUAGE)
+    .filter(
+      ([, lang]) =>
+        !["markdown", "rst", "text", "json", "yaml", "toml", "ini", "xml", "html"].includes(lang),
+    )
+    .map(([ext]) => ext),
+);
+
+const DOC_EXTENSIONS = new Set([".md", ".mdx", ".rst", ".txt"]);
+const CONFIG_EXTENSIONS = new Set([".json", ".yaml", ".yml", ".toml", ".ini", ".cfg", ".xml"]);
+
+const CONTENT_EXTENSIONS: Record<ContentType, Set<string>> = {
+  code: CODE_EXTENSIONS,
+  docs: DOC_EXTENSIONS,
+  config: CONFIG_EXTENSIONS,
+};
+
+export function getExtensions(types: ContentType[]): string[] {
+  const exts = new Set<string>();
+  for (const t of types) {
+    for (const e of CONTENT_EXTENSIONS[t] ?? []) {
+      exts.add(e);
+    }
+  }
+  return [...exts].sort();
+}
+
+export function detectLanguage(filePath: string): string | null {
+  const ext = extname(filePath).toLowerCase();
+  if (basename(filePath).toLowerCase() === "dockerfile") {
+    return "dockerfile";
+  }
+  return EXTENSION_TO_LANGUAGE[ext] ?? null;
+}
+
+export type FileStatus = "too_large" | "empty" | "valid";
+
+export async function readFileText(filePath: string): Promise<string> {
+  const file = Bun.file(filePath);
+  return file.text();
+}
+
+export async function getFileStatus(filePath: string): Promise<FileStatus> {
+  const file = Bun.file(filePath);
+  const stat = await file.stat();
+  if (stat.size > MAX_FILE_BYTES) {
+    return "too_large";
+  }
+  if (stat.size < EMPTY_FILE_BYTES) {
+    const text = await file.text();
+    if (!text.trim()) {
+      return "empty";
+    }
+  }
+  return "valid";
+}

package/src/index/persistence.ts

@@ -0,0 +1,100 @@
+import { mkdir } from "node:fs/promises";
+import { join } from "node:path";
+import { BM25Index } from "./bm25.ts";
+import { VectorIndex } from "./dense.ts";
+import { QuantizedVectorIndex } from "./quantize.ts";
+import type { SemanticIndex } from "./semantic-index.ts";
+import { resolveSemanticStorage, semanticStorageOf } from "./vector-storage.ts";
+
+export interface PersistencePaths {
+  root: string;
+  bm25Index: string;
+  semanticIndex: string;
+  chunks: string;
+  metadata: string;
+}
+
+export function persistencePaths(root: string): PersistencePaths {
+  return {
+    root,
+    bm25Index: join(root, "bm25_index.json"),
+    semanticIndex: join(root, "semantic_index"),
+    chunks: join(root, "chunks.json"),
+    metadata: join(root, "metadata.json"),
+  };
+}
+
+export async function pathsExist(paths: PersistencePaths): Promise<string[]> {
+  const missing: string[] = [];
+  const checks: [string, string][] = [
+    ["bm25", paths.bm25Index],
+    ["semantic", `${paths.semanticIndex}/meta.json`],
+    ["chunks", paths.chunks],
+    ["metadata", paths.metadata],
+  ];
+  for (const [name, p] of checks) {
+    if (!(await Bun.file(p).exists())) {
+      missing.push(name);
+    }
+  }
+  return missing;
+}
+
+export async function saveBm25(index: BM25Index, path: string): Promise<void> {
+  await Bun.write(path, JSON.stringify(index.toJSON()));
+}
+
+export async function loadBm25(path: string): Promise<BM25Index> {
+  return BM25Index.fromJSON(JSON.parse(await Bun.file(path).text()));
+}
+
+export async function saveSemantic(index: SemanticIndex, path: string): Promise<void> {
+  if (index instanceof QuantizedVectorIndex) {
+    await index.save(path);
+    return;
+  }
+  if (index instanceof VectorIndex) {
+    await index.save(path);
+    return;
+  }
+  throw new Error("Unsupported semantic index type");
+}
+
+export async function loadSemantic(path: string): Promise<SemanticIndex> {
+  const meta = JSON.parse(await Bun.file(`${path}/meta.json`).text()) as {
+    storage?: string;
+  };
+  if (meta.storage === "int8") {
+    return QuantizedVectorIndex.load(path);
+  }
+  return VectorIndex.load(path);
+}
+
+export async function saveIndexBundle(options: {
+  paths: PersistencePaths;
+  bm25: BM25Index;
+  semantic: SemanticIndex;
+  chunks: unknown[];
+  metadata: Record<string, unknown>;
+}): Promise<void> {
+  const { paths, bm25, semantic, chunks, metadata } = options;
+  await mkdir(paths.root, { recursive: true });
+  await saveBm25(bm25, paths.bm25Index);
+  await saveSemantic(semantic, paths.semanticIndex);
+  await Bun.write(
+    paths.metadata,
+    JSON.stringify({
+      ...metadata,
+      vector_storage: semanticStorageOf(semantic),
+    }),
+  );
+  await Bun.write(paths.chunks, JSON.stringify(chunks));
+}
+
+export async function semanticIndexMatchesStorage(path: string): Promise<boolean> {
+  const meta = JSON.parse(await Bun.file(`${path}/meta.json`).text()) as {
+    storage?: string;
+  };
+  const stored = meta.storage === "int8" ? "int8" : "float32";
+  return stored === resolveSemanticStorage();
+}

package/src/index/quantize.ts

@@ -0,0 +1,168 @@
+import type { SemanticIndex } from "./semantic-index.ts";
+import { selectTopKByDistance } from "./topk.ts";
+
+export interface QuantizedVector {
+  codes: Int8Array;
+  scale: number;
+}
+
+/** Symmetric int8 quantization per vector (embeddings are L2-normalized upstream). */
+export function quantizeVector(vector: Float32Array): QuantizedVector {
+  let maxAbs = 0;
+  for (let i = 0; i < vector.length; i++) {
+    const value = vector[i] ?? 0;
+    const abs = Math.abs(value);
+    if (abs > maxAbs) {
+      maxAbs = abs;
+    }
+  }
+  const scale = maxAbs > 0 ? maxAbs / 127 : 1;
+  const codes = new Int8Array(vector.length);
+  if (maxAbs > 0) {
+    const inv = 127 / maxAbs;
+    for (let i = 0; i < vector.length; i++) {
+      const value = vector[i] ?? 0;
+      codes[i] = Math.max(-127, Math.min(127, Math.round(value * inv)));
+    }
+  }
+  return { codes, scale };
+}
+
+function quantizedDot(query: QuantizedVector, docCodes: Int8Array, docScale: number): number {
+  let sum = 0;
+  const q = query.codes;
+  for (let i = 0; i < q.length; i++) {
+    sum += (q[i] ?? 0) * (docCodes[i] ?? 0);
+  }
+  return sum * query.scale * docScale;
+}
+
+/** int8 codes + per-vector scale; ~4x smaller than float32 at same dimensionality. */
+export class QuantizedVectorIndex implements SemanticIndex {
+  private readonly codes: Int8Array[];
+  private readonly scales: Float32Array;
+  private readonly dim: number;
+
+  constructor(vectors: Float32Array[]) {
+    if (vectors.length === 0) {
+      this.dim = 0;
+      this.codes = [];
+      this.scales = new Float32Array(0);
+      return;
+    }
+    this.dim = vectors[0]?.length ?? 0;
+    this.codes = new Array(vectors.length);
+    this.scales = new Float32Array(vectors.length);
+    for (let i = 0; i < vectors.length; i++) {
+      const vector = vectors[i];
+      if (!vector) {
+        continue;
+      }
+      const { codes, scale } = quantizeVector(vector);
+      this.codes[i] = codes;
+      this.scales[i] = scale;
+    }
+  }
+
+  static fromPersisted(
+    codes: Int8Array[],
+    scales: Float32Array,
+    dim: number,
+  ): QuantizedVectorIndex {
+    const index = Object.create(QuantizedVectorIndex.prototype) as QuantizedVectorIndex;
+    Object.assign(index, { dim, codes, scales });
+    return index;
+  }
+
+  get size(): number {
+    return this.codes.length;
+  }
+
+  get dimensions(): number {
+    return this.dim;
+  }
+
+  memoryBytes(): number {
+    return this.codes.length * this.dim + this.scales.byteLength;
+  }
+
+  query(
+    queryVector: Float32Array,
+    k: number,
+    selector?: number[],
+  ): { indices: number[]; distances: number[] } {
+    if (k < 1) {
+      throw new Error(`k should be >= 1, is now ${k}`);
+    }
+    if (this.size === 0) {
+      return { indices: [], distances: [] };
+    }
+
+    const q = quantizeVector(queryVector);
+    const indices = selector ?? Array.from({ length: this.size }, (_, i) => i);
+    const effectiveK = Math.min(k, indices.length);
+    if (effectiveK === 0) {
+      return { indices: [], distances: [] };
+    }
+
+    const top = selectTopKByDistance(
+      indices.flatMap((idx) => {
+        const codes = this.codes[idx];
+        if (!codes) {
+          return [];
+        }
+        const scale = this.scales[idx];
+        if (scale === undefined) {
+          return [];
+        }
+        const similarity = quantizedDot(q, codes, scale);
+        return [{ index: idx, distance: 1 - similarity }];
+      }),
+      effectiveK,
+    );
+
+    return {
+      indices: top.map((t) => t.index),
+      distances: top.map((t) => t.distance),
+    };
+  }
+
+  async save(dir: string): Promise<void> {
+    const { mkdir } = await import("node:fs/promises");
+    await mkdir(dir, { recursive: true });
+
+    const count = this.size;
+    const dim = this.dim;
+    const flatCodes = new Int8Array(count * dim);
+    for (let i = 0; i < count; i++) {
+      const codes = this.codes[i];
+      if (codes) {
+        flatCodes.set(codes, i * dim);
+      }
+    }
+    await Bun.write(`${dir}/codes.bin`, flatCodes);
+    await Bun.write(`${dir}/scales.bin`, this.scales);
+    await Bun.write(
+      `${dir}/meta.json`,
+      JSON.stringify({ count, dimensions: dim, storage: "int8" }),
+    );
+  }
+
+  static async load(dir: string): Promise<QuantizedVectorIndex> {
+    const meta = JSON.parse(await Bun.file(`${dir}/meta.json`).text()) as {
+      count: number;
+      dimensions: number;
+      storage?: string;
+    };
+    const rawCodes = new Int8Array(await Bun.file(`${dir}/codes.bin`).arrayBuffer());
+    const scales = new Float32Array(await Bun.file(`${dir}/scales.bin`).arrayBuffer());
+    const index = QuantizedVectorIndex.fromPersisted(
+      Array.from({ length: meta.count }, (_, i) =>
+        rawCodes.subarray(i * meta.dimensions, (i + 1) * meta.dimensions),
+      ),
+      scales,
+      meta.dimensions,
+    );
+    return index;
+  }
+}

package/src/index/semantic-index.ts

@@ -0,0 +1,11 @@
+/** Shared query interface for float and quantized vector indexes. */
+export interface SemanticIndex {
+  readonly size: number;
+  readonly dimensions: number;
+  memoryBytes(): number;
+  query(
+    queryVector: Float32Array,
+    k: number,
+    selector?: number[],
+  ): { indices: number[]; distances: number[] };
+}

package/src/index/sparse.ts

@@ -0,0 +1,25 @@
+import type { Chunk } from "../types.ts";
+
+export function selectorToMask(
+  selector: number[] | null | undefined,
+  size: number,
+): boolean[] | undefined {
+  if (!selector || selector.length === 0) {
+    return undefined;
+  }
+  const mask = new Array<boolean>(size).fill(false);
+  for (const idx of selector) {
+    if (idx >= 0 && idx < size) {
+      mask[idx] = true;
+    }
+  }
+  return mask;
+}
+
+export function enrichForBm25(chunk: Chunk): string {
+  const parts = chunk.file_path.replace(/\\/g, "/").split("/");
+  const stem = parts[parts.length - 1]?.replace(/\.[^.]+$/, "") ?? "";
+  const dirParts = parts.slice(0, -1).filter((p) => p && p !== "." && p !== "..");
+  const dirText = dirParts.slice(-3).join(" ");
+  return `${chunk.content} ${stem} ${stem} ${dirText}`;
+}

package/src/index/topk.ts

@@ -0,0 +1,52 @@
+export interface TopKDistanceEntry {
+  index: number;
+  distance: number;
+}
+
+/** Keep the k smallest distances (best cosine matches). */
+export function selectTopKByDistance(
+  entries: Iterable<TopKDistanceEntry>,
+  k: number,
+): TopKDistanceEntry[] {
+  const top: TopKDistanceEntry[] = [];
+  for (const entry of entries) {
+    if (top.length < k) {
+      top.push(entry);
+      top.sort((a, b) => b.distance - a.distance);
+      continue;
+    }
+    const worst = top[0];
+    if (worst && entry.distance < worst.distance) {
+      top[0] = entry;
+      top.sort((a, b) => b.distance - a.distance);
+    }
+  }
+  top.sort((a, b) => a.distance - b.distance);
+  return top;
+}
+
+/** Indices of the k highest scores in a dense array. */
+export function selectTopKScoreIndices(scores: number[], k: number): number[] {
+  if (k <= 0 || scores.length === 0) {
+    return [];
+  }
+  const top: { v: number; i: number }[] = [];
+  for (let i = 0; i < scores.length; i++) {
+    const v = scores[i];
+    if (v === undefined) {
+      continue;
+    }
+    if (top.length < k) {
+      top.push({ v, i });
+      top.sort((a, b) => a.v - b.v);
+      continue;
+    }
+    const worst = top[0];
+    if (worst && v > worst.v) {
+      top[0] = { v, i };
+      top.sort((a, b) => a.v - b.v);
+    }
+  }
+  top.sort((a, b) => b.v - a.v);
+  return top.map((x) => x.i);
+}

package/src/index/vector-storage.ts

@@ -0,0 +1,27 @@
+import { VectorIndex } from "./dense.ts";
+import { QuantizedVectorIndex } from "./quantize.ts";
+import type { SemanticIndex } from "./semantic-index.ts";
+
+export type SemanticStorage = "int8" | "float32";
+
+/** Default: int8 quantized vectors (~4x less RAM). Set MIRU_FLOAT_VECTORS=1 for float32. */
+export function resolveSemanticStorage(): SemanticStorage {
+  return process.env.MIRU_FLOAT_VECTORS === "1" || process.env.SEMBLE_FLOAT_VECTORS === "1"
+    ? "float32"
+    : "int8";
+}
+
+export function semanticStorageFromMetadata(metadata: Record<string, unknown>): SemanticStorage {
+  return metadata.vector_storage === "float32" ? "float32" : "int8";
+}
+
+export function buildSemanticIndex(vectors: Float32Array[]): SemanticIndex {
+  if (resolveSemanticStorage() === "int8") {
+    return new QuantizedVectorIndex(vectors);
+  }
+  return new VectorIndex(vectors);
+}
+
+export function semanticStorageOf(index: SemanticIndex): SemanticStorage {
+  return index instanceof QuantizedVectorIndex ? "int8" : "float32";
+}

package/src/index.ts

@@ -0,0 +1,3 @@
+export { clearCache, findIndexCachePath, resolveCacheFolder } from "./cache.ts";
+export { MiruIndex } from "./miru-index.ts";
+export type { Chunk, ContentType, SearchResult } from "./types.ts";

package/src/mcp/index-cache.ts

@@ -0,0 +1,95 @@
+import { watch } from "node:fs";
+import { resolve } from "node:path";
+import { MiruIndex } from "../miru-index.ts";
+import type { ContentType } from "../types.ts";
+import { computeSourceCacheKey, isGitUrl } from "../utils.ts";
+
+const CACHE_MAX_SIZE = 10;
+
+type IndexTask = Promise<MiruIndex>;
+
+export class IndexCache {
+  private readonly content: ContentType[];
+  private readonly tasks = new Map<string, IndexTask>();
+  private watcher: ReturnType<typeof watch> | null = null;
+
+  constructor(content: ContentType[] = ["code"]) {
+    this.content = content;
+  }
+
+  evict(source: string, ref?: string | null): void {
+    this.tasks.delete(computeSourceCacheKey(source, ref));
+  }
+
+  async get(source: string, ref?: string | null): Promise<MiruIndex> {
+    const cacheKey = computeSourceCacheKey(source, ref);
+    let task = this.tasks.get(cacheKey);
+    if (!task) {
+      if (this.tasks.size >= CACHE_MAX_SIZE) {
+        const oldest = this.tasks.keys().next().value;
+        if (oldest) {
+          this.tasks.delete(oldest);
+        }
+      }
+      task = MiruIndex.fromSource(source, this.content, undefined, ref);
+      this.tasks.set(cacheKey, task);
+    }
+
+    try {
+      return await task;
+    } catch (err) {
+      if (this.tasks.get(cacheKey) === task) {
+        this.tasks.delete(cacheKey);
+      }
+      throw err;
+    }
+  }
+
+  startWatcher(path: string): void {
+    const resolved = resolve(path);
+    if (this.watcher) {
+      this.watcher.close();
+    }
+
+    this.watcher = watch(resolved, { recursive: true }, () => {
+      this.evict(resolved);
+      void this.get(resolved).catch(() => undefined);
+    });
+  }
+
+  close(): void {
+    this.watcher?.close();
+    this.watcher = null;
+  }
+}
+
+export async function getIndexForRepo(
+  repo: string | null | undefined,
+  defaultSource: string | null,
+  cache: IndexCache,
+  ref?: string | null,
+): Promise<MiruIndex> {
+  if (repo && isGitUrl(repo) && !repo.startsWith("https://") && !repo.startsWith("http://")) {
+    throw new Error(
+      `Only https://, http://, or local directory paths are accepted as \`repo\`. Got: ${repo}`,
+    );
+  }
+
+  const source = repo ?? defaultSource;
+  if (!source) {
+    throw new Error(
+      "No repo specified and no default index. Pass an https:// or http:// git URL or local directory path as `repo`.",
+    );
+  }
+
+  try {
+    return await cache.get(source, ref);
+  } catch (err) {
+    const message = err instanceof Error ? err.message : String(err);
+    throw new Error(`Failed to index ${source}: ${message}`);
+  }
+}
+
+export function toolText(content: string): { content: Array<{ type: "text"; text: string }> } {
+  return { content: [{ type: "text", text: content }] };
+}

package/src/mcp/serve.ts

@@ -0,0 +1,24 @@
+import { resolve } from "node:path";
+import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import type { ContentType } from "../types.ts";
+import { isGitUrl } from "../utils.ts";
+import { IndexCache } from "./index-cache.ts";
+import { createMcpServer } from "./server.ts";
+
+export async function serveMcp(options: {
+  path?: string | null;
+  ref?: string | null;
+  content?: ContentType[];
+}): Promise<void> {
+  const cache = new IndexCache(options.content ?? ["code"]);
+  const defaultSource = options.path ?? resolve(process.cwd());
+
+  void cache.get(defaultSource, options.ref).catch(() => undefined);
+  if (!isGitUrl(defaultSource)) {
+    cache.startWatcher(defaultSource);
+  }
+
+  const server = createMcpServer(cache, defaultSource);
+  const transport = new StdioServerTransport();
+  await server.connect(transport);
+}