karst 0.1.0__py3-none-any.whl

karst/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

karst/analyze.py

@@ -0,0 +1,39 @@
+"""End-to-end analyze pipeline: walk → parse → chunk.
+
+Holds the public surface that the CLI (and later, agents) will call.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+from .chunker import chunk_file
+from .models import Chunk
+from .parser import ParsedFile, ParserRegistry, parse_file
+from .walker import iter_source_files
+
+
+@dataclass
+class FileResult:
+    parsed: ParsedFile
+    chunks: list[Chunk]
+
+
+def analyze_repo(root: str | Path) -> Iterator[FileResult]:
+    """Iterate over every supported source file under `root`, yielding the
+    parsed file + its extracted chunks.
+
+    Streaming — callers can write JSONL as it flows, without holding the
+    whole repo in memory.
+    """
+    root_path = Path(root).resolve()
+    registry = ParserRegistry()
+
+    for file_path in iter_source_files(root_path):
+        parsed = parse_file(file_path, repo_root=root_path, registry=registry)
+        if parsed is None:
+            continue
+        chunks = chunk_file(parsed)
+        yield FileResult(parsed=parsed, chunks=chunks)

karst/ask.py

@@ -0,0 +1,149 @@
+"""Repo Q&A — question in, cited answer out.
+
+Pipeline:
+  question → embed → Qdrant top-k → assemble prompt → LLM → answer
+
+Citation discipline (spec §33): the prompt forces the model to anchor every
+claim to `file:start-end`. If no LLM is configured, callers can render the
+retrieved hits directly — still useful, just no prose synthesis.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from .embedder import DEFAULT_MODEL, Embedder
+from .llm import LLM, LLMResponse, default_llm
+from .store import DEFAULT_COLLECTION, ChunkStore, SearchHit
+
+
+@dataclass
+class _LabeledHit:
+    """Internal: a SearchHit with a source label ('vector' or 'graph')."""
+    hit: SearchHit
+    source: str
+
+# How long any single retrieved chunk is allowed to be inside the prompt.
+# Beyond this we cut — the citation still points at the full file.
+_MAX_CHUNK_CHARS = 2_000
+
+
+@dataclass
+class AskResult:
+    question: str
+    hits: list[SearchHit]
+    answer: str | None
+    llm: LLMResponse | None
+
+
+SYSTEM_PROMPT = """\
+You are an AI Staff Engineer answering questions about a specific code repository.
+
+You are given the user's question and the top retrieved chunks of code from
+the repo. Each chunk header looks like [N] path/to/file.ts:start-end. You must:
+
+1. Answer concisely and concretely.
+2. Cite every claim with a bracketed reference in the form [path/to/file.ts:start-end].
+   Use the same path/range shown in the chunk header. Never invent files or line ranges.
+3. If the retrieved chunks do not contain enough information, say so plainly and
+   suggest what to look at next. Do not guess.
+4. Prefer evidence from the retrieved chunks over background knowledge.
+"""
+
+
+def ask(
+    question: str,
+    *,
+    storage_path: str | Path,
+    collection: str = DEFAULT_COLLECTION,
+    embedding_model: str = DEFAULT_MODEL,
+    embedder_cache_dir: str | Path | None = None,
+    top_k: int = 8,
+    llm: LLM | None = None,
+    use_llm: bool = True,
+    graph_path: str | Path | None = None,
+    graph_extra: int = 6,
+    pack_ids: list[str] | None = None,
+) -> AskResult:
+    """Question → embed → Qdrant top-k → (optional graph expansion) → LLM.
+
+    When `pack_ids` is provided, retrieval is scoped to chunks tagged with
+    any of those packs (spec §22). This is the single largest token-cost
+    lever in Phase 4 — 60-80% input reduction on big repos.
+    """
+    embedder = Embedder(
+        embedding_model,
+        cache_dir=str(embedder_cache_dir) if embedder_cache_dir else None,
+    )
+    store = ChunkStore(location=storage_path, collection=collection)
+    try:
+        (query_vec,) = embedder.embed_texts([question])
+        seed_hits = store.search(query_vec, limit=top_k, pack_ids=pack_ids)
+
+        if graph_path is not None:
+            hits = _expand_with_graph(seed_hits, graph_path, store, extra=graph_extra)
+        else:
+            hits = seed_hits
+    finally:
+        store.close()
+
+    if not use_llm:
+        return AskResult(question=question, hits=hits, answer=None, llm=None)
+
+    used_llm = llm or default_llm()
+    user_prompt = _build_user_prompt(question, hits)
+    resp = used_llm.generate(SYSTEM_PROMPT, user_prompt)
+    return AskResult(question=question, hits=hits, answer=resp.text, llm=resp)
+
+
+def _expand_with_graph(
+    seed_hits: list[SearchHit],
+    graph_path: str | Path,
+    qdrant: ChunkStore,
+    *,
+    extra: int,
+) -> list[SearchHit]:
+    """Lazy import so plain `ask` doesn't pay for networkx unnecessarily."""
+    from .graph.graphrag import expand_with_graph
+    from .graph.store import GraphStore
+
+    graph = GraphStore.load(graph_path)
+    expanded = expand_with_graph(
+        seed_hits,
+        graph=graph,
+        qdrant=qdrant,
+        max_extra=extra,
+    )
+    # Collapse back into SearchHit list so the downstream prompt builder
+    # doesn't need to learn a new type. Source label is encoded in the score
+    # rank order; graph hits will already be lower-scored than seeds.
+    return [SearchHit(chunk=h.chunk, score=h.score) for h in expanded]
+
+
+def _build_user_prompt(question: str, hits: list[SearchHit]) -> str:
+    if not hits:
+        return (
+            "No chunks were retrieved from the index for this question.\n\n"
+            f"Question: {question}\n\n"
+            "Tell the user the index is empty or the question matches nothing, "
+            "and recommend re-running `karst index <path>` or rephrasing."
+        )
+
+    parts: list[str] = ["# Retrieved chunks", ""]
+    for i, hit in enumerate(hits, start=1):
+        c = hit.chunk
+        code = c.code
+        if len(code) > _MAX_CHUNK_CHARS:
+            code = code[:_MAX_CHUNK_CHARS] + "\n… (truncated)"
+        parts.append(
+            f"[{i}] {c.citation}  "
+            f"({c.kind.value} {c.qualified_name}, score={hit.score:.3f})"
+        )
+        parts.append(f"```{c.language}")
+        parts.append(code)
+        parts.append("```")
+        parts.append("")
+    parts.append("# User question")
+    parts.append(question)
+    return "\n".join(parts)

karst/chunker.py

@@ -0,0 +1,205 @@
+"""AST-aware chunker.
+
+Given a ParsedFile, walks the tree-sitter tree and emits Chunk objects, one
+per function / class / method / interface / etc. Each chunk preserves its
+exact byte range and line range so it doubles as a citation.
+
+Design notes:
+- The chunker is intentionally one-pass and stateless beyond the parent stack.
+- We DO emit both a container (e.g. class) and its members (methods) — the
+  container chunk gives architectural shape; the member chunks give the
+  retrieval-friendly units the spec calls for ("each chunk is a complete
+  function, class, or top-level statement"; spec §7).
+- "decorated_definition" in Python wraps the real function/class. We treat
+  the decorated form as the chunk and skip the inner duplicate.
+
+API note:
+- The tree-sitter Python bindings shipped with tree-sitter-language-pack on
+  Windows expose Node attributes as methods (e.g. `node.kind()`,
+  `node.child_count()`, `node.start_position().row`). The helpers in this
+  module call those methods directly rather than the property-style API of
+  upstream py-tree-sitter, so we work with whichever wheel is installed.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+from .languages import LanguageSpec, get_language
+from .models import Chunk, ChunkKind
+from .parser import ParsedFile
+
+
+_SIGNATURE_MAX_BYTES = 240
+
+
+def chunk_file(parsed: ParsedFile) -> list[Chunk]:
+    """Extract AST-aware chunks from a parsed file."""
+    lang = get_language(parsed.language)
+    if lang is None or not lang.chunk_nodes:
+        return []
+
+    chunks: list[Chunk] = []
+    root = parsed.tree.root_node()
+    _walk(root, lang, parsed, parent_qname=None, out=chunks, skip_children_of=set())
+    return chunks
+
+
+def _iter_children(node) -> Iterator:
+    count = node.child_count()
+    for i in range(count):
+        yield node.child(i)
+
+
+def _walk(
+    node,
+    lang: LanguageSpec,
+    parsed: ParsedFile,
+    *,
+    parent_qname: str | None,
+    out: list[Chunk],
+    skip_children_of: set[int],
+) -> None:
+    """Recursively walk the tree, emitting chunks for chunkable nodes.
+
+    `skip_children_of` carries Python object ids of nodes whose chunkable
+    descendants we've already processed via a wrapper (e.g.
+    decorated_definition wraps function_definition; we don't want both).
+    """
+    for child in _iter_children(node):
+        child_kind = child.kind()
+        if id(child) in skip_children_of:
+            continue
+
+        chunk_kind = lang.chunk_nodes.get(child_kind)
+        if chunk_kind is not None:
+            chunk = _emit_chunk(child, chunk_kind, lang, parsed, parent_qname=parent_qname)
+            next_parent = chunk.qualified_name if chunk is not None else parent_qname
+            if chunk is not None:
+                out.append(chunk)
+
+            # Python decorated_definition wraps function_definition /
+            # class_definition. Mark the inner node so we don't double-emit.
+            if child_kind == "decorated_definition":
+                for grand in _iter_children(child):
+                    if grand.kind() in {"function_definition", "class_definition"}:
+                        skip_children_of.add(id(grand))
+
+            if child_kind in lang.container_nodes:
+                _walk(
+                    child,
+                    lang,
+                    parsed,
+                    parent_qname=next_parent,
+                    out=out,
+                    skip_children_of=skip_children_of,
+                )
+        else:
+            # Not a chunk node; keep descending — methods may be wrapped in a
+            # class_body / declaration_list node we don't emit ourselves.
+            _walk(
+                child,
+                lang,
+                parsed,
+                parent_qname=parent_qname,
+                out=out,
+                skip_children_of=skip_children_of,
+            )
+
+
+def _emit_chunk(
+    node,
+    kind: ChunkKind,
+    lang: LanguageSpec,
+    parsed: ParsedFile,
+    *,
+    parent_qname: str | None,
+) -> Chunk | None:
+    name = _extract_name(node, lang, parsed.source)
+    if name is None:
+        return None
+
+    if kind == ChunkKind.FUNCTION and parent_qname is not None:
+        kind = ChunkKind.METHOD
+
+    qualified = f"{parent_qname}.{name}" if parent_qname else name
+
+    start_byte = node.start_byte()
+    end_byte = node.end_byte()
+    code_bytes = parsed.source[start_byte:end_byte]
+    code = code_bytes.decode("utf-8", errors="replace")
+
+    start_point = node.start_position()
+    end_point = node.end_position()
+
+    return Chunk(
+        file_relpath=parsed.relpath,
+        language=parsed.language,
+        kind=kind,
+        name=name,
+        qualified_name=qualified,
+        start_line=start_point.row + 1,
+        end_line=end_point.row + 1,
+        start_byte=start_byte,
+        end_byte=end_byte,
+        code=code,
+        file_sha=parsed.sha,
+        parent=parent_qname,
+        signature=_extract_signature(code),
+    )
+
+
+def _extract_name(node, lang: LanguageSpec, source: bytes) -> str | None:
+    kind = node.kind()
+
+    # Python decorated_definition: name lives on the wrapped function/class.
+    if kind == "decorated_definition":
+        for child in _iter_children(node):
+            if child.kind() in {"function_definition", "class_definition"}:
+                return _extract_name(child, lang, source)
+        return None
+
+    # Rust impl_item: prefer the "type" being implemented (or the trait).
+    if kind == "impl_item":
+        for fname in ("type", "trait"):
+            named = node.child_by_field_name(fname)
+            if named is not None:
+                return _node_text(named, source)
+        return None
+
+    # Go type_declaration wraps one or more type_specs; take the first.
+    if kind == "type_declaration":
+        for child in _iter_children(node):
+            if child.kind() == "type_spec":
+                named = child.child_by_field_name("name")
+                if named is not None:
+                    return _node_text(named, source)
+        return None
+
+    named = node.child_by_field_name(lang.name_field)
+    if named is not None:
+        return _node_text(named, source)
+
+    for child in _iter_children(node):
+        if child.kind() in {"identifier", "property_identifier", "type_identifier"}:
+            return _node_text(child, source)
+    return None
+
+
+def _node_text(node, source: bytes) -> str:
+    return source[node.start_byte():node.end_byte()].decode("utf-8", errors="replace")
+
+
+def _extract_signature(code: str) -> str:
+    for line in code.splitlines():
+        stripped = line.strip()
+        if stripped:
+            if len(stripped) > _SIGNATURE_MAX_BYTES:
+                return stripped[:_SIGNATURE_MAX_BYTES] + "…"
+            return stripped
+    return ""
+
+
+def chunk_files(parsed_files: Iterator[ParsedFile]) -> Iterator[Chunk]:
+    for parsed in parsed_files:
+        yield from chunk_file(parsed)