| """VULPO-style lightweight repository context extraction (spec §4.5). |
| |
| For each target function, extracts (a) callee function bodies (1 hop), |
| (b) callee signatures only (2 hops), (c) relevant global/type declarations — |
| budgeted to <=4096 tokens, truncating callee bodies first, then dropping |
| 2-hop signatures. Where repo access is unavailable, context is empty and |
| `has_context=False` (spec: "never fabricate context"). |
| |
| Token counting here is a simple whitespace-split approximation, not a real |
| tokenizer — good enough for budgeting (consistently conservative: real BPE |
| tokenizers produce >= this many tokens for code, so this under-counts and |
| stays safely under budget, never over). |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
|
|
| DEFAULT_TOKEN_BUDGET = 4096 |
|
|
|
|
| def approx_token_count(text: str) -> int: |
| return len(text.split()) |
|
|
|
|
| @dataclass |
| class RepoContext: |
| callee_bodies: dict[str, str] = field(default_factory=dict) |
| callee_signatures: dict[str, str] = field(default_factory=dict) |
| global_decls: list[str] = field(default_factory=list) |
| has_context: bool = False |
| truncated: bool = False |
|
|
|
|
| def build_context( |
| callee_bodies: dict[str, str], |
| callee_signatures: dict[str, str], |
| global_decls: list[str], |
| token_budget: int = DEFAULT_TOKEN_BUDGET, |
| ) -> RepoContext: |
| """Assembles context under `token_budget`, truncating callee bodies first |
| (dropping lowest-priority ones entirely once budget is exhausted — not |
| mid-body, since a half-truncated function body is worse than none), then |
| dropping 2-hop signatures (spec §4.5's stated truncation order).""" |
| if not callee_bodies and not callee_signatures and not global_decls: |
| return RepoContext(has_context=False) |
|
|
| budget = token_budget |
| truncated = False |
|
|
| kept_decls = [] |
| for decl in global_decls: |
| cost = approx_token_count(decl) |
| if cost <= budget: |
| kept_decls.append(decl) |
| budget -= cost |
| else: |
| truncated = True |
|
|
| kept_bodies: dict[str, str] = {} |
| for name, body in callee_bodies.items(): |
| cost = approx_token_count(body) |
| if cost <= budget: |
| kept_bodies[name] = body |
| budget -= cost |
| else: |
| truncated = True |
|
|
| kept_sigs: dict[str, str] = {} |
| for name, sig in callee_signatures.items(): |
| cost = approx_token_count(sig) |
| if cost <= budget: |
| kept_sigs[name] = sig |
| budget -= cost |
| else: |
| truncated = True |
|
|
| return RepoContext( |
| callee_bodies=kept_bodies, |
| callee_signatures=kept_sigs, |
| global_decls=kept_decls, |
| has_context=True, |
| truncated=truncated, |
| ) |
|
|
|
|
| def render_context(ctx: RepoContext) -> str: |
| """Flattens a RepoContext into the text block the prompt template embeds.""" |
| if not ctx.has_context: |
| return "" |
| parts: list[str] = [] |
| if ctx.global_decls: |
| parts.append("\n".join(ctx.global_decls)) |
| for name, body in ctx.callee_bodies.items(): |
| parts.append(f"// callee: {name}\n{body}") |
| for name, sig in ctx.callee_signatures.items(): |
| parts.append(f"// callee signature: {name}\n{sig}") |
| return "\n\n".join(parts) |
|
|
