Local SE Agent Masterclass

from collections import defaultdict
from typing import List, Dict, Any, Optional
from pathlib import Path

class HybridRetrievalEngine:
    """
    เอนจินสืบค้นโค้ดแบบผสมผสาน (Hybrid Retrieval Engine)
    ผสานพลังระหว่าง:
    1. Vector Search (ChromaDB): ค้นหาตามความหมายของปัญหา (Semantic Intent)
    2. AST Symbol Index (Tree-sitter): ค้นหาตามชื่อคลาส/เมธอดตรงจุด (Exact Identifiers)
    ผ่านอัลกอริทึม Reciprocal Rank Fusion (RRF, k=60)
    """
    def __init__(self, vector_store=None, ast_parser=None, rrf_k: int = 60):
        self.vector_store = vector_store
        self.ast_parser = ast_parser
        self.rrf_k = rrf_k
        # เก็บตาราง Symbol -> List of Filepaths
        self.symbol_index: Dict[str, List[str]] = defaultdict(list)

    def register_symbol(self, symbol_name: str, filepath: str) -> None:
        """บันทึกชื่อสัญลักษณ์ (คลาส, เมธอด) เข้าสู่ตารางดัชนีทางไวยากรณ์"""
        symbol_clean = symbol_name.lower().strip()
        if filepath not in self.symbol_index[symbol_clean]:
            self.symbol_index[symbol_clean].append(filepath)

    def index_ast_symbols(self, files: List[Path]) -> int:
        """สกัดชื่อคลาสและเมธอดจากไฟล์ด้วย Tree-sitter เพื่อสร้าง Symbol Index"""
        total_symbols = 0
        for file in files:
            code = file.read_text(encoding="utf-8", errors="replace")
            # ถ้ามี ast_parser ใช้ Tree-sitter แกะชื่อเมธอด
            if self.ast_parser:
                tree, source_bytes = self.ast_parser.parse_code(code)
                def visit(node):
                    nonlocal total_symbols
                    if node.type in ("class_declaration", "method_declaration"):
                        name_node = node.child_by_field_name("name")
                        if name_node:
                            sym_text = source_bytes[name_node.start_byte:name_node.end_byte].decode("utf8")
                            self.register_symbol(sym_text, file.name)
                            total_symbols += 1
                    for child in node.children:
                        visit(child)
                visit(tree.root_node)
            else:
                # Fallback พื้นฐาน
                self.register_symbol(file.stem, file.name)
                total_symbols += 1
        return total_symbols

    def search_symbols(self, target_symbols: List[str]) -> List[str]:
        """ค้นหาไฟล์ที่มีสัญลักษณ์เป้าหมายตรงกัน คืนค่าเรียงตามความถี่"""
        file_hits: Dict[str, int] = defaultdict(int)
        for sym in target_symbols:
            sym_clean = sym.lower().strip()
            # ตรวจสอบว่าตรงเป๊ะ หรือเป็นคำย่อย
            for indexed_sym, files in self.symbol_index.items():
                if sym_clean in indexed_sym or indexed_sym in sym_clean:
                    for f in files:
                        file_hits[f] += 1

        # เรียงตามจำนวนสัญลักษณ์ที่ตรงมากที่สุด
        sorted_files = sorted(file_hits.keys(), key=lambda f: file_hits[f], reverse=True)
        return sorted_files

    def reciprocal_rank_fusion(
        self,
        vector_candidates: List[str],
        symbol_candidates: List[str]
    ) -> List[Dict[str, Any]]:
        """
        [RRF Algorithm Master Formula]:
        Score(d) = sum( 1.0 / (k + rank) )
        :param vector_candidates: รายชื่อไฟล์เรียงจากคะแนนเวกเตอร์สูงสุด
        :param symbol_candidates: รายชื่อไฟล์เรียงจากความตรงของสัญลักษณ์สูงสุด
        :return: รายการผลลัพธ์จัดอันดับใหม่ (Re-ranked) พร้อมคะแนน RRF รวม
        """
        rrf_scores: Dict[str, float] = defaultdict(float)
        rank_details: Dict[str, Dict[str, int]] = defaultdict(dict)

        # 1. ให้คะแนนจาก Vector Search Rank
        for rank, file_id in enumerate(vector_candidates):
            score = 1.0 / (self.rrf_k + rank + 1)
            rrf_scores[file_id] += score
            rank_details[file_id]["vector_rank"] = rank + 1

        # 2. ให้คะแนนจาก AST Symbol Search Rank
        for rank, file_id in enumerate(symbol_candidates):
            score = 1.0 / (self.rrf_k + rank + 1)
            rrf_scores[file_id] += score
            rank_details[file_id]["symbol_rank"] = rank + 1

        # 3. จัดอันดับไฟล์ที่มีคะแนน RRF สูงสุด
        sorted_candidates = sorted(rrf_scores.items(), key=lambda item: item[1], reverse=True)

        final_results = []
        for rank_idx, (file_id, total_score) in enumerate(sorted_candidates, start=1):
            final_results.append({
                "final_rank": rank_idx,
                "file": file_id,
                "rrf_score": round(total_score, 6),
                "vector_rank": rank_details[file_id].get("vector_rank", None),
                "symbol_rank": rank_details[file_id].get("symbol_rank", None)
            })

        return final_results

    def hybrid_search(
        self,
        query_text: str,
        target_symbols: Optional[List[str]] = None,
        top_k: int = 3
    ) -> List[Dict[str, Any]]:
        """
        ค้นหาแบบรวมพลัง (Hybrid Execution):
        รัน Vector Search + Symbol Search ➔ ผสานด้วย RRF ➔ ตัดเฉพาะ top_k
        """
        # 1. ค้นหาทางเวกเตอร์ (ถ้าเชื่อมต่อ ChromaDB ไว้)
        vector_candidates = []
        if self.vector_store:
            raw_vec = self.vector_store.query_similar_code(query_text, top_k=10)
            vector_candidates = [hit["filename"] for hit in raw_vec]

        # 2. ค้นหาทางสัญลักษณ์
        symbol_candidates = self.search_symbols(target_symbols or [query_text])

        # 3. รวมพลังด้วย RRF
        ranked_hits = self.reciprocal_rank_fusion(vector_candidates, symbol_candidates)
        return ranked_hits[:top_k]

"""
Intelligence Package: รวบรวมเอนจิน AST Parsing, Code Skeletonization, Call Graph และ Hybrid Search
[อัปเดต Day 20]: เพิ่ม HybridRetrievalEngine
"""
from intelligence.ast_parser import JavaASTParser
from intelligence.call_graph import CallHierarchyGraph
from intelligence.hybrid_retrieval import HybridRetrievalEngine

__all__ = ["JavaASTParser", "CallHierarchyGraph", "HybridRetrievalEngine"]

import time
import shutil
from pathlib import Path
from intelligence.ast_parser import JavaASTParser
from intelligence.hybrid_retrieval import HybridRetrievalEngine

def setup_mock_project(base_dir: Path):
    """จำลองโปรเจกต์ 5 ไฟล์ที่มีทั้งไฟล์ DTO, Job, Service, และ Controller"""
    base_dir.mkdir(parents=True, exist_ok=True)

    # 1. ไฟล์เป้าหมายจริง: ตัดรอบบิลรายเดือน
    (base_dir / "SubscriptionJob.java").write_text("""
    package com.enterprise.job;
    public class SubscriptionJob {
        public void executeMonthlyBilling(Long subId) {
            System.out.println("Executing subscription recurring payment...");
        }
    }
    """, encoding="utf-8")

    # 2. ไฟล์อธิบายแพลน (มักหลอก Vector Search ให้มาจับอันดับ 1)
    (base_dir / "BillingPlanDoc.java").write_text("""
    package com.enterprise.doc;
    public class BillingPlanDoc {
        // Document: This class describes customer monthly recurring billing policies
        public String getTerms() { return "Terms and conditions"; }
    }
    """, encoding="utf-8")

    # 3. ไฟล์ชำระเงินทั่วไป
    (base_dir / "PaymentService.java").write_text("""
    package com.enterprise.service;
    public class PaymentService {
        public boolean processPayment(Double amount) { return true; }
    }
    """, encoding="utf-8")

    # 4. คอนโทรลเลอร์รับคำสั่ง
    (base_dir / "OrderController.java").write_text("""
    package com.enterprise.controller;
    public class OrderController {
        public void checkout() { System.out.println("Checkout"); }
    }
    """, encoding="utf-8")

def main():
    print("=" * 70)
    print("🚀 เริ่มทดสอบ Day 20: Hybrid Retrieval Engine (Vector + AST Symbol Index)")
    print("=" * 70)

    mock_dir = Path("./temp_hybrid_test").resolve()
    setup_mock_project(mock_dir)
    print(f"📁 1. สร้างโปรเจกต์จำลองสำเร็จที่: {mock_dir.name}")

    try:
        # 2. เริ่มต้นระบบ AST Parser และ Hybrid Retrieval Engine
        parser = JavaASTParser()
        hybrid_engine = HybridRetrievalEngine(ast_parser=parser, rrf_k=60)

        # 3. สแกนสร้าง AST Symbol Index จากไฟล์ Java ทั้งหมด
        java_files = list(mock_dir.glob("*.java"))
        start_time = time.perf_counter()
        total_syms = hybrid_engine.index_ast_symbols(java_files)
        elapsed_ms = (time.perf_counter() - start_time) * 1000

        print(f"⚡ 2. สกัด AST Symbols สำเร็จ: {total_syms} สัญลักษณ์ ในเวลา {elapsed_ms:.2f} ms!")

        # 4. จำลองสถานการณ์การค้นหา (Simulation):
        # ปัญหา: ผู้ใช้ถามแนวคิดภาษาธรรมชาติ + อ้างชื่อฟังก์ชันเฉพาะเจาะจง
        # Vector Search ลำพังมักจะเทใจให้ 'BillingPlanDoc.java'
        # AST Symbol Index ลำพังจับคู่ 'SubscriptionJob.java' เพราะมี executeMonthlyBilling
        print("\n" + "-" * 50)
        print("🎯 3. ทดสอบการจำลองค้นหา (Simulation Case):")
        print("   💬 User Query: \"fix recurring billing calculation error\"")
        print("   🔍 Extracted Symbol: [\"executeMonthlyBilling\"]")
        print("-" * 50)

        # จำลองผลลัพธ์จาก Vector Search (Dense)
        simulated_vector_rank = [
            "BillingPlanDoc.java",   # Rank 1 (โดนหลอกเพราะมีคำว่า monthly recurring billing ใน doc)
            "SubscriptionJob.java",  # Rank 2 (เป้าหมายแท้จริง)
            "PaymentService.java",   # Rank 3
            "OrderController.java"   # Rank 4
        ]

        # ค้นหาจาก AST Symbol Index (Sparse)
        symbol_hits = hybrid_engine.search_symbols(["executeMonthlyBilling"])
        print(f"📊 ผลลัพธ์จาก AST Symbol Search: {symbol_hits}")

        # 5. รวมพลังด้วย Reciprocal Rank Fusion (RRF)
        start_fusion = time.perf_counter()
        final_ranks = hybrid_engine.reciprocal_rank_fusion(
            vector_candidates=simulated_vector_rank,
            symbol_candidates=symbol_hits
        )
        fusion_ms = (time.perf_counter() - start_fusion) * 1000

        print(f"\n🏆 4. ผลลัพธ์การจัดอันดับใหม่ (RRF Final Re-ranking) ในเวลา {fusion_ms:.2f} ms:")
        print(f"{'Final Rank':<12} | {'File Name':<22} | {'RRF Score':<10} | {'Vec Rank':<9} | {'Sym Rank':<9}")
        print("-" * 70)
        for r in final_ranks:
            v_rank = str(r['vector_rank']) if r['vector_rank'] else "-"
            s_rank = str(r['symbol_rank']) if r['symbol_rank'] else "-"
            print(f"#{r['final_rank']:<11} | {r['file']:<22} | {r['rrf_score']:<10} | {v_rank:<9} | {s_rank:<9}")

        # 6. ยืนยันผลลัพธ์ว่าเป้าหมายแท้จริงถูกยกขึ้นอันดับ 1
        assert final_ranks[0]["file"] == "SubscriptionJob.java", "❌ ล้มเหลว! SubscriptionJob ไม่ได้ขึ้นอันดับ 1"
        print("-" * 70)
        print("\n🎉 ความสำเร็จ: SubscriptionJob.java พุ่งขึ้นเป็นอันดับที่ 1 ทันที!")
        print("   เพราะได้คะแนนหนุนทั้งจาก Vector (Rank 2) และ AST Symbol (Rank 1) รวมพลังกันอย่างแม่นยำ!")

        print("\n" + "=" * 70)
        print("🏆 ผ่านการทดสอบ Hybrid Retrieval Engine สมบูรณ์แบบ 100%!")
        print("🚀 คุณพร้อมเข้าสู่ Day 21: Week 1 Architecture Review & Integration Milestone!")
        print("=" * 70)

    finally:
        shutil.rmtree(mock_dir, ignore_errors=True)

if __name__ == "__main__":
    main()