Bài 3 — Đọc code thông minh: AST, chữ ký hàm và ngữ cảnh vừa đủ

Ở Bài 2 ta ném cả file vào prompt. Cách đó tốn token và dễ khiến model bịa. Bài này ta dạy agent đọc có cấu trúc bằng AST (Abstract Syntax Tree), để lấy đúng chữ ký hàm, tham số, kiểu trả về — nguồn sự thật chống hallucination tốt nhất.

Vì sao AST tốt hơn raw text

Raw text: model phải "đoán" đâu là tham số. AST: bạn trích chính xác danh sách tham số, default value, decorator, type hint — dữ liệu có thật, không đoán. Bạn đưa dữ liệu này vào prompt như "sự thật đã xác minh", model chỉ việc diễn giải.

Trích chữ ký bằng module ast của Python

import ast

def extract_functions(path):
    tree = ast.parse(open(path).read())
    out = []
    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef):
            args = [a.arg for a in node.args.args]
            has_doc = ast.get_docstring(node) is not None
            out.append({
                "name": node.name,
                "args": args,
                "lineno": node.lineno,
                "has_doc": has_doc,
                "source": ast.get_source_segment(open(path).read(), node),
            })
    return out

Giờ bạn biết chắc hàm fetch_user có [user_id, retries, timeout] — không cần model đoán.

Các bước xây Source Reader theo AST

• Parse file thành AST bằng ast.parse.
• Duyệt cây, lọc FunctionDef/AsyncFunctionDef/ClassDef.
• Với mỗi node, trích: tên, danh sách tham số, default, type hint, decorator, docstring hiện có.
• Đánh dấu node thiếu docstring hoặc docstring không khớp tham số (candidate để cập nhật).
• Chỉ gửi những node cần xử lý + ngữ cảnh liên quan cho LLM.
• Đưa danh sách tham số như một "contract" trong prompt để model không được thêm bớt.

Prompt kết hợp AST

def build_prompt(fn):
    return (
        f"Ham: {fn['name']}\n"
        f"Tham so CHINH XAC (khong them bot): {fn['args']}\n\n"
        f"Ma nguon:\n{fn['source']}\n\n"
        "Viet docstring kieu Google. Moi tham so trong danh sach "
        "tren PHAI xuat hien dung 1 lan, khong duoc them tham so la."
    )

Khung tư duy "ngữ cảnh vừa đủ" (right-sized context)

• Quá ít (chỉ tên hàm) → model bịa.
• Quá nhiều (cả repo) → tốn tiền, model lạc, vượt context window.
• Vừa đủ = hàm đích + chữ ký + các symbol nó gọi trực tiếp + type liên quan.

Quy tắc: đưa những gì cần để mô tả đúng, không hơn. Với hàm gọi hàm khác, kèm chữ ký của hàm được gọi (không cần thân).

Với ngôn ngữ khác Python

Dùng tree-sitter (đa ngôn ngữ) hoặc TypeScript Compiler API cho TS/JS. Nguyên tắc giống nhau: trích cấu trúc thật rồi mới đưa vào prompt. tree-sitter đặc biệt hữu ích khi agent cần chạy trên codebase đa ngôn ngữ (backend Python + frontend TS): bạn dùng chung một lớp reader với nhiều grammar, trả về cùng một cấu trúc {name, args, source, has_doc} để phần generator/verifier phía sau không phải quan tâm ngôn ngữ.

Trích cả type hint và giá trị mặc định

Chữ ký hàm không chỉ có tên tham số. node.args.defaults cho bạn giá trị mặc định, node.returns cho kiểu trả về khai báo, và các arg.annotation cho type hint từng tham số. Càng nhiều dữ liệu thật đưa vào contract, model càng ít phải đoán:

def arg_details(node):
    args = node.args.args
    defaults = [None]*(len(args)-len(node.args.defaults)) + node.args.defaults
    out = []
    for a, d in zip(args, defaults):
        ann = ast.unparse(a.annotation) if a.annotation else None
        dv = ast.unparse(d) if d else None
        out.append({"name": a.arg, "type": ann, "default": dv})
    return out

Giờ prompt có thể nói: "tham số timeout kiểu int, mặc định 30" — không còn chỗ cho model tưởng tượng.

Bảng quyết định: node nào cần xử lý

Bảng này chính là bộ lọc giúp agent chỉ gọi LLM cho những gì thực sự cần, tiết kiệm chi phí đáng kể trên repo lớn.

Sai lầm thường gặp

• Bỏ qua AST vì "ném cả file cho nhanh". Nhanh lúc đầu, nhưng bịa nhiều và tốn token về lâu dài.
• Quên xử lý async/nested/decorator. ast.walk bắt được nhưng cần phân biệt để mô tả đúng (ví dụ decorator @property).
• Không đối chiếu tham số docstring cũ vs chữ ký mới. Đây chính là nơi phát hiện drift — đừng bỏ.
• Đọc file mỗi node một lần như ví dụ trên (gọi open() lặp). Trong thực tế hãy đọc source một lần, truyền vào — tránh I/O thừa.
• Bỏ sót args/kwargs/tham số keyword-only.* node.args có vararg, kwarg, kwonlyargs riêng — quên chúng khiến contract thiếu và model mô tả sai.

Kết bài: agent giờ đọc code có cấu trúc và biết node nào cần docs. Bài 4 ta dùng khả năng này để sinh docstring chất lượng cao hàng loạt.