@takara-ai/miru-code 0.1.0

package/src/ranking/penalties.ts

@@ -0,0 +1,110 @@
+import { basename } from "node:path";
+import type { Chunk } from "../types.ts";
+
+const TEST_FILE_RE =
+  /(?:^|\/)(?:test_[^/]*\.py|[^/]*_test\.py|[^/]*_test\.go|[^/]*Tests?\.java|[^/]*Test\.php|[^/]*_spec\.rb|[^/]*_test\.rb|[^/]*\.test\.[jt]sx?|[^/]*\.spec\.[jt]sx?|[^/]*Tests?\.kt|[^/]*Spec\.kt|[^/]*Tests?\.swift|[^/]*Spec\.swift|[^/]*Tests?\.cs|test_[^/]*\.cpp|[^/]*_test\.cpp|test_[^/]*\.c|[^/]*_test\.c|[^/]*Spec\.scala|[^/]*Suite\.scala|[^/]*Test\.scala|[^/]*_test\.dart|test_[^/]*\.dart|[^/]*_spec\.lua|[^/]*_test\.lua|test_[^/]*\.lua|test_helpers?[^/]*\.\w+)$/;
+
+const TEST_DIR_RE = /(?:^|\/)(?:tests?|__tests__|spec|testing)(?:\/|$)/;
+const COMPAT_DIR_RE = /(?:^|\/)(?:compat|_compat|legacy)(?:\/|$)/;
+const EXAMPLES_DIR_RE = /(?:^|\/)(?:_?examples?|docs?_src)(?:\/|$)/;
+const TYPE_DEFS_RE = /\.d\.ts$/;
+
+const STRONG_PENALTY = 0.3;
+const MODERATE_PENALTY = 0.5;
+const MILD_PENALTY = 0.7;
+
+const REEXPORT_FILENAMES = new Set(["__init__.py", "package-info.java"]);
+const FILE_SATURATION_THRESHOLD = 1;
+const FILE_SATURATION_DECAY = 0.5;
+
+function filePathPenalty(filePath: string): number {
+  const normalised = filePath.replace(/\\/g, "/");
+  let penalty = 1.0;
+  if (TEST_FILE_RE.test(normalised) || TEST_DIR_RE.test(normalised)) {
+    penalty *= STRONG_PENALTY;
+  }
+  if (REEXPORT_FILENAMES.has(basename(filePath))) {
+    penalty *= MODERATE_PENALTY;
+  }
+  if (COMPAT_DIR_RE.test(normalised)) {
+    penalty *= STRONG_PENALTY;
+  }
+  if (EXAMPLES_DIR_RE.test(normalised)) {
+    penalty *= STRONG_PENALTY;
+  }
+  if (TYPE_DEFS_RE.test(normalised)) {
+    penalty *= MILD_PENALTY;
+  }
+  return penalty;
+}
+
+export function rerankTopk(
+  scores: Map<string, number>,
+  chunksByKey: Map<string, Chunk>,
+  topK: number,
+  penalisePaths = true,
+): [Chunk, number][] {
+  if (scores.size === 0) {
+    return [];
+  }
+
+  const penaltyCache = new Map<string, number>();
+  const penalised = new Map<string, number>();
+
+  for (const [key, score] of scores) {
+    const chunk = chunksByKey.get(key);
+    if (!chunk) {
+      continue;
+    }
+    let pen = score;
+    if (penalisePaths) {
+      let mult = penaltyCache.get(chunk.file_path);
+      if (mult === undefined) {
+        mult = filePathPenalty(chunk.file_path);
+        penaltyCache.set(chunk.file_path, mult);
+      }
+      pen *= mult;
+    }
+    penalised.set(key, pen);
+  }
+
+  const ranked = [...penalised.entries()].sort((a, b) => b[1] - a[1]);
+
+  const fileSelected = new Map<string, number>();
+  const selected: { score: number; key: string }[] = [];
+  let minSelected = Infinity;
+
+  for (const [key, penScore] of ranked) {
+    const chunk = chunksByKey.get(key);
+    if (!chunk) {
+      continue;
+    }
+
+    if (selected.length >= topK && penScore <= minSelected) {
+      break;
+    }
+
+    const already = fileSelected.get(chunk.file_path) ?? 0;
+    let effScore = penScore;
+    if (already >= FILE_SATURATION_THRESHOLD) {
+      const excess = already - FILE_SATURATION_THRESHOLD + 1;
+      effScore *= FILE_SATURATION_DECAY ** excess;
+    }
+
+    selected.push({ score: effScore, key });
+    fileSelected.set(chunk.file_path, already + 1);
+
+    if (selected.length >= topK) {
+      minSelected = Math.min(...selected.map((s) => s.score));
+    }
+  }
+
+  selected.sort((a, b) => b.score - a.score);
+  return selected.slice(0, topK).flatMap(({ key, score }) => {
+    const chunk = chunksByKey.get(key);
+    if (!chunk) {
+      return [];
+    }
+    return [[chunk, score] as [Chunk, number]];
+  });
+}

package/src/ranking/weighting.ts

@@ -0,0 +1,11 @@
+import { isSymbolQuery } from "./boosting.ts";
+
+const ALPHA_SYMBOL = 0.3;
+const ALPHA_NL = 0.5;
+
+export function resolveAlpha(query: string, alpha: number | null | undefined): number {
+  if (alpha != null) {
+    return alpha;
+  }
+  return isSymbolQuery(query) ? ALPHA_SYMBOL : ALPHA_NL;
+}

package/src/search.ts

@@ -0,0 +1,180 @@
+import type { EmbeddingBackend } from "./embeddings/openai.ts";
+import type { BM25Index } from "./index/bm25.ts";
+import type { SemanticIndex } from "./index/semantic-index.ts";
+import { selectorToMask } from "./index/sparse.ts";
+import { selectTopKScoreIndices } from "./index/topk.ts";
+import { applyQueryBoost, boostMultiChunkFiles } from "./ranking/boosting.ts";
+import { rerankTopk } from "./ranking/penalties.ts";
+import { resolveAlpha } from "./ranking/weighting.ts";
+import { tokenize } from "./tokens.ts";
+import type { Chunk, SearchResult } from "./types.ts";
+import { chunkKey } from "./types.ts";
+
+const RRF_K = 60;
+
+function rrfScores(scores: Map<string, number>): Map<string, number> {
+  if (scores.size === 0) {
+    return scores;
+  }
+  const ranked = [...scores.entries()].sort((a, b) => b[1] - a[1]);
+  const out = new Map<string, number>();
+  ranked.forEach(([key], i) => {
+    out.set(key, 1.0 / (RRF_K + i + 1));
+  });
+  return out;
+}
+
+function sortTopK(arr: number[], topK: number): number[] {
+  return selectTopKScoreIndices(arr, topK);
+}
+
+function semanticFromQueryVector(
+  queryVec: Float32Array,
+  semanticIndex: SemanticIndex,
+  chunks: Chunk[],
+  topK: number,
+  selector?: number[],
+): SearchResult[] {
+  const { indices, distances } = semanticIndex.query(queryVec, topK, selector);
+  return indices.flatMap((index, i) => {
+    const chunk = chunks[index];
+    if (!chunk) {
+      return [];
+    }
+    return [{ chunk, score: 1.0 - (distances[i] ?? 0) }];
+  });
+}
+
+async function searchBm25(
+  query: string,
+  bm25Index: BM25Index,
+  chunks: Chunk[],
+  topK: number,
+  selector?: number[],
+): Promise<SearchResult[]> {
+  const tokens = tokenize(query);
+  if (tokens.length === 0) {
+    return [];
+  }
+  const mask = selectorToMask(selector, chunks.length);
+  const scores = await bm25Index.getScoresAsync(tokens, mask);
+  const indices = sortTopK(scores, topK);
+  return indices.flatMap((i) => {
+    const score = scores[i];
+    const chunk = chunks[i];
+    if (score === undefined || chunk === undefined || score <= 0) {
+      return [];
+    }
+    return [{ chunk, score }];
+  });
+}
+
+export async function hybridSearch(options: {
+  query: string;
+  embeddings: EmbeddingBackend;
+  semanticIndex: SemanticIndex;
+  bm25Index: BM25Index;
+  chunks: Chunk[];
+  topK: number;
+  alpha?: number | null;
+  selector?: number[];
+  rerank?: boolean;
+}): Promise<SearchResult[]> {
+  const {
+    query,
+    embeddings,
+    semanticIndex,
+    bm25Index,
+    chunks,
+    topK,
+    alpha,
+    selector,
+    rerank = true,
+  } = options;
+
+  const alphaWeight = resolveAlpha(query, alpha);
+  const candidateCount = topK * 5;
+  const chunksByKey = new Map(chunks.map((c) => [chunkKey(c), c]));
+
+  const queryVecPromise = embeddings.embedQuery(query);
+  const bm25Promise = searchBm25(query, bm25Index, chunks, candidateCount, selector);
+  const queryVec = await queryVecPromise;
+  const bm25Hits = await bm25Promise;
+  const semantic = semanticFromQueryVector(
+    queryVec,
+    semanticIndex,
+    chunks,
+    candidateCount,
+    selector,
+  );
+
+  const semanticScores = new Map<string, number>();
+  for (const r of semantic) {
+    semanticScores.set(chunkKey(r.chunk), r.score);
+  }
+
+  const bm25Scores = new Map<string, number>();
+  for (const r of bm25Hits) {
+    if (r.score) {
+      bm25Scores.set(chunkKey(r.chunk), r.score);
+    }
+  }
+
+  const normalizedSemantic = rrfScores(semanticScores);
+  const normalizedBm25 = rrfScores(bm25Scores);
+
+  const allKeys = new Set([...normalizedSemantic.keys(), ...normalizedBm25.keys()]);
+  const sortedKeys = [...allKeys].sort((a, b) => {
+    const ca = chunksByKey.get(a);
+    const cb = chunksByKey.get(b);
+    if (!ca || !cb) {
+      return 0;
+    }
+    return ca.start_line - cb.start_line;
+  });
+
+  const combinedScores = new Map<string, number>();
+  for (const key of sortedKeys) {
+    combinedScores.set(
+      key,
+      alphaWeight * (normalizedSemantic.get(key) ?? 0) +
+        (1 - alphaWeight) * (normalizedBm25.get(key) ?? 0),
+    );
+  }
+
+  if (rerank) {
+    boostMultiChunkFiles(combinedScores, chunksByKey);
+    applyQueryBoost(combinedScores, query, chunks, chunksByKey);
+    const ranked = rerankTopk(combinedScores, chunksByKey, topK, alphaWeight < 1.0);
+    return ranked.map(([chunk, score]) => ({ chunk, score }));
+  }
+
+  return [...combinedScores.entries()]
+    .sort((a, b) => b[1] - a[1])
+    .slice(0, topK)
+    .flatMap(([key, score]) => {
+      const chunk = chunksByKey.get(key);
+      if (!chunk) {
+        return [];
+      }
+      return [{ chunk, score }];
+    });
+}
+
+export async function searchSemanticOnly(options: {
+  query: string;
+  embeddings: EmbeddingBackend;
+  semanticIndex: SemanticIndex;
+  chunks: Chunk[];
+  topK: number;
+  selector?: number[];
+}): Promise<SearchResult[]> {
+  const queryVec = await options.embeddings.embedQuery(options.query);
+  return semanticFromQueryVector(
+    queryVec,
+    options.semanticIndex,
+    options.chunks,
+    options.topK,
+    options.selector,
+  );
+}

package/src/tokens.ts

@@ -0,0 +1,27 @@
+const TOKEN_RE = /[a-zA-Z_][a-zA-Z0-9_]*/g;
+const CAMEL_RE = /[A-Z]+(?=[A-Z][a-z])|[A-Z]?[a-z]+|[A-Z]+|[0-9]+/g;
+
+export function splitIdentifier(token: string): string[] {
+  const lower = token.toLowerCase();
+  let parts: string[];
+
+  if (token.includes("_")) {
+    parts = lower.split("_").filter(Boolean);
+  } else {
+    parts = (token.match(CAMEL_RE) ?? []).map((m) => m.toLowerCase());
+  }
+
+  if (parts.length >= 2) {
+    return [lower, ...parts];
+  }
+  return [lower];
+}
+
+export function tokenize(text: string): string[] {
+  const raw = text.match(TOKEN_RE) ?? [];
+  const result: string[] = [];
+  for (const tok of raw) {
+    result.push(...splitIdentifier(tok));
+  }
+  return result;
+}

package/src/types.ts

@@ -0,0 +1,47 @@
+export type ContentType = "code" | "docs" | "config";
+
+export interface Chunk {
+  content: string;
+  file_path: string;
+  start_line: number;
+  end_line: number;
+  language: string | null;
+}
+
+export interface SearchResult {
+  chunk: Chunk;
+  score: number;
+}
+
+export function chunkKey(chunk: Chunk): string {
+  return `${chunk.file_path}:${chunk.start_line}:${chunk.end_line}`;
+}
+
+export function chunkToDict(chunk: Chunk): Record<string, unknown> {
+  return {
+    content: chunk.content,
+    file_path: chunk.file_path,
+    start_line: chunk.start_line,
+    end_line: chunk.end_line,
+    language: chunk.language,
+    location: `${chunk.file_path}:${chunk.start_line}-${chunk.end_line}`,
+  };
+}
+
+export function chunkFromDict(data: Record<string, unknown>): Chunk {
+  const { location: _loc, ...rest } = data;
+  return {
+    content: String(rest.content),
+    file_path: String(rest.file_path),
+    start_line: Number(rest.start_line),
+    end_line: Number(rest.end_line),
+    language: rest.language == null ? null : String(rest.language),
+  };
+}
+
+export function searchResultToDict(result: SearchResult): Record<string, unknown> {
+  return {
+    chunk: chunkToDict(result.chunk),
+    score: result.score,
+  };
+}

package/src/utils.ts

@@ -0,0 +1,64 @@
+import { resolve } from "node:path";
+import type { Chunk, ContentType, SearchResult } from "./types.ts";
+import { searchResultToDict } from "./types.ts";
+
+const GIT_URL_SCHEMES = [
+  "https://",
+  "http://",
+  "ssh://",
+  "git://",
+  "git+ssh://",
+  "file://",
+] as const;
+const SCP_GIT_URL_RE = /^[\w.-]+@[\w.-]+:(?!\/)/;
+
+export function isGitUrl(path: string): boolean {
+  return GIT_URL_SCHEMES.some((scheme) => path.startsWith(scheme)) || SCP_GIT_URL_RE.test(path);
+}
+
+export function isAllowedRepoSource(repo: string): boolean {
+  if (!isGitUrl(repo)) {
+    return true;
+  }
+  return repo.startsWith("https://") || repo.startsWith("http://");
+}
+
+export function resolveChunk(chunks: Chunk[], filePath: string, line: number): Chunk | null {
+  let fallback: Chunk | null = null;
+  for (const chunk of chunks) {
+    if (chunk.file_path === filePath && chunk.start_line <= line && line <= chunk.end_line) {
+      if (line < chunk.end_line) {
+        return chunk;
+      }
+      if (fallback === null) {
+        fallback = chunk;
+      }
+    }
+  }
+  return fallback;
+}
+
+export function formatResults(
+  query: string,
+  results: SearchResult[],
+): { query: string; results: Record<string, unknown>[] } {
+  return {
+    query,
+    results: results.map(searchResultToDict),
+  };
+}
+
+export function resolveContent(raw: string[]): ContentType[] {
+  if (raw.includes("all")) {
+    return ["code", "docs", "config"];
+  }
+  const valid: ContentType[] = ["code", "docs", "config"];
+  return raw.filter((item): item is ContentType => valid.includes(item as ContentType));
+}
+
+export function computeSourceCacheKey(source: string, ref?: string | null): string {
+  if (isGitUrl(source)) {
+    return ref ? `${source}@${ref}` : source;
+  }
+  return resolve(source);
+}