RAG (Retrieval-Augmented Generation — sinh nội dung có tăng cường truy hồi) ra đời để giải quyết bài toán hỏi-đáp trên tài liệu văn bản, nhưng khi áp dụng vào codebase, kỹ thuật này còn phát huy sức mạnh lớn hơn nhiều. Một hệ thống RAG cho code có thể truy hồi động những đoạn code liên quan nhất cho từng câu hỏi cụ thể, đưa đúng phần context cần thiết vào context window (cửa sổ ngữ cảnh) của AI agent mà không phải nạp toàn bộ repo. Nếu các file context tĩnh như CLAUDE.md hay repo map cung cấp "kiến thức nền" cố định, thì RAG cung cấp "kiến thức tiền tuyến" — thay đổi theo từng truy vấn, chỉ lấy đúng cái cần cho câu hỏi đang được xử lý.

Nghe thì giống RAG cho tài liệu, nhưng RAG cho code là một bài toán khác hẳn về bản chất. Bài viết này đi sâu vào lý do vì sao code cần một kiến trúc retrieval (truy hồi) riêng, các chiến lược chunking (chia nhỏ) tôn trọng cấu trúc code, cách chọn embedding model phù hợp, kỹ thuật hybrid retrieval kết hợp semantic và keyword search, retrieval được tăng cường bằng dependency graph (đồ thị phụ thuộc), và cuối cùng là một pipeline RAG hoàn chỉnh dựng bằng LangChain mà kỹ sư, QA và PM đều có thể áp dụng vào công việc hàng ngày.

Vì Sao RAG cho Code Khác với RAG cho Tài liệu

Với tài liệu văn bản thông thường (hợp đồng, bài báo, tài liệu hướng dẫn), RAG hoạt động theo một công thức khá đơn giản: cắt văn bản thành các đoạn có độ dài cố định (fixed-size chunk), embed từng đoạn thành vector, rồi khi có câu hỏi thì tìm những đoạn có vector gần nhất bằng cosine similarity. Công thức này hoạt động tốt vì văn bản tự nhiên có tính "liên tục" — cắt một đoạn văn ở giữa câu vẫn còn giữ được phần lớn ý nghĩa nhờ ngữ cảnh xung quanh.

Code thì không như vậy. Code có ba đặc điểm khiến việc áp dụng nguyên công thức RAG-cho-văn-bản trở thành một sai lầm tốn kém:

1. Code có cấu trúc cú pháp (syntax) chặt, không phải văn xuôi. Một function, một class, một block if/else là một đơn vị hoàn chỉnh về mặt logic — cắt đôi nó ra (ví dụ dừng chunk ngay giữa thân function) sẽ tạo ra một đoạn code không compile được, không chạy được, và với LLM thì cũng khó hiểu được trọn vẹn ý nghĩa. Trong văn bản, cắt giữa câu vẫn còn ngữ nghĩa; trong code, cắt giữa function là phá vỡ hoàn toàn đơn vị ngữ nghĩa.

2. Code có tham chiếu chéo file (cross-file reference) dày đặc. Một function gọi một function khác được định nghĩa ở file khác; một class kế thừa từ class ở module khác; một biến được import từ một package riêng. Định nghĩa của calculateDiscount() là vô nghĩa nếu tách rời khỏi nơi nó bị gọi, và ngược lại. Semantic search thuần túy (dựa trên độ tương đồng nội dung văn bản) hoàn toàn không "biết" về các cạnh (edge) phụ thuộc này — nó chỉ so khớp nội dung có "giống nhau về ngữ nghĩa" hay không, trong khi function definition và call site thường có nội dung văn bản rất khác nhau.

3. Function/class là đơn vị ngữ nghĩa tự nhiên, không phải một đoạn text tùy ý. Đây là điểm khác biệt quan trọng nhất: với code, ranh giới ngữ nghĩa tự nhiên đã được compiler/parser định nghĩa sẵn — đó là function, method, class, hoặc module. Một hệ RAG tốt nên tôn trọng ranh giới này, thay vì áp một cửa sổ trượt (sliding window) cố định 512 token như cách người ta vẫn làm với tài liệu.

Hệ quả thực tế: nếu bạn dùng một pipeline RAG "mặc định" (built cho tài liệu) để index một codebase, bạn sẽ gặp các lỗi kinh điển — retrieve về nửa function không có phần khai báo, retrieve về đoạn code gọi hàm nhưng thiếu định nghĩa hàm, hoặc bỏ sót hoàn toàn những file liên quan chỉ vì chúng dùng identifier khác tên biến (dù về mặt logic vẫn liên quan chặt). Với một AI agent coding, hậu quả không chỉ là "câu trả lời chưa hay" — mà có thể là code được sinh ra sai hoàn toàn vì thiếu context, hoặc token bị lãng phí vì agent phải "dò" thêm nhiều vòng để tự tìm phần còn thiếu.

Mẹo

• Trước khi build RAG cho một codebase, hãy tự hỏi: "Nếu tôi cắt ngẫu nhiên ở đây, đoạn code còn lại có đứng được một mình không?" Nếu câu trả lời là không, chunking của bạn cần điều chỉnh.
• Đừng copy nguyên pipeline RAG cho tài liệu (ví dụ từ một tutorial PDF Q&A) rồi áp thẳng vào code — tỷ lệ retrieve sai sẽ rất cao và rất khó debug vì lỗi "ẩn" (silent failure), agent vẫn trả lời nhưng dựa trên context thiếu.
• Với team đa persona (dev, QA, PM), hãy nhớ rằng PM và QA thường hỏi ở mức "tính năng" hoặc "hành vi", không phải mức "dòng code" — RAG cho code cần retrieve đủ rộng (cả comment, docstring, test case) để phục vụ những câu hỏi này, không chỉ tối ưu cho câu hỏi kiểu "function X nằm ở đâu".

Các Chiến lược Chunking Giữ Nguyên Ý nghĩa Code

Chunking (chia nhỏ) là bước quyết định 80% chất lượng của một hệ RAG cho code. Chunk sai, mọi bước sau (embedding, retrieval, ranking) đều xây trên nền cát.

AST-based chunking — chuẩn vàng cho code

Cách tốt nhất để chunk code là dựa trên AST (Abstract Syntax Tree — cây cú pháp trừu tượng) thay vì đếm ký tự. Thư viện phổ biến nhất cho việc này là tree-sitter — một parser đa ngôn ngữ, nhanh, và có thể chạy incremental (chỉ parse lại phần thay đổi). Ý tưởng cốt lõi: parse file thành AST, sau đó chunk theo ranh giới node — mỗi function, mỗi method, mỗi class là một chunk (hoặc nếu class quá lớn, mỗi method trong class là một chunk, kèm theo signature của class cha để giữ ngữ cảnh).

from tree_sitter import Language, Parser
import tree_sitter_python as tspython

PY_LANGUAGE = Language(tspython.language())
parser = Parser(PY_LANGUAGE)

CHUNKABLE_NODE_TYPES = {"function_definition", "class_definition"}

def extract_chunks(source_code: bytes, file_path: str):
    tree = parser.parse(source_code)
    chunks = []

    def walk(node, class_context=None):
        if node.type in CHUNKABLE_NODE_TYPES:
            chunk_text = source_code[node.start_byte:node.end_byte].decode("utf-8")

            # Gắn docstring / comment ngay phía trên node vào cùng chunk
            leading_comment = _get_leading_comment(node, source_code)

            # Nếu là method trong class, giữ lại class signature để không mất ngữ cảnh
            prefix = f"# class {class_context}:\n" if class_context and node.type == "function_definition" else ""

            chunks.append({
                "file_path": file_path,
                "node_type": node.type,
                "name": _get_node_name(node, source_code),
                "start_line": node.start_point[0] + 1,
                "end_line": node.end_point[0] + 1,
                "content": prefix + leading_comment + chunk_text,
            })

            # Với class, tiếp tục đi sâu để chunk từng method riêng
            if node.type == "class_definition":
                class_name = _get_node_name(node, source_code)
                for child in node.children:
                    walk(child, class_context=class_name)
            return

        for child in node.children:
            walk(child, class_context)

    walk(tree.root_node)
    return chunks

Với logic này, mỗi chunk luôn là một function hoặc method hoàn chỉnh, có kèm docstring/comment liền trước nó (rất quan trọng — docstring thường chứa mô tả ý định bằng ngôn ngữ tự nhiên, giúp embedding "hiểu" chunk tốt hơn nhiều so với chỉ có code thuần), và nếu là method thì có thêm một dòng ghi chú class cha để LLM không bị lạc ngữ cảnh khi chunk được retrieve đứng một mình.

Import và top-level statement — đừng bỏ quên

Một sai lầm phổ biến: chỉ chunk function/class mà bỏ qua phần import ở đầu file. Hệ quả là chunk retrieve về có gọi pd.DataFrame nhưng không biết pd là alias của pandas. Cách xử lý phổ biến: tạo một "file header chunk" riêng chứa toàn bộ import + các biến/constant ở top-level, rồi tùy chọn nối (prepend) header này vào mỗi chunk function khi build prompt — hoặc lưu header như metadata để agent tự quyết định có cần lấy thêm không.

Sliding window với overlap — phương án dự phòng

AST-based chunking cần một parser hỗ trợ đúng ngôn ngữ. Với file config (YAML, JSON), file markdown, hoặc ngôn ngữ chưa có parser tree-sitter ổn định, bạn cần phương án dự phòng: sliding window với overlap. Nguyên tắc: chunk theo số dòng cố định (ví dụ 100 dòng) nhưng cho hai chunk liên tiếp overlap khoảng 15-20 dòng, để nếu một khối logic nằm vắt qua ranh giới chunk, nó vẫn xuất hiện trọn vẹn ở ít nhất một trong hai chunk.

def sliding_window_chunk(lines: list[str], window_size=100, overlap=20):
    chunks = []
    step = window_size - overlap
    for start in range(0, len(lines), step):
        end = min(start + window_size, len(lines))
        chunks.append({
            "start_line": start + 1,
            "end_line": end,
            "content": "\n".join(lines[start:end]),
        })
        if end == len(lines):
            break
    return chunks

Đánh đổi kích thước chunk (chunk size)

• Chunk nhỏ (một function ngắn, 10-30 dòng): embedding chính xác hơn (ít "nhiễu" ngữ nghĩa), nhưng số lượng chunk tăng, chi phí embedding và lưu trữ vector store tăng, và có nguy cơ mất ngữ cảnh liên quan (ví dụ helper function nhỏ tách rời khỏi nơi dùng nó).
• Chunk lớn (cả file, hoặc class lớn): giữ được nhiều ngữ cảnh hơn trong một chunk, nhưng embedding bị "pha loãng" — vector đại diện cho một đoạn quá dài thường không phân biệt tốt các truy vấn cụ thể, và khi retrieve về, bạn tốn nhiều token hơn cho phần không liên quan.
• Quy tắc thực dụng: với hầu hết codebase enterprise, chunk ở cấp function/method (kèm docstring + class context) là điểm cân bằng tốt nhất. Với file cấu hình hoặc file rất nhỏ (dưới ~50 dòng), chunk cả file làm một đơn vị.

Mẹo

• Luôn gắn metadata đầy đủ vào mỗi chunk: file_path, tên function/class, start_line/end_line, ngôn ngữ. Đây là dữ liệu bắt buộc cho bước graph-enhanced retrieval và hybrid search ở phần sau.
• Test chunking bằng cách tự đọc lại một vài chunk ngẫu nhiên: nếu bạn (con người) đọc một chunk mà không hiểu nó đang làm gì vì thiếu ngữ cảnh, LLM cũng sẽ gặp vấn đề tương tự.
• Với monorepo nhiều ngôn ngữ, đừng dùng một chiến lược chunking cho tất cả — mỗi ngôn ngữ cần grammar tree-sitter riêng (Python, TypeScript, Go, Java... đều có package tree-sitter riêng).

Embedding Model Tối ưu cho Code

Embedding (vector biểu diễn ngữ nghĩa) là bước biến mỗi chunk code thành một vector số học để có thể so sánh độ tương đồng. Chọn sai embedding model là một trong những nguyên nhân phổ biến nhất khiến retrieval "trông đúng logic nhưng kết quả tệ".

General-purpose embedding vs. embedding chuyên cho code

Các embedding model đa dụng như OpenAI text-embedding-3-large/text-embedding-3-small được huấn luyện chủ yếu trên văn bản tự nhiên (có một phần code trong dữ liệu huấn luyện, nhưng không phải trọng tâm). Chúng vẫn hoạt động khá tốt cho các câu hỏi mang tính "mô tả" (ví dụ "hàm xử lý thanh toán ở đâu") vì embedding này giỏi bắt ngữ nghĩa ngôn ngữ tự nhiên trong docstring/comment.

Ngược lại, các embedding model chuyên huấn luyện cho code như Voyage AI voyage-code-3 (hoặc các model tương tự như CodeBERT-family, jina-embeddings-v2-code) được huấn luyện trên lượng lớn cặp (code, code) và (natural language, code) — chúng nắm bắt tốt hơn các mẫu hình đặc trưng của code: quan hệ giữa tên biến và kiểu dữ liệu, cấu trúc lồng nhau, các pattern lập trình lặp lại (ví dụ loop, error handling). Trong benchmark truy hồi code (ví dụ CodeSearchNet), embedding chuyên biệt thường vượt embedding đa dụng đáng kể về độ chính xác top-k.

Bảng đánh đổi chi phí — chất lượng — số chiều

Số chiều (dimensionality) ảnh hưởng trực tiếp đến chi phí lưu trữ vector store và tốc độ tìm kiếm (search): vector 3072 chiều tốn gấp 12 lần dung lượng so với 256 chiều, nhưng độ chính xác tăng không tuyến tính — thường lợi ích giảm dần sau một ngưỡng nào đó (diminishing return). Với hầu hết trường hợp thực tế, 512–1024 chiều là điểm cân bằng hợp lý.

from openai import OpenAI

client = OpenAI()

def embed_openai(code_chunks: list[str]) -> list[list[float]]:
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=code_chunks,
        dimensions=1024,  # cắt giảm số chiều để tiết kiệm chi phí lưu trữ
    )
    return [item.embedding for item in response.data]


import voyageai

vo = voyageai.Client()

def embed_voyage_code(code_chunks: list[str]) -> list[list[float]]:
    result = vo.embed(
        code_chunks,
        model="voyage-code-3",
        input_type="document",  # dùng "query" khi embed câu hỏi tìm kiếm
    )
    return result.embeddings

Một điểm cần lưu ý: đa số embedding API cho phép chỉ định input_type là document (khi embed dữ liệu để lưu vào vector store) hoặc query (khi embed câu hỏi của người dùng) — hai chế độ này thường dùng cách encode khác nhau bên trong model để tối ưu độ khớp giữa query ngắn và document dài. Bỏ qua chi tiết này là lỗi rất phổ biến khiến retrieval kém hơn kỳ vọng.

Mẹo

• Nếu ngân sách hạn chế hoặc codebase nhỏ (dưới vài nghìn file), bắt đầu với OpenAI embedding đa dụng — đủ tốt và dễ triển khai.
• Nếu retrieval cho các câu hỏi kỹ thuật sâu (API signature, pattern implementation) liên tục cho kết quả "gần đúng nhưng không trúng", hãy thử A/B so sánh với một embedding chuyên code như voyage-code-3 trước khi đổ lỗi cho bước chunking hoặc ranking.
• Luôn benchmark bằng chính dữ liệu và truy vấn thật của bạn (không chỉ tin vào leaderboard công khai) — đặc thù ngôn ngữ lập trình, style code của team, và loại câu hỏi thường gặp ảnh hưởng lớn đến việc model nào thực sự phù hợp.
• Đừng đổi embedding model giữa đường mà không re-embed lại toàn bộ vector store — vector từ hai model khác nhau không so sánh được với nhau.

Hybrid Retrieval: Kết hợp Semantic Search và Keyword Search

Semantic search (tìm kiếm theo ngữ nghĩa qua embedding) rất mạnh khi câu hỏi diễn đạt bằng ngôn ngữ tự nhiên và không trùng từ khóa chính xác với code (ví dụ hỏi "chỗ nào xử lý giảm giá cho khách VIP" mà code đặt tên là applyLoyaltyDiscount). Nhưng nó có một điểm yếu chí mạng: semantic search kém trong việc tìm chính xác một identifier, tên biến, tên function, hoặc error message cụ thể.

Lý do: embedding hoạt động dựa trên độ tương đồng ngữ nghĩa tổng quát, không phải khớp ký tự (exact match). Nếu bạn hỏi "hàm calculatePriceWithTax nằm ở đâu", một embedding model có thể trả về những chunk "nghe có vẻ liên quan" đến giá và thuế, nhưng lại bỏ sót chính chunk chứa định nghĩa calculatePriceWithTax nếu chunk đó không có đủ ngôn ngữ tự nhiên mô tả xung quanh. Đây chính là lúc keyword search (dựa trên BM25 — một thuật toán ranking dựa trên tần suất từ khóa, rất phổ biến trong search engine truyền thống) tỏa sáng: BM25 tìm chính xác token calculatePriceWithTax xuất hiện ở đâu, không quan tâm ngữ nghĩa.

Reciprocal Rank Fusion (RRF) — cách kết hợp hai kết quả

Cách phổ biến và hiệu quả để trộn hai danh sách kết quả (một từ semantic search, một từ BM25) là Reciprocal Rank Fusion (RRF) — thay vì cộng điểm số trực tiếp (vốn có scale khác nhau, khó so sánh), RRF chỉ dựa vào hạng (rank) của mỗi kết quả trong từng danh sách, sau đó cộng nghịch đảo hạng lại với nhau.

def reciprocal_rank_fusion(
    semantic_results: list[str],   # danh sách chunk_id, đã sắp theo độ liên quan giảm dần
    keyword_results: list[str],
    k: int = 60,                   # hằng số làm mượt, giá trị 60 là mặc định phổ biến
) -> list[str]:
    scores: dict[str, float] = {}

    for rank, chunk_id in enumerate(semantic_results):
        scores[chunk_id] = scores.get(chunk_id, 0) + 1 / (k + rank + 1)

    for rank, chunk_id in enumerate(keyword_results):
        scores[chunk_id] = scores.get(chunk_id, 0) + 1 / (k + rank + 1)

    return sorted(scores, key=lambda cid: scores[cid], reverse=True)


def hybrid_retrieve(query: str, top_k: int = 8) -> list[dict]:
    # Semantic search: embed query, tìm k lân cận gần nhất trong vector store
    query_vector = embed_voyage_code([query])[0]
    semantic_hits = vector_store.similarity_search(query_vector, k=20)

    # Keyword search: BM25 trên toàn bộ chunk (dùng thư viện như rank_bm25 hoặc Elasticsearch)
    keyword_hits = bm25_index.search(query, k=20)

    fused_ids = reciprocal_rank_fusion(
        [c.id for c in semantic_hits],
        [c.id for c in keyword_hits],
    )
    return [chunk_store.get(cid) for cid in fused_ids[:top_k]]


results = hybrid_retrieve("hàm nào tính discount cho khách VIP", top_k=8)

Trong thực tế triển khai, nhiều vector store hiện đại (Weaviate, Qdrant, Elasticsearch với dense+sparse vector) đã hỗ trợ hybrid search built-in, bạn không cần tự viết BM25 index từ đầu — nhưng hiểu cơ chế RRF vẫn rất cần thiết để tune tham số k và debug khi kết quả không như mong đợi.

Mẹo

• Nếu câu hỏi chứa một identifier viết đúng chuẩn (camelCase, snake_case, hoặc trong dấu backtick), hãy tăng trọng số cho keyword search — đây là dấu hiệu rõ ràng người dùng đang tìm chính xác một symbol.
• Đừng bỏ semantic search hoàn toàn dù bạn thấy BM25 "có vẻ đủ" — với các câu hỏi diễn đạt tự nhiên (rất phổ biến từ PM/QA, ví dụ "chỗ nào xử lý khi user hủy đơn giữa lúc thanh toán"), semantic search vẫn là kênh chính.
• Benchmark hybrid retrieval bằng một bộ câu hỏi thật (golden set) có sẵn đáp án đúng — đo recall@k trước và sau khi thêm hybrid để chứng minh giá trị thực, tránh tối ưu "cảm tính".

Graph-Enhanced Retrieval: Truy hồi Theo Các Cạnh Phụ thuộc

Ngay cả với hybrid retrieval tốt, vẫn còn một lớp vấn đề mà cả semantic search và keyword search không giải quyết được: quan hệ cấu trúc (structural relationship). Khi bạn retrieve về một function gọi tới một function khác, bạn cũng cần định nghĩa của function được gọi đó — nhưng nó có thể nằm ở một file hoàn toàn khác, dùng từ ngữ hoàn toàn khác, và không hề "gần" về semantic hay khớp keyword với câu hỏi ban đầu.

Đây là lúc kết hợp RAG với dependency graph (đồ thị phụ thuộc — kỹ thuật đã đề cập ở các bài trước trong module này về repo map và targeted file inclusion) mang lại giá trị lớn. Ý tưởng: sau khi có kết quả retrieval ban đầu (từ hybrid search), thực hiện thêm một bước graph expansion — đi theo các cạnh phụ thuộc từ chunk vừa tìm được để lấy thêm các chunk "hàng xóm" quan trọng:

• Định nghĩa của các function/class được gọi trong chunk retrieve về (outgoing call edge).
• Nơi gọi đến function này (incoming call edge / reverse reference) — hữu ích khi câu hỏi dạng "nếu tôi sửa hàm này thì ảnh hưởng gì" (rất phổ biến trong câu hỏi của QA hoặc khi đánh giá rủi ro thay đổi).
• File import chunk đó, và các file mà chunk đó import.
• Test file tương ứng, nếu tồn tại quy ước đặt tên (ví dụ foo.py ↔ test_foo.py) — cực kỳ giá trị cho câu hỏi của QA engineer.

def graph_expand(initial_chunks: list[dict], call_graph: dict, max_hops: int = 1) -> list[dict]:
    """
    call_graph: dict ánh xạ chunk_id -> {"calls": [chunk_id...], "called_by": [chunk_id...], "imports": [chunk_id...]}
    """
    seen = {c["id"] for c in initial_chunks}
    expanded = list(initial_chunks)
    frontier = list(initial_chunks)

    for _ in range(max_hops):
        next_frontier = []
        for chunk in frontier:
            edges = call_graph.get(chunk["id"], {})
            neighbor_ids = edges.get("calls", []) + edges.get("called_by", []) + edges.get("imports", [])
            for nid in neighbor_ids:
                if nid not in seen:
                    seen.add(nid)
                    neighbor_chunk = chunk_store.get(nid)
                    expanded.append(neighbor_chunk)
                    next_frontier.append(neighbor_chunk)
        frontier = next_frontier

    return expanded


def graph_enhanced_retrieve(query: str, top_k: int = 8, max_hops: int = 1) -> list[dict]:
    base_results = hybrid_retrieve(query, top_k=top_k)
    return graph_expand(base_results, call_graph, max_hops=max_hops)

Lưu ý về max_hops: mỗi hop mở rộng làm tăng số chunk (và token) theo cấp lũy — hop=1 thường đủ cho hầu hết câu hỏi ("hàm này gọi gì, cái gì gọi hàm này"), hop=2 có thể cần cho câu hỏi phân tích tác động sâu (impact analysis) nhưng phải cân nhắc kỹ vì token cost tăng nhanh và tỷ lệ chunk không liên quan cũng tăng theo.

Mẹo

• Xây call_graph một lần (offline, khi index codebase) bằng cách parse AST và trích các lời gọi function/import, đừng tính toán lại graph theo từng query — quá tốn thời gian và token.
• Luôn đặt giới hạn cứng (hard limit) cho số chunk sau graph expansion (ví dụ tối đa 15-20 chunk) trước khi đưa vào prompt — expansion không kiểm soát dễ làm context window bị "ngập" bởi các chunk liên quan xa, làm loãng tín hiệu quan trọng.
• Ưu tiên expand theo cạnh "called_by" khi câu hỏi liên quan đến rủi ro/tác động thay đổi (rất phù hợp với câu hỏi QA/PM kiểu "nếu sửa cái này thì ảnh hưởng chỗ nào"), và ưu tiên "calls" khi câu hỏi liên quan đến việc hiểu logic implementation (phù hợp câu hỏi của developer).

Xây dựng Pipeline RAG cho Code với LangChain

Phần này ghép toàn bộ các khối đã học thành một pipeline hoàn chỉnh, dùng LangChain — framework phổ biến để build ứng dụng LLM có RAG. Pipeline gồm 5 bước: load codebase → chunk theo AST → embed → lưu vào vector store → xây retrieval chain và kiểm thử với truy vấn thật của từng persona.

from langchain_community.document_loaders import DirectoryLoader
from langchain_core.documents import Document
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever

def load_and_chunk_repo(repo_path: str) -> list[Document]:
    documents = []
    for file_path in _walk_source_files(repo_path, extensions=[".py"]):
        source_code = open(file_path, "rb").read()
        for chunk in extract_chunks(source_code, file_path):
            documents.append(Document(
                page_content=chunk["content"],
                metadata={
                    "file_path": chunk["file_path"],
                    "name": chunk["name"],
                    "start_line": chunk["start_line"],
                    "end_line": chunk["end_line"],
                    "node_type": chunk["node_type"],
                },
            ))
    return documents

documents = load_and_chunk_repo("./my-project")

embeddings = OpenAIEmbeddings(model="text-embedding-3-small", dimensions=1024)
vector_store = Chroma.from_documents(documents, embeddings, collection_name="codebase_rag")

vector_retriever = vector_store.as_retriever(search_kwargs={"k": 10})

bm25_retriever = BM25Retriever.from_documents(documents)
bm25_retriever.k = 10

hybrid_retriever = EnsembleRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    weights=[0.6, 0.4],  # ưu tiên semantic nhẹ hơn, nhưng vẫn giữ trọng số cho keyword match
)

llm = ChatOpenAI(model="gpt-4.1-mini", temperature=0)

qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=hybrid_retriever,
    chain_type="stuff",
    return_source_documents=True,
)


engineer_query = "Hàm applyLoyaltyDiscount đang tính discount theo công thức nào, và nó gọi tới hàm nào khác?"
engineer_answer = qa_chain.invoke({"query": engineer_query})
print(engineer_answer["result"])
for doc in engineer_answer["source_documents"]:
    print(f"- {doc.metadata['file_path']}:{doc.metadata['start_line']}")

qa_query = "Khi user hủy đơn hàng ngay lúc đang xử lý thanh toán, hệ thống xử lý race condition này như thế nào?"
qa_answer = qa_chain.invoke({"query": qa_query})
print(qa_answer["result"])

pm_query = "Tính năng giảm giá cho khách VIP hiện đang áp dụng cho những nhóm khách hàng nào, và có giới hạn số lần dùng không?"
pm_answer = qa_chain.invoke({"query": pm_query})
print(pm_answer["result"])

Vài điểm cần chú ý khi vận hành pipeline này trong thực tế:

• return_source_documents=True không chỉ để debug — nó cho phép hiển thị "trích dẫn nguồn" (đường dẫn file, số dòng) kèm câu trả lời, giúp mọi persona (đặc biệt là PM/QA không đọc code hàng ngày) tự kiểm chứng được câu trả lời có đáng tin không, thay vì tin tuyệt đối vào LLM.
• Trọng số ensemble (weights=[0.6, 0.4]) nên được tinh chỉnh dựa trên loại truy vấn phổ biến nhất trong team bạn — team có nhiều câu hỏi kiểu "tính năng X hoạt động ra sao" (PM/QA) có thể cần tăng trọng số semantic; team có nhiều câu hỏi tra cứu symbol cụ thể (dev) có thể cần tăng trọng số BM25.
• Chain "stuff" (nhồi toàn bộ chunk retrieve về vào một prompt) chỉ phù hợp khi tổng token của các chunk còn nằm trong giới hạn context window an toàn. Với codebase lớn và graph expansion sâu, hãy dùng chain "map-reduce" hoặc "refine" của LangChain để tránh tràn context, hoặc giới hạn top_k chặt hơn.
• Muốn thêm graph-enhanced retrieval vào pipeline này, chỉ cần bọc hybrid_retriever bằng một custom retriever gọi graph_expand() sau bước retrieve gốc, trước khi trả kết quả cho qa_chain.

Mẹo

• Bắt đầu nhỏ: index một module hoặc một service trước, đo chất lượng retrieval bằng một bộ câu hỏi thật của cả ba persona, rồi mới scale ra toàn bộ monorepo — tránh tốn chi phí embedding cho toàn repo trước khi biết pipeline có hoạt động đúng không.
• Log lại mọi câu query cùng với chunk được retrieve và câu trả lời cuối — đây là dữ liệu quý để sau này fine-tune trọng số ensemble hoặc phát hiện lỗ hổng trong chunking.
• Với PM/QA, xây một lớp giao diện đơn giản (chatbot nội bộ, Slack bot) gọi qa_chain thay vì để họ tự chạy code Python — mục tiêu của RAG cho code là dân chủ hóa quyền truy cập kiến thức codebase cho toàn bộ team, không chỉ cho developer.
• Định kỳ re-index (incremental, chỉ re-embed các file thay đổi qua git diff) để vector store không bị "lệch" so với code thực tế — codebase RAG lỗi thời gây hại nhiều hơn lợi, vì câu trả lời trông đáng tin nhưng lại dựa trên code đã cũ.