@takara-ai/miru-code 0.1.0

package/src/mcp/server.ts

@@ -0,0 +1,100 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import * as z from "zod";
+import type { ContentType } from "../types.ts";
+import { formatResults, resolveChunk } from "../utils.ts";
+import { getIndexForRepo, type IndexCache, toolText } from "./index-cache.ts";
+
+const REPO_DESCRIPTION =
+  "https:// or http:// git URL (e.g. https://github.com/org/repo) or local directory path to index and search. " +
+  "Required when no default index was configured at startup. " +
+  "The index is cached after the first call, so repeat queries are fast.";
+
+const SERVER_INSTRUCTIONS =
+  "Instant code search for any local or remote git repository. " +
+  "Call `search` to find relevant code; call `find_related` on a result to discover similar code elsewhere. " +
+  "When working in a local project, pass the project root as `repo`. " +
+  "For remote repos, pass an explicit https:// URL. Never guess or infer URLs. " +
+  "Prefer these tools over Grep, Glob, or Read for any question about how code works.";
+
+export function createMcpServer(cache: IndexCache, defaultSource: string | null = null): McpServer {
+  const server = new McpServer(
+    {
+      name: "miru",
+      version: "0.1.0",
+    },
+    {
+      instructions: SERVER_INSTRUCTIONS,
+    },
+  );
+
+  server.registerTool(
+    "search",
+    {
+      description:
+        "Search a codebase with a natural-language or code query. " +
+        "Pass a git URL or local path as `repo` to index it on demand; indexes are cached for the session.",
+      inputSchema: {
+        query: z.string().describe("Natural language or code query."),
+        repo: z.string().nullable().optional().describe(REPO_DESCRIPTION),
+        top_k: z.number().int().min(1).optional().describe("Number of results to return."),
+      },
+    },
+    async ({ query, repo, top_k: topK }) => {
+      try {
+        const index = await getIndexForRepo(repo ?? null, defaultSource, cache);
+        const results = await index.search({ query, topK: topK ?? 5 });
+        if (results.length === 0) {
+          return toolText(JSON.stringify({ error: "No results found." }));
+        }
+        return toolText(JSON.stringify(formatResults(query, results)));
+      } catch (err) {
+        return toolText(err instanceof Error ? err.message : String(err));
+      }
+    },
+  );
+
+  server.registerTool(
+    "find_related",
+    {
+      description:
+        "Find code chunks semantically similar to a specific location in a file. " +
+        "Use after `search` to explore related implementations or callers.",
+      inputSchema: {
+        file_path: z
+          .string()
+          .describe(
+            "Path to the file as stored in the index (use file_path from a search result).",
+          ),
+        line: z.number().int().describe("Line number (1-indexed)."),
+        repo: z.string().nullable().optional().describe(REPO_DESCRIPTION),
+        top_k: z.number().int().min(1).optional().describe("Number of similar chunks to return."),
+      },
+    },
+    async ({ file_path: filePath, line, repo, top_k: topK }) => {
+      try {
+        const index = await getIndexForRepo(repo ?? null, defaultSource, cache);
+        const chunk = resolveChunk(index.chunks, filePath, line);
+        if (!chunk) {
+          return toolText(
+            `No chunk found at ${filePath}:${line}. Make sure the file is indexed and the line number is within a known chunk.`,
+          );
+        }
+        const results = await index.findRelated(chunk, topK ?? 5);
+        if (results.length === 0) {
+          return toolText(
+            JSON.stringify({ error: `No related chunks found for ${filePath}:${line}.` }),
+          );
+        }
+        return toolText(
+          JSON.stringify(formatResults(`Chunks related to ${filePath}:${line}`, results)),
+        );
+      } catch (err) {
+        return toolText(err instanceof Error ? err.message : String(err));
+      }
+    },
+  );
+
+  return server;
+}
+
+export type { ContentType };

package/src/miru-index.ts

@@ -0,0 +1,269 @@
+import { mkdir, rm } from "node:fs/promises";
+import { resolve } from "node:path";
+import { findIndexCachePath, getValidatedCache, loadCachedIndex } from "./cache.ts";
+import {
+  type EmbeddingBackend,
+  getEmbeddingBackend,
+  resolveEmbeddingDimensions,
+  resolveEmbeddingModel,
+} from "./embeddings/openai.ts";
+import { cloneGitRepository } from "./git.ts";
+import type { BM25Index } from "./index/bm25.ts";
+import { createIndexFromPath } from "./index/create.ts";
+import { persistencePaths, saveIndexBundle } from "./index/persistence.ts";
+import type { SemanticIndex } from "./index/semantic-index.ts";
+import { hybridSearch, searchSemanticOnly } from "./search.ts";
+import type { Chunk, ContentType, SearchResult } from "./types.ts";
+import { chunkKey, chunkToDict } from "./types.ts";
+import { computeSourceCacheKey, isGitUrl } from "./utils.ts";
+
+export class MiruIndex {
+  readonly embeddings: EmbeddingBackend;
+  readonly chunks: Chunk[];
+  readonly loadedFromDisk: boolean;
+
+  private readonly bm25Index: BM25Index;
+  private readonly semanticIndex: SemanticIndex;
+  private readonly embeddingModel: string;
+  private readonly root: string | null;
+  private readonly content: ContentType[];
+  private readonly fileMapping: Map<string, number[]>;
+  private readonly languageMapping: Map<string, number[]>;
+
+  constructor(options: {
+    embeddings: EmbeddingBackend;
+    bm25Index: BM25Index;
+    semanticIndex: SemanticIndex;
+    chunks: Chunk[];
+    embeddingModel: string;
+    root?: string | null;
+    content?: ContentType[];
+    loadedFromDisk?: boolean;
+  }) {
+    this.embeddings = options.embeddings;
+    this.bm25Index = options.bm25Index;
+    this.semanticIndex = options.semanticIndex;
+    this.chunks = options.chunks;
+    this.embeddingModel = options.embeddingModel;
+    this.root = options.root ?? null;
+    this.content = options.content ?? ["code"];
+    this.loadedFromDisk = options.loadedFromDisk ?? false;
+
+    this.fileMapping = new Map();
+    this.languageMapping = new Map();
+    for (let i = 0; i < this.chunks.length; i++) {
+      const chunk = this.chunks[i];
+      if (!chunk) {
+        continue;
+      }
+      const fp = chunk.file_path;
+      if (!this.fileMapping.has(fp)) {
+        this.fileMapping.set(fp, []);
+      }
+      this.fileMapping.get(fp)?.push(i);
+      if (chunk.language) {
+        if (!this.languageMapping.has(chunk.language)) {
+          this.languageMapping.set(chunk.language, []);
+        }
+        this.languageMapping.get(chunk.language)?.push(i);
+      }
+    }
+  }
+
+  static async fromGit(
+    url: string,
+    content: ContentType[] = ["code"],
+    embeddingModel?: string,
+    ref?: string | null,
+  ): Promise<MiruIndex> {
+    const cacheKey = computeSourceCacheKey(url, ref);
+    const model = embeddingModel ?? resolveEmbeddingModel();
+    const cached = await getValidatedCache(cacheKey, model, content);
+    if (cached) {
+      return MiruIndex.loadFromDisk(cached, model);
+    }
+
+    const cloneDir = await cloneGitRepository(url, ref);
+    try {
+      const embeddings = getEmbeddingBackend(model);
+      const { bm25, semantic, chunks } = await createIndexFromPath(
+        cloneDir,
+        embeddings,
+        content,
+        cloneDir,
+      );
+      const index = new MiruIndex({
+        embeddings,
+        bm25Index: bm25,
+        semanticIndex: semantic,
+        chunks,
+        embeddingModel: model,
+        root: null,
+        content,
+      });
+      await index.save(findIndexCachePath(cacheKey));
+      return index;
+    } finally {
+      await rm(cloneDir, { recursive: true, force: true });
+    }
+  }
+
+  static async fromSource(
+    source: string,
+    content: ContentType[] = ["code"],
+    embeddingModel?: string,
+    ref?: string | null,
+  ): Promise<MiruIndex> {
+    if (isGitUrl(source)) {
+      return MiruIndex.fromGit(source, content, embeddingModel, ref);
+    }
+    return MiruIndex.fromPath(source, content, embeddingModel);
+  }
+
+  static async fromPath(
+    path: string,
+    content: ContentType[] = ["code"],
+    embeddingModel?: string,
+  ): Promise<MiruIndex> {
+    const resolved = resolve(path);
+    try {
+      const st = await Bun.file(resolved).stat();
+      if (!st.isDirectory()) {
+        throw new Error(`Path is not a directory: ${path}`);
+      }
+    } catch (err) {
+      if (err instanceof Error && err.message.startsWith("Path is not a directory")) {
+        throw err;
+      }
+      throw new Error(`Path does not exist: ${path}`);
+    }
+
+    const model = embeddingModel ?? resolveEmbeddingModel();
+    const cached = await getValidatedCache(resolved, model, content);
+    if (cached) {
+      return MiruIndex.loadFromDisk(cached, model);
+    }
+
+    const embeddings = getEmbeddingBackend(model);
+    const { bm25, semantic, chunks } = await createIndexFromPath(
+      resolved,
+      embeddings,
+      content,
+      resolved,
+    );
+
+    return new MiruIndex({
+      embeddings,
+      bm25Index: bm25,
+      semanticIndex: semantic,
+      chunks,
+      embeddingModel: model,
+      root: resolved,
+      content,
+    });
+  }
+
+  static async loadFromDisk(path: string, embeddingModel?: string): Promise<MiruIndex> {
+    const model = embeddingModel ?? resolveEmbeddingModel();
+    const { bm25, semantic, chunks, metadata } = await loadCachedIndex(path);
+    const content = (metadata.content_type as ContentType[]) ?? ["code"];
+    const root = metadata.root_path ? String(metadata.root_path) : null;
+
+    return new MiruIndex({
+      embeddings: getEmbeddingBackend(model),
+      bm25Index: bm25,
+      semanticIndex: semantic,
+      chunks,
+      embeddingModel: model,
+      root,
+      content,
+      loadedFromDisk: true,
+    });
+  }
+
+  async save(path: string): Promise<void> {
+    const paths = persistencePaths(path);
+    await mkdir(paths.root, { recursive: true });
+    await saveIndexBundle({
+      paths,
+      bm25: this.bm25Index,
+      semantic: this.semanticIndex,
+      chunks: this.chunks.map(chunkToDict),
+      metadata: {
+        root_path: this.root,
+        time: Date.now() / 1000,
+        embedding_model: this.embeddingModel,
+        embedding_dimensions:
+          this.embeddings.dimensions || resolveEmbeddingDimensions(this.embeddingModel),
+        embedding_provider: "openai",
+        content_type: this.content,
+        file_paths: [...this.fileMapping.keys()].sort(),
+      },
+    });
+  }
+
+  async saveToDefaultCache(sourcePath: string): Promise<void> {
+    if (!this.loadedFromDisk) {
+      await this.save(findIndexCachePath(sourcePath));
+    }
+  }
+
+  private getSelector(filterLanguages?: string[], filterPaths?: string[]): number[] | undefined {
+    const selector: number[] = [];
+    for (const lang of filterLanguages ?? []) {
+      selector.push(...(this.languageMapping.get(lang) ?? []));
+    }
+    for (const fp of filterPaths ?? []) {
+      selector.push(...(this.fileMapping.get(fp) ?? []));
+    }
+    if (selector.length === 0) {
+      return undefined;
+    }
+    return [...new Set(selector)];
+  }
+
+  async search(options: {
+    query: string;
+    topK?: number;
+    alpha?: number | null;
+    filterLanguages?: string[];
+    filterPaths?: string[];
+    rerank?: boolean;
+  }): Promise<SearchResult[]> {
+    const { query, topK = 10, alpha, filterLanguages, filterPaths, rerank } = options;
+
+    if (!this.chunks.length || !query.trim()) {
+      return [];
+    }
+
+    const resolvedRerank = rerank ?? this.content.includes("code");
+
+    return hybridSearch({
+      query,
+      embeddings: this.embeddings,
+      semanticIndex: this.semanticIndex,
+      bm25Index: this.bm25Index,
+      chunks: this.chunks,
+      topK,
+      alpha,
+      selector: this.getSelector(filterLanguages, filterPaths),
+      rerank: resolvedRerank,
+    });
+  }
+
+  async findRelated(source: Chunk | SearchResult, topK = 5): Promise<SearchResult[]> {
+    const target = "chunk" in source ? source.chunk : source;
+    const selector = target.language ? this.getSelector([target.language]) : undefined;
+
+    const results = await searchSemanticOnly({
+      query: target.content,
+      embeddings: this.embeddings,
+      semanticIndex: this.semanticIndex,
+      chunks: this.chunks,
+      topK: topK + 1,
+      selector,
+    });
+
+    return results.filter((r) => chunkKey(r.chunk) !== chunkKey(target)).slice(0, topK);
+  }
+}

package/src/ranking/boosting.ts

@@ -0,0 +1,329 @@
+import { basename, dirname } from "node:path";
+import { splitIdentifier } from "../tokens.ts";
+import type { Chunk } from "../types.ts";
+import { chunkKey } from "../types.ts";
+
+const SYMBOL_QUERY_RE =
+  /^(?:(?:[A-Za-z_][A-Za-z0-9_]*(?:(?:::|\\|->|\.)[A-Za-z_][A-Za-z0-9_]*)+)|_(?:[A-Za-z0-9_]*)|(?:[A-Za-z][A-Za-z0-9]*[A-Z_][A-Za-z0-9_]*)|(?:[A-Z][A-Za-z0-9]*))$/;
+
+const EMBEDDED_SYMBOL_RE =
+  /\b(?:[A-Z][a-z][a-zA-Z0-9]*[A-Z][a-zA-Z0-9]*|[a-z][a-zA-Z0-9]*[A-Z][a-zA-Z0-9]+)\b/g;
+
+const EMBEDDED_STEM_MIN_LEN = 4;
+const EMBEDDED_SYMBOL_BOOST_SCALE = 0.5;
+
+const DEFINITION_KEYWORDS = [
+  "class",
+  "module",
+  "defmodule",
+  "def",
+  "interface",
+  "struct",
+  "enum",
+  "trait",
+  "type",
+  "func",
+  "function",
+  "object",
+  "abstract class",
+  "data class",
+  "fn",
+  "fun",
+  "package",
+  "namespace",
+  "protocol",
+  "record",
+  "typedef",
+];
+
+const SQL_DEFINITION_KEYWORDS = [
+  "CREATE TABLE",
+  "CREATE VIEW",
+  "CREATE PROCEDURE",
+  "CREATE FUNCTION",
+];
+
+const DEFINITION_BOOST_MULTIPLIER = 3.0;
+const STEM_BOOST_MULTIPLIER = 1.0;
+const FILE_COHERENCE_BOOST_FRAC = 0.2;
+
+const STOPWORDS = new Set(
+  "a an and are as at be by do does for from has have how if in is it not of on or the to was what when where which who why with".split(
+    " ",
+  ),
+);
+
+const definitionPatternCache = new Map<string, [RegExp, RegExp]>();
+
+function definitionPatterns(symbolName: string): [RegExp, RegExp] {
+  let cached = definitionPatternCache.get(symbolName);
+  if (cached) {
+    return cached;
+  }
+  const escaped = symbolName.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
+  const nsPrefix = "(?:[A-Za-z_][A-Za-z0-9_]*(?:\\.|::))*";
+  const suffix = `)\\s+${nsPrefix}${escaped}(?:\\s|[<({:\\[;]|$)`;
+  const keywordBody = DEFINITION_KEYWORDS.map((k) => k.replace(/ /g, "\\s+")).join("|");
+  const sqlBody = SQL_DEFINITION_KEYWORDS.join("|");
+  const keywordPrefix = "(?:^|(?<=\\s))(?:";
+  cached = [
+    new RegExp(keywordPrefix + keywordBody + suffix, "m"),
+    new RegExp(keywordPrefix + sqlBody + suffix, "im"),
+  ];
+  definitionPatternCache.set(symbolName, cached);
+  return cached;
+}
+
+function chunkDefinesSymbol(chunk: Chunk, symbolName: string): boolean {
+  const [general, sql] = definitionPatterns(symbolName);
+  return general.test(chunk.content) || sql.test(chunk.content);
+}
+
+function stemMatches(stem: string, name: string): boolean {
+  const stemNorm = stem.replace(/_/g, "");
+  return (
+    stem === name ||
+    stemNorm === name ||
+    stem.replace(/s$/, "") === name ||
+    stemNorm.replace(/s$/, "") === name
+  );
+}
+
+export function isSymbolQuery(query: string): boolean {
+  return SYMBOL_QUERY_RE.test(query.trim());
+}
+
+function extractSymbolName(query: string): string {
+  for (const sep of ["::", "\\", "->", "."]) {
+    if (query.includes(sep)) {
+      return query.split(sep).pop()?.trim() ?? query.trim();
+    }
+  }
+  return query.trim();
+}
+
+function definitionTier(chunk: Chunk, names: Set<string>, boostUnit: number): number {
+  if (![...names].some((n) => chunkDefinesSymbol(chunk, n))) {
+    return 0;
+  }
+  const stem = basename(chunk.file_path)
+    .replace(/\.[^.]+$/, "")
+    .toLowerCase();
+  return boostUnit * ([...names].some((n) => stemMatches(stem, n.toLowerCase())) ? 1.5 : 1.0);
+}
+
+export function boostMultiChunkFiles(
+  scores: Map<string, number>,
+  chunksByKey: Map<string, Chunk>,
+): void {
+  if (scores.size === 0) {
+    return;
+  }
+  const maxScore = Math.max(...scores.values());
+  if (maxScore === 0) {
+    return;
+  }
+
+  const fileSum = new Map<string, number>();
+  const bestChunk = new Map<string, string>();
+
+  for (const [key, score] of scores) {
+    const chunk = chunksByKey.get(key);
+    if (!chunk) {
+      continue;
+    }
+    const fp = chunk.file_path;
+    fileSum.set(fp, (fileSum.get(fp) ?? 0) + score);
+    const best = bestChunk.get(fp);
+    if (!best || score > (scores.get(best) ?? 0)) {
+      bestChunk.set(fp, key);
+    }
+  }
+
+  const maxFileSum = Math.max(...fileSum.values());
+  const boostUnit = maxScore * FILE_COHERENCE_BOOST_FRAC;
+
+  for (const [fp, key] of bestChunk) {
+    const current = scores.get(key) ?? 0;
+    scores.set(key, current + boostUnit * ((fileSum.get(fp) ?? 0) / maxFileSum));
+  }
+}
+
+export function applyQueryBoost(
+  combinedScores: Map<string, number>,
+  query: string,
+  allChunks: Chunk[],
+  chunksByKey: Map<string, Chunk>,
+): Map<string, number> {
+  if (combinedScores.size === 0) {
+    return combinedScores;
+  }
+
+  const maxScore = Math.max(...combinedScores.values());
+  if (isSymbolQuery(query)) {
+    boostSymbolDefinitions(combinedScores, query, maxScore, allChunks, chunksByKey);
+  } else {
+    boostStemMatches(combinedScores, query, maxScore, chunksByKey);
+    boostEmbeddedSymbols(combinedScores, query, maxScore, allChunks, chunksByKey);
+  }
+  return combinedScores;
+}
+
+function boostSymbolDefinitions(
+  boosted: Map<string, number>,
+  query: string,
+  maxScore: number,
+  allChunks: Chunk[],
+  chunksByKey: Map<string, Chunk>,
+): void {
+  const symbolName = extractSymbolName(query);
+  const names = new Set([symbolName]);
+  if (symbolName !== query.trim()) {
+    names.add(query.trim());
+  }
+  const boostUnit = maxScore * DEFINITION_BOOST_MULTIPLIER;
+
+  for (const key of [...boosted.keys()]) {
+    const chunk = chunksByKey.get(key);
+    if (!chunk) {
+      continue;
+    }
+    const tier = definitionTier(chunk, names, boostUnit);
+    if (tier) {
+      boosted.set(key, (boosted.get(key) ?? 0) + tier);
+    }
+  }
+
+  for (const chunk of allChunks) {
+    const key = chunkKey(chunk);
+    if (boosted.has(key)) {
+      continue;
+    }
+    const stem = basename(chunk.file_path)
+      .replace(/\.[^.]+$/, "")
+      .toLowerCase();
+    if (![...names].some((n) => stemMatches(stem, n.toLowerCase()))) {
+      continue;
+    }
+    const tier = definitionTier(chunk, names, boostUnit);
+    if (tier) {
+      boosted.set(key, tier);
+    }
+  }
+}
+
+function boostEmbeddedSymbols(
+  boosted: Map<string, number>,
+  query: string,
+  maxScore: number,
+  allChunks: Chunk[],
+  chunksByKey: Map<string, Chunk>,
+): void {
+  const names = new Set(query.match(EMBEDDED_SYMBOL_RE) ?? []);
+  if (names.size === 0) {
+    return;
+  }
+
+  const boostUnit = maxScore * DEFINITION_BOOST_MULTIPLIER * EMBEDDED_SYMBOL_BOOST_SCALE;
+  const symbolsLower = new Set([...names].map((s) => s.toLowerCase()));
+
+  for (const key of [...boosted.keys()]) {
+    const chunk = chunksByKey.get(key);
+    if (!chunk) {
+      continue;
+    }
+    const tier = definitionTier(chunk, names, boostUnit);
+    if (tier) {
+      boosted.set(key, (boosted.get(key) ?? 0) + tier);
+    }
+  }
+
+  for (const chunk of allChunks) {
+    const key = chunkKey(chunk);
+    if (boosted.has(key)) {
+      continue;
+    }
+    const stem = basename(chunk.file_path)
+      .replace(/\.[^.]+$/, "")
+      .toLowerCase();
+    const stemNorm = stem.replace(/_/g, "");
+    const ok = [...symbolsLower].some(
+      (symbolLower) =>
+        stem === symbolLower ||
+        stemNorm === symbolLower ||
+        (stem.length >= EMBEDDED_STEM_MIN_LEN && symbolLower.startsWith(stem)) ||
+        (stemNorm.length >= EMBEDDED_STEM_MIN_LEN && symbolLower.startsWith(stemNorm)),
+    );
+    if (!ok) {
+      continue;
+    }
+    const tier = definitionTier(chunk, names, boostUnit);
+    if (tier) {
+      boosted.set(key, tier);
+    }
+  }
+}
+
+function countKeywordMatches(keywords: Set<string>, parts: Set<string>): number {
+  const exact = [...keywords].filter((k) => parts.has(k));
+  if (exact.length === keywords.size) {
+    return exact.length;
+  }
+  let nMatches = exact.length;
+  for (const keyword of [...keywords].filter((k) => !parts.has(k))) {
+    for (const part of parts) {
+      const shorter = keyword.length <= part.length ? keyword : part;
+      const longer = keyword.length <= part.length ? part : keyword;
+      if (shorter.length >= 3 && longer.startsWith(shorter)) {
+        nMatches++;
+        break;
+      }
+    }
+  }
+  return nMatches;
+}
+
+function boostStemMatches(
+  boosted: Map<string, number>,
+  query: string,
+  maxScore: number,
+  chunksByKey: Map<string, Chunk>,
+): void {
+  const keywords = new Set(
+    (query.match(/[a-zA-Z_][a-zA-Z0-9_]*/g) ?? [])
+      .filter((w) => w.length > 2 && !STOPWORDS.has(w.toLowerCase()))
+      .map((w) => w.toLowerCase()),
+  );
+  if (keywords.size === 0) {
+    return;
+  }
+
+  const boost = maxScore * STEM_BOOST_MULTIPLIER;
+  const pathCache = new Map<string, Set<string>>();
+
+  for (const key of [...boosted.keys()]) {
+    const chunk = chunksByKey.get(key);
+    if (!chunk) {
+      continue;
+    }
+    let parts = pathCache.get(chunk.file_path);
+    if (!parts) {
+      const stem = basename(chunk.file_path).replace(/\.[^.]+$/, "");
+      parts = new Set(splitIdentifier(stem));
+      const parent = basename(dirname(chunk.file_path));
+      if (parent && parent !== "." && parent !== "..") {
+        for (const p of splitIdentifier(parent)) {
+          parts.add(p);
+        }
+      }
+      pathCache.set(chunk.file_path, parts);
+    }
+    const nMatches = countKeywordMatches(keywords, parts);
+    if (nMatches > 0) {
+      const matchRatio = nMatches / keywords.size;
+      if (matchRatio >= 0.1) {
+        boosted.set(key, (boosted.get(key) ?? 0) + boost * matchRatio);
+      }
+    }
+  }
+}