Local SE Agent Masterclass

# 1.1 ไฟล์ requirements.txt (อัปเดตเพิ่ม chromadb)
ollama>=0.4.0
chromadb>=0.4.24

# -------------------------------------------------------------

# 1.2 ไฟล์ memory/__init__.py (อัปเดตให้ Export LocalVectorStore ด้วย)
"""
Package บริหารจัดการหน่วยความจำทั้งระยะสั้น (Sliding Window) และระยะยาว (ChromaDB)
"""
from memory.context import SlidingWindowContext
from memory.vector_store import LocalVectorStore

__all__ = ["SlidingWindowContext", "LocalVectorStore"]

import chromadb
from pathlib import Path
from typing import List, Dict, Any, Optional

class LocalVectorStore:
    """
    คลังความจำระยะยาว (Long-Term Semantic Code Store)
    รันบนฮาร์ดแวร์ Mac 16GB แบบ 100% Local โดยไม่ต้องต่ออินเทอร์เน็ตภายนอก
    """
    def __init__(self, persist_dir: str = "./.chroma_db", collection_name: str = "codebase_index"):
        """
        :param persist_dir: ไดเรกทอรีบนดิสก์สำหรับบันทึก SQLite/Parquet
        :param collection_name: ชื่อคอลเลกชันที่ใช้จัดเก็บข้อมูลโค้ด
        """
        self.persist_path = Path(persist_dir).resolve()
        # 1. สร้าง Persistent Client: บันทึกข้อมูลถาวรลงดิสก์ ไม่หายเมื่อปิดโปรแกรม
        self.client = chromadb.PersistentClient(path=str(self.persist_path))
        
        # 2. ดึงหรือสร้าง Collection ใหม่ (ChromaDB จะใช้โมเดล Embedding เริ่มต้นแบบ Built-in อัตโนมัติ)
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"description": "Source code semantic vectors for SE Agent"}
        )

    def index_directory(self, root_dir: str, file_extensions: Optional[List[str]] = None) -> int:
        """
        สแกนไดเรกทอรีโปรเจกต์ อ่านไฟล์โค้ด แปลงเป็น Embeddings และบันทึกลง ChromaDB
        
        :param root_dir: ไดเรกทอรีโปรเจกต์เป้าหมาย เช่น './sample_maven_project'
        :param file_extensions: นามสกุลไฟล์ที่ต้องการทำดัชนี เช่น ['.java', '.py', '.xml']
        :return: จำนวนไฟล์ที่ถูกทำดัชนีสำเร็จ
        """
        if file_extensions is None:
            file_extensions = [".java", ".py", ".xml", ".md"]

        target_root = Path(root_dir).resolve()
        if not target_root.exists():
            raise FileNotFoundError(f"ไม่พบไดเรกทอรีที่ต้องการ Index: {target_root}")

        documents: List[str] = []
        metadatas: List[Dict[str, Any]] = []
        ids: List[str] = []

        # ท่องไฟล์ทั้งหมดแบบ Recursive (ข้ามโฟลเดอร์ build/cache)
        for filepath in target_root.rglob("*"):
            if filepath.is_file() and filepath.suffix in file_extensions:
                # ข้ามไฟล์ในโฟลเดอร์ที่ไม่จำเป็น เช่น target/, .git/, .venv/
                if any(ignored in filepath.parts for ignored in ["target", ".git", ".venv", ".chroma_db"]):
                    continue

                try:
                    content = filepath.read_text(encoding="utf-8", errors="replace").strip()
                    if not content:
                        continue

                    # ใช้ relative path เป็น ID และ Metadata
                    rel_path = str(filepath.relative_to(target_root))
                    documents.append(content)
                    metadatas.append({"filepath": rel_path, "filename": filepath.name})
                    ids.append(rel_path)
                except Exception as err:
                    print(f"⚠️ ข้ามไฟล์ {filepath.name} เนื่องจากอ่านไม่สำเร็จ: {err}")

        if documents:
            # upsert: เพิ่มข้อมูลใหม่ หรืออัปเดตทับถ้ามี ID เดิมอยู่แล้ว (Idempotent)
            self.collection.upsert(
                documents=documents,
                metadatas=metadatas,
                ids=ids
            )

        return len(documents)

    def query_similar_code(self, query_text: str, top_k: int = 2) -> List[Dict[str, Any]]:
        """
        ค้นหาไฟล์โค้ดที่มีความหมายสอดคล้องกับคำถาม (Semantic Code Search)
        
        :param query_text: คำค้นหา เช่น 'ฟังก์ชันคำนวณส่วนลดราคาสินค้า' หรือ 'database connection error'
        :param top_k: จำนวนผลลัพธ์ที่ใกล้เคียงที่สุดที่ต้องการดึงออกมา
        :return: ลิสต์ผลลัพธ์ที่มี filepath, snippet, และ distance
        """
        if self.collection.count() == 0:
            return []

        # ยิงค้นหาด้วย Semantic Vector Distance (Cosine/L2)
        results = self.collection.query(
            query_texts=[query_text],
            n_results=min(top_k, self.collection.count())
        )

        formatted_results: List[Dict[str, Any]] = []
        if results and results.get("ids"):
            matched_ids = results["ids"][0]
            matched_metadatas = results["metadatas"][0] if results.get("metadatas") else []
            matched_documents = results["documents"][0] if results.get("documents") else []
            matched_distances = results["distances"][0] if results.get("distances") else []

            for idx in range(len(matched_ids)):
                formatted_results.append({
                    "id": matched_ids[idx],
                    "filepath": matched_metadatas[idx].get("filepath", matched_ids[idx]),
                    "filename": matched_metadatas[idx].get("filename", ""),
                    "content_snippet": matched_documents[idx][:300], # ตัดทอนตัวอย่างโค้ด
                    "distance": matched_distances[idx] if idx < len(matched_distances) else 0.0
                })

        return formatted_results

    def count(self) -> int:
        """นับจำนวนชิ้นส่วนโค้ดที่มีอยู่ใน Vector Store"""
        return self.collection.count()

import shutil
from pathlib import Path
from memory.vector_store import LocalVectorStore

def setup_mock_project(mock_dir: Path):
    """สร้างคลังไฟล์โค้ดตัวอย่าง 4 ไฟล์เพื่อทดสอบความฉลาดของ Semantic Search"""
    mock_dir.mkdir(parents=True, exist_ok=True)
    
    # ไฟล์ 1: บริการคิดเงินและส่วนลด
    (mock_dir / "Calculator.java").write_text("""
    public class Calculator {
        public double applyDiscount(double price, double discount) {
            // คำนวณส่วนลดราคาสินค้าตามโปรโมชั่น
            return price - discount;
        }
    }
    """, encoding="utf-8")

    # ไฟล์ 2: บริการจัดการผู้ใช้งาน
    (mock_dir / "UserService.java").write_text("""
    public class UserService {
        public void registerUser(String username, String email) {
            // สมัครสมาชิกใหม่และบันทึกลงระบบ
            System.out.println("User registered: " + username);
        }
    }
    """, encoding="utf-8")

    # ไฟล์ 3: บริการเชื่อมต่อฐานข้อมูล
    (mock_dir / "DatabaseConfig.java").write_text("""
    public class DatabaseConfig {
        public String getJdbcUrl() {
            // ตั้งค่า URL สำหรับเชื่อมต่อ PostgreSQL Database
            return "jdbc:postgresql://localhost:5432/production_db";
        }
    }
    """, encoding="utf-8")

    # ไฟล์ 4: คู่มือโปรเจกต์
    (mock_dir / "README.md").write_text("""
    # E-Commerce Core Engine
    ระบบแกนกลางสำหรับประมวลผลคำสั่งซื้อและตะกร้าสินค้า
    """, encoding="utf-8")

def main():
    print("🚀 เริ่มทดสอบ Day 9: ChromaDB Local Vector Store บน Mac M3...")

    mock_project_dir = Path("./temp_mock_codebase").resolve()
    mock_db_dir = Path("./temp_test_chroma_db").resolve()

    try:
        # 1. จัดเตรียมโฟลเดอร์โปรเจกต์จำลอง
        setup_mock_project(mock_project_dir)
        print(f"📁 สร้างคลังโค้ดจำลองที่: {mock_project_dir.name}")

        # 2. เริ่มต้นระบบ LocalVectorStore
        vector_store = LocalVectorStore(
            persist_dir=str(mock_db_dir),
            collection_name="test_code_collection"
        )

        # 3. สั่งสแกนและทำ Index ไฟล์โค้ดทั้งหมด
        indexed_count = vector_store.index_directory(str(mock_project_dir))
        print(f"📦 ทำดัชนีโค้ดสำเร็จ: {indexed_count} ไฟล์ (นับรวมใน Vector DB: {vector_store.count()} ชิ้น)")

        # 4. ทดสอบค้นหาด้วยภาษาธรรมชาติ (Semantic Queries ที่ไม่มีคำในโค้ดตรงๆ!)
        test_queries = [
            ("ฟังก์ชันลดราคาสินค้าและคำนวณเงินโปรโมชั่น", "Calculator.java"),
            ("ระบบสร้างบัญชีสมาชิกและข้อมูลอีเมลผู้ใช้", "UserService.java"),
            ("การตั้งค่าพอร์ตและเชื่อมต่อ PostgreSQL", "DatabaseConfig.java")
        ]

        print("\n" + "=" * 60)
        print("🔍 เริ่มทดสอบ Semantic Retrieval (ค้นหาตามความหมาย):")
        print("=" * 60)

        for query_text, expected_file in test_queries:
            results = vector_store.query_similar_code(query_text, top_k=1)
            assert len(results) > 0, f"❌ ไม่พบผลลัพธ์สำหรับ: '{query_text}'"

            top_match = results[0]
            print(f"💬 Query: \"{query_text}\"")
            print(f"  🎯 Top Matched: {top_match['filename']} (Distance: {top_match['distance']:.4f})")
            
            # ตรวจสอบความถูกต้อง
            assert top_match["filename"] == expected_file, (
                f"❌ ล้มเหลว! คาดหวัง '{expected_file}' แต่ระบบดึง '{top_match['filename']}'"
            )
            print("  ✅ ถูกต้อง 100%!\n")

        print("=" * 60)
        print("🎉 สำเร็จ 100%! ChromaDB สามารถค้นหาไฟล์โค้ดได้ตามความหมายโดยไม่ต้องเดาชื่อไฟล์!")
        print("💡 เมื่อ Agent ได้รับโจทย์ มันจะสกัดเฉพาะไฟล์เป้าหมายนี้ส่งเข้า LLM ไม่ต้องอ่านทั้งโปรเจกต์!")

    finally:
        # ทำความสะอาดโฟลเดอร์ทดสอบชั่วคราว
        shutil.rmtree(mock_project_dir, ignore_errors=True)
        shutil.rmtree(mock_db_dir, ignore_errors=True)

if __name__ == "__main__":
    main()