nabla-cli 0.1.0__py3-none-any.whl

nabla/prompt_recon.py

@@ -0,0 +1,542 @@
+"""T2 — prompt template reconstruction. See cli-plan.md T2.
+
+Owner fills this module. Re-parses source from raw_site.file/span; does NOT
+import scanner internals.
+
+Approach: pure static backward dataflow over the AST. Starting from the
+`messages=` argument of the SDK call, each message's `content` expression is
+resolved to a template string by chasing Name references through a small
+stack of "frames" -- one per function body currently being inlined:
+
+  * module-level constants (simple `NAME = <expr>` assignments) are inlined
+    directly;
+  * local assignments (`x = <expr>`) are inlined, with `<expr>` evaluated in
+    the same frame;
+  * calls to other functions defined in the same module are inlined by
+    binding the callee's parameters to the caller's argument expressions
+    (evaluated in the *caller's* frame) and resolving the callee's `return`
+    expression;
+  * a function parameter that has no binding (i.e. we've walked all the way
+    back to the function that directly contains -- or, for a wrapper site,
+    is one hop away from -- the SDK call) becomes a named `{slot}`;
+  * anything else (attribute access on an unknown object, calls to
+    functions we can't locate, subscripts, comprehensions, ...) becomes an
+    opaque `{UNRESOLVED_n}` slot and sets `partial=True`. This function never
+    raises; degenerate/unfamiliar code always degrades to a partial result.
+
+For `wrapper_indirect` sites, `raw_site` describes the *caller* (the
+reported site) and `raw_site.wrapper` describes where the SDK call actually
+lives. We resolve the wrapper function's messages, but treat its own
+parameters as bound to the caller's argument expressions at the one call
+site (within `raw_site.symbol`) that invokes the wrapper -- i.e. exactly one
+cross-module hop, as required.
+"""
+
+from __future__ import annotations
+
+import ast
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from .contract import PromptTemplate, RawSite, Slot
+
+_FALLBACK_NAME = "UNRESOLVED_1"
+
+
+def reconstruct(raw_site: RawSite, repo_path: str) -> PromptTemplate:
+    """Recover the prompt template with named {slot} markers.
+
+    Must: resolve messages back through assignments, f-strings, .format(),
+    %, concatenation, and function parameters (module + one call-hop);
+    split system vs user messages; mark unresolvable fragments as
+    {UNRESOLVED_n} slots and set partial=True instead of failing.
+    """
+    try:
+        result = _reconstruct_inner(raw_site, repo_path)
+        if result is not None:
+            return result
+    except Exception:
+        pass
+    return _fallback_template()
+
+
+def _fallback_template() -> PromptTemplate:
+    return PromptTemplate(
+        template="{" + _FALLBACK_NAME + "}",
+        slots=[Slot(name=_FALLBACK_NAME, inferred_type="unresolved")],
+        system_prompt=None,
+        embedded_few_shot=[],
+        partial=True,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Analysis state
+# ---------------------------------------------------------------------------
+
+
+class _Ctx:
+    def __init__(self) -> None:
+        self.partial = False
+        self.unresolved_count = 0
+
+
+@dataclass
+class _ModuleInfo:
+    path: Path
+    tree: ast.Module
+    module_consts: dict[str, ast.expr]
+    functions: dict[str, ast.FunctionDef | ast.AsyncFunctionDef]
+
+
+@dataclass
+class _Frame:
+    module: _ModuleInfo
+    func: ast.FunctionDef | ast.AsyncFunctionDef | None
+    env: dict[str, tuple[ast.expr, "_Frame"]] = field(default_factory=dict)
+    # is_root frames are those whose own parameters have no known caller
+    # binding -- their unbound parameters become the actual named slots.
+    is_root: bool = False
+
+
+# ---------------------------------------------------------------------------
+# Top-level driver
+# ---------------------------------------------------------------------------
+
+
+def _reconstruct_inner(raw_site: RawSite, repo_path: str) -> PromptTemplate | None:
+    cache: dict[str, _ModuleInfo] = {}
+    ctx = _Ctx()
+    slots: dict[str, Slot] = {}
+
+    if raw_site.wrapper is not None:
+        wrapper_module = _load_module(repo_path, raw_site.wrapper.file, cache)
+        wrapper_func = _find_function(wrapper_module.tree, raw_site.wrapper.symbol)
+        if wrapper_func is None:
+            return None
+        target_call = _find_call(wrapper_func, raw_site.wrapper.span, raw_site.kind)
+        if target_call is None:
+            return None
+
+        caller_module = _load_module(repo_path, raw_site.file, cache)
+        caller_func = _find_function(caller_module.tree, raw_site.symbol)
+        if caller_func is None:
+            return None
+        caller_frame = _make_root_frame(caller_module, caller_func)
+
+        wrapper_name = _bare_name(raw_site.wrapper.symbol)
+        inner_call = _find_named_call(caller_func, wrapper_name)
+        if inner_call is None:
+            return None
+        message_frame = _make_callee_frame(wrapper_module, wrapper_func, inner_call, caller_frame)
+    else:
+        module = _load_module(repo_path, raw_site.file, cache)
+        func = _find_function(module.tree, raw_site.symbol)
+        if func is None:
+            return None
+        target_call = _find_call(func, raw_site.span, raw_site.kind)
+        if target_call is None:
+            return None
+        message_frame = _make_root_frame(module, func)
+
+    messages_node = _get_kw(target_call, "messages")
+    if messages_node is None or not isinstance(messages_node, ast.List):
+        return None
+
+    parsed: list[tuple[str, str]] = [
+        _resolve_message(elt, message_frame, slots, ctx) for elt in messages_node.elts
+    ]
+
+    system_prompt: str | None = None
+    if parsed and parsed[0][0] == "system":
+        system_prompt = parsed[0][1]
+        parsed = parsed[1:]
+
+    embedded_few_shot: list[dict[str, str]] = []
+    if parsed:
+        *lead, last_msg = parsed
+        template = last_msg[1]
+        embedded_few_shot = [{"role": role, "content": content} for role, content in lead]
+    else:
+        template = _new_unresolved(slots, ctx)
+
+    return PromptTemplate(
+        template=template,
+        slots=list(slots.values()),
+        system_prompt=system_prompt,
+        embedded_few_shot=embedded_few_shot,
+        partial=ctx.partial,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Module loading / AST lookups
+# ---------------------------------------------------------------------------
+
+
+def _load_module(repo_path: str, rel_file: str, cache: dict[str, _ModuleInfo]) -> _ModuleInfo:
+    if rel_file in cache:
+        return cache[rel_file]
+    full_path = Path(repo_path) / rel_file
+    source = full_path.read_text(encoding="utf-8")
+    tree = ast.parse(source, filename=str(full_path))
+    module_consts: dict[str, ast.expr] = {}
+    functions: dict[str, ast.FunctionDef | ast.AsyncFunctionDef] = {}
+    for stmt in tree.body:
+        if isinstance(stmt, ast.Assign) and len(stmt.targets) == 1 and isinstance(stmt.targets[0], ast.Name):
+            module_consts[stmt.targets[0].id] = stmt.value
+        elif isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            functions[stmt.name] = stmt
+    info = _ModuleInfo(path=full_path, tree=tree, module_consts=module_consts, functions=functions)
+    cache[rel_file] = info
+    return info
+
+
+def _bare_name(qualname: str) -> str:
+    name = qualname.strip()
+    if name.endswith("()"):
+        name = name[:-2]
+    return name.split(".")[-1]
+
+
+def _find_function(tree: ast.AST, qualname: str) -> ast.FunctionDef | ast.AsyncFunctionDef | None:
+    target = _bare_name(qualname)
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == target:
+            return node
+    return None
+
+
+def _find_call(func: ast.AST, span: tuple[int, int], kind: str) -> ast.Call | None:
+    candidates = [
+        node
+        for node in ast.walk(func)
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == kind
+    ]
+    if not candidates:
+        return None
+    lo, hi = span
+    for node in candidates:
+        node_lo = node.lineno
+        node_hi = getattr(node, "end_lineno", node.lineno)
+        if node_lo <= hi and node_hi >= lo:
+            return node
+    return candidates[0]
+
+
+def _find_named_call(func: ast.AST, name: str) -> ast.Call | None:
+    for node in ast.walk(func):
+        if isinstance(node, ast.Call):
+            f = node.func
+            if isinstance(f, ast.Name) and f.id == name:
+                return node
+            if isinstance(f, ast.Attribute) and f.attr == name:
+                return node
+    return None
+
+
+def _get_kw(call: ast.Call, name: str) -> ast.expr | None:
+    for kw in call.keywords:
+        if kw.arg == name:
+            return kw.value
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Scope walking (stays within one function; does not descend into nested
+# defs/classes/lambdas, so their locals don't leak into the outer scope)
+# ---------------------------------------------------------------------------
+
+
+def _scope_statements(stmts: list[ast.stmt]):
+    for stmt in stmts:
+        yield stmt
+        for _, field_val in ast.iter_fields(stmt):
+            if isinstance(field_val, list):
+                for item in field_val:
+                    if isinstance(item, ast.stmt) and not isinstance(
+                        item, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)
+                    ):
+                        yield from _scope_statements([item])
+            elif isinstance(field_val, ast.stmt) and not isinstance(
+                field_val, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)
+            ):
+                yield from _scope_statements([field_val])
+
+
+def _collect_local_assigns(func: ast.FunctionDef | ast.AsyncFunctionDef) -> dict[str, ast.expr]:
+    out: dict[str, ast.expr] = {}
+    for stmt in _scope_statements(func.body):
+        if isinstance(stmt, ast.Assign) and len(stmt.targets) == 1 and isinstance(stmt.targets[0], ast.Name):
+            out[stmt.targets[0].id] = stmt.value
+        elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name) and stmt.value is not None:
+            out[stmt.target.id] = stmt.value
+    return out
+
+
+def _find_return_expr(func: ast.FunctionDef | ast.AsyncFunctionDef) -> ast.expr | None:
+    result: ast.expr | None = None
+    for stmt in _scope_statements(func.body):
+        if isinstance(stmt, ast.Return) and stmt.value is not None:
+            result = stmt.value
+    return result
+
+
+def _param_names(func: ast.FunctionDef | ast.AsyncFunctionDef) -> list[str]:
+    a = func.args
+    return [p.arg for p in a.posonlyargs] + [p.arg for p in a.args] + [p.arg for p in a.kwonlyargs]
+
+
+# ---------------------------------------------------------------------------
+# Frame construction
+# ---------------------------------------------------------------------------
+
+
+def _make_root_frame(module: _ModuleInfo, func: ast.FunctionDef | ast.AsyncFunctionDef) -> _Frame:
+    frame = _Frame(module=module, func=func, env={}, is_root=True)
+    for name, val in _collect_local_assigns(func).items():
+        frame.env[name] = (val, frame)
+    return frame
+
+
+def _make_callee_frame(
+    module: _ModuleInfo,
+    func: ast.FunctionDef | ast.AsyncFunctionDef,
+    call_node: ast.Call,
+    calling_frame: _Frame,
+) -> _Frame:
+    frame = _Frame(module=module, func=func, env={}, is_root=False)
+    a = func.args
+    positional_params = [p.arg for p in a.posonlyargs] + [p.arg for p in a.args]
+    kwonly_params = [p.arg for p in a.kwonlyargs]
+
+    for i, arg_node in enumerate(call_node.args):
+        if i < len(positional_params):
+            frame.env[positional_params[i]] = (arg_node, calling_frame)
+    for kw in call_node.keywords:
+        if kw.arg and (kw.arg in positional_params or kw.arg in kwonly_params):
+            frame.env[kw.arg] = (kw.value, calling_frame)
+
+    if a.defaults:
+        default_params = positional_params[-len(a.defaults):]
+        for p, d in zip(default_params, a.defaults):
+            frame.env.setdefault(p, (d, frame))
+    for p, d in zip(kwonly_params, a.kw_defaults):
+        if d is not None:
+            frame.env.setdefault(p, (d, frame))
+
+    for name, val in _collect_local_assigns(func).items():
+        frame.env.setdefault(name, (val, frame))
+    return frame
+
+
+# ---------------------------------------------------------------------------
+# Expression resolution
+# ---------------------------------------------------------------------------
+
+
+def _escape(text: str) -> str:
+    return text.replace("{", "{{").replace("}", "}}")
+
+
+def _new_slot(name: str, slots: dict[str, Slot]) -> str:
+    if name not in slots:
+        slots[name] = Slot(name=name)
+    return "{" + name + "}"
+
+
+def _new_unresolved(slots: dict[str, Slot], ctx: _Ctx) -> str:
+    ctx.partial = True
+    ctx.unresolved_count += 1
+    name = f"UNRESOLVED_{ctx.unresolved_count}"
+    slots[name] = Slot(name=name, inferred_type="unresolved")
+    return "{" + name + "}"
+
+
+def _resolve_message(elt: ast.expr, frame: _Frame, slots: dict[str, Slot], ctx: _Ctx) -> tuple[str, str]:
+    if not isinstance(elt, ast.Dict):
+        return ("unknown", _new_unresolved(slots, ctx))
+    role = "unknown"
+    content_node: ast.expr | None = None
+    for k, v in zip(elt.keys, elt.values):
+        if isinstance(k, ast.Constant) and k.value == "role":
+            role = v.value if isinstance(v, ast.Constant) and isinstance(v.value, str) else "unknown"
+        elif isinstance(k, ast.Constant) and k.value == "content":
+            content_node = v
+    if content_node is None:
+        return (role, _new_unresolved(slots, ctx))
+    return (role, _resolve_expr(content_node, frame, slots, ctx))
+
+
+def _resolve_expr(node: ast.expr, frame: _Frame, slots: dict[str, Slot], ctx: _Ctx, depth: int = 0) -> str:
+    if depth > 60:
+        return _new_unresolved(slots, ctx)
+
+    if isinstance(node, ast.Constant):
+        return _escape(node.value) if isinstance(node.value, str) else _escape(str(node.value))
+
+    if isinstance(node, ast.JoinedStr):
+        parts = []
+        for v in node.values:
+            if isinstance(v, ast.Constant):
+                parts.append(_escape(v.value) if isinstance(v.value, str) else _escape(str(v.value)))
+            elif isinstance(v, ast.FormattedValue):
+                parts.append(_resolve_expr(v.value, frame, slots, ctx, depth + 1))
+            else:
+                parts.append(_new_unresolved(slots, ctx))
+        return "".join(parts)
+
+    if isinstance(node, ast.Name):
+        return _resolve_name(node.id, frame, slots, ctx, depth)
+
+    if isinstance(node, ast.BinOp):
+        if isinstance(node.op, ast.Add):
+            return _resolve_expr(node.left, frame, slots, ctx, depth + 1) + _resolve_expr(
+                node.right, frame, slots, ctx, depth + 1
+            )
+        if isinstance(node.op, ast.Mod):
+            raw = _raw_literal(node.left, frame)
+            if raw is not None:
+                return _apply_percent(raw, node.right, frame, slots, ctx)
+        return _new_unresolved(slots, ctx)
+
+    if isinstance(node, ast.Call):
+        return _resolve_call(node, frame, slots, ctx, depth)
+
+    return _new_unresolved(slots, ctx)
+
+
+def _resolve_name(name: str, frame: _Frame, slots: dict[str, Slot], ctx: _Ctx, depth: int) -> str:
+    if name in frame.env:
+        val_node, val_frame = frame.env[name]
+        return _resolve_expr(val_node, val_frame, slots, ctx, depth + 1)
+    if name in frame.module.module_consts:
+        val_node = frame.module.module_consts[name]
+        mod_frame = _Frame(module=frame.module, func=None, env={}, is_root=False)
+        return _resolve_expr(val_node, mod_frame, slots, ctx, depth + 1)
+    if frame.func is not None and frame.is_root and name in _param_names(frame.func):
+        return _new_slot(name, slots)
+    return _new_unresolved(slots, ctx)
+
+
+def _resolve_call(node: ast.Call, frame: _Frame, slots: dict[str, Slot], ctx: _Ctx, depth: int) -> str:
+    func = node.func
+
+    if isinstance(func, ast.Attribute):
+        if func.attr in ("strip", "lstrip", "rstrip") and not node.args and not node.keywords:
+            inner = _resolve_expr(func.value, frame, slots, ctx, depth + 1)
+            if func.attr == "strip":
+                return inner.strip()
+            if func.attr == "lstrip":
+                return inner.lstrip()
+            return inner.rstrip()
+        if func.attr == "format":
+            raw = _raw_literal(func.value, frame)
+            if raw is not None:
+                return _apply_format(raw, node.args, node.keywords, frame, slots, ctx)
+            return _new_unresolved(slots, ctx)
+
+    if isinstance(func, ast.Name) and func.id in frame.module.functions:
+        callee = frame.module.functions[func.id]
+        callee_frame = _make_callee_frame(frame.module, callee, node, frame)
+        ret = _find_return_expr(callee)
+        if ret is None:
+            return _new_unresolved(slots, ctx)
+        return _resolve_expr(ret, callee_frame, slots, ctx, depth + 1)
+
+    return _new_unresolved(slots, ctx)
+
+
+def _raw_literal(node: ast.expr, frame: _Frame, depth: int = 0) -> str | None:
+    """Best-effort *unescaped* literal string, for use as a .format()/%
+    base template (whose braces/`%s` are meaningful, not literal text).
+    Returns None if it doesn't bottom out at a plain string constant."""
+    if depth > 20:
+        return None
+    if isinstance(node, ast.Constant) and isinstance(node.value, str):
+        return node.value
+    if isinstance(node, ast.Name):
+        if node.id in frame.env:
+            val_node, val_frame = frame.env[node.id]
+            return _raw_literal(val_node, val_frame, depth + 1)
+        if node.id in frame.module.module_consts:
+            val_node = frame.module.module_consts[node.id]
+            mod_frame = _Frame(module=frame.module, func=None, env={}, is_root=False)
+            return _raw_literal(val_node, mod_frame, depth + 1)
+    return None
+
+
+_FORMAT_FIELD_RE = re.compile(r"\{([^{}]*)\}")
+
+
+def _apply_format(
+    raw: str,
+    call_args: list[ast.expr],
+    call_keywords: list[ast.keyword],
+    frame: _Frame,
+    slots: dict[str, Slot],
+    ctx: _Ctx,
+) -> str:
+    kwmap = {kw.arg: kw.value for kw in call_keywords if kw.arg}
+    pos_args = list(call_args)
+    auto_idx = 0
+    out: list[str] = []
+    last = 0
+    for m in _FORMAT_FIELD_RE.finditer(raw):
+        out.append(_escape(raw[last:m.start()]))
+        last = m.end()
+        field = m.group(1).split("!", 1)[0].split(":", 1)[0]
+        arg_node: ast.expr | None
+        if field == "":
+            arg_node = pos_args[auto_idx] if auto_idx < len(pos_args) else None
+            auto_idx += 1
+        elif field.isdigit():
+            idx = int(field)
+            arg_node = pos_args[idx] if idx < len(pos_args) else None
+        else:
+            arg_node = kwmap.get(field)
+        if arg_node is None:
+            out.append(_new_unresolved(slots, ctx))
+        else:
+            out.append(_resolve_expr(arg_node, frame, slots, ctx))
+    out.append(_escape(raw[last:]))
+    return "".join(out)
+
+
+_PERCENT_RE = re.compile(r"%(\(([^)]+)\))?[-+#0 ]*\d*(\.\d+)?[a-zA-Z%]")
+
+
+def _apply_percent(raw: str, right: ast.expr, frame: _Frame, slots: dict[str, Slot], ctx: _Ctx) -> str:
+    positional: list[ast.expr] | None = None
+    dict_map: dict[str, ast.expr] = {}
+    if isinstance(right, ast.Tuple):
+        positional = list(right.elts)
+    elif isinstance(right, ast.Dict):
+        for k, v in zip(right.keys, right.values):
+            if isinstance(k, ast.Constant):
+                dict_map[k.value] = v
+    else:
+        positional = [right]
+
+    out: list[str] = []
+    last = 0
+    idx = 0
+    for m in _PERCENT_RE.finditer(raw):
+        out.append(_escape(raw[last:m.start()]))
+        last = m.end()
+        if m.group(0) == "%%":
+            out.append("%")
+            continue
+        name = m.group(2)
+        if name is not None:
+            arg_node = dict_map.get(name)
+        else:
+            arg_node = positional[idx] if positional is not None and idx < len(positional) else None
+            idx += 1
+        if arg_node is None:
+            out.append(_new_unresolved(slots, ctx))
+        else:
+            out.append(_resolve_expr(arg_node, frame, slots, ctx))
+    out.append(_escape(raw[last:]))
+    return "".join(out)