code_parser.py

import ast
import sqlite3
from sentence_transformers import SentenceTransformer

# Khởi tạo mô hình nhúng vector (Sử dụng mô hình nhỏ gọn, nhẹ cho máy)
model = SentenceTransformer('all-MiniLM-L6-v2')

def extract_functions_from_code(code_content):
    """Bóc tách các hàm từ file code .py sử dụng thư viện AST"""
    chunks = []
    try:
        tree = ast.parse(code_content)
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                # Lấy toàn bộ đoạn code của hàm đó
                func_code = ast.get_source_segment(code_content, node)
                chunks.append({
                    "name": node.name,
                    "line": node.lineno,
                    "code": func_code
                })
    except Exception as e:
        print(f"Lỗi parse code: {e}")
    return chunks

def process_and_store_file(file_content, file_name, user_id):
    """Xử lý file code, kiểm tra trùng lặp từng hàm, trả về kết quả chi tiết dưới dạng dict"""
    chunks = extract_functions_from_code(file_content)
    if not chunks:
        return {"status": "error", "message": "❌ Không tìm thấy hàm Python nào hợp lệ trong file này!"}
        
    conn = sqlite3.connect("code_rag_repository.db")
    cursor = conn.cursor()
    
    # [TỐI ƯU HIỆU NĂNG]: Lấy toàn bộ vector hiện tại của user ra bộ nhớ trước để đối chiếu, tránh quét DB trong vòng lặp
    cursor.execute("SELECT file_name, class_name, line_number, vector_data FROM code_vectors WHERE user_id = ?", (user_id,))
    all_rows = cursor.fetchall()
    
    saved_count = 0
    duplicated_functions = []
    SIMILARITY_THRESHOLD = 0.38 # Ngưỡng phát hiện trùng
    
    for chunk in chunks:
        # 1. Biến đổi đoạn code hiện tại thành vector để chuẩn bị đối chiếu
        current_vector = model.encode(chunk['code']).tolist()
        
        is_duplicated = False
        duplicate_info = ""
        
        # 2. Đối chiếu trực tiếp trên mảng dữ liệu đã lưu ở RAM
        for db_file, db_class, db_line, db_vector_str in all_rows:
            try:
                db_vector = [float(x) for x in db_vector_str.split(",")]
            except:
                continue
            
            # Tính độ tương đồng Cosine
            dot_product = sum(a * b for a, b in zip(current_vector, db_vector))
            norm_a = sum(a * a for a in current_vector) ** 0.5
            norm_b = sum(b * b for b in db_vector) ** 0.5
            score = dot_product / (norm_a * norm_b) if (norm_a * norm_b) > 0 else 0
            
            # Nếu phát hiện độ tương đồng vượt ngưỡng, đánh dấu là trùng lặp
            if score >= SIMILARITY_THRESHOLD:
                is_duplicated = True
                duplicate_info = f"Hàm '{chunk['name']}' (dòng {chunk['line']}) giống {score:.1%} với hàm '{db_class}' trong file '{db_file}' (dòng {db_line})"
                break
        
        # 3. Quyết định lưu hay chặn
        if is_duplicated:
            duplicated_functions.append(duplicate_info)
        else:
            # Nếu code mới hoàn toàn thì lưu vào DB
            vector_str = ",".join(map(str, current_vector))
            cursor.execute('''
                INSERT INTO code_vectors (user_id, file_name, class_name, line_number, embedding_text, vector_data)
                VALUES (?, ?, ?, ?, ?, ?)
            ''', (user_id, file_name, chunk['name'], chunk['line'], chunk['code'], vector_str))
            saved_count += 1
            
            # Đồng thời cập nhật luôn vào bộ nhớ tạm all_rows để các hàm phía sau trong cùng file đối chiếu chéo được luôn
            all_rows.append((file_name, chunk['name'], chunk['line'], vector_str))
            
    conn.commit()
    conn.close()
    
    # 4. Trả về kết quả cấu trúc tường minh để app.py bắt lỗi chính xác
    if saved_count == 0 and duplicated_functions:
        return {
            "status": "duplicated",
            "message": "Toàn bộ hàm trong file đều bị trùng lặp!",
            "duplicate_info": "\n".join(duplicated_functions)
        }
    
    return {
        "status": "success",
        "saved_count": saved_count,
        "duplicated_count": len(duplicated_functions),
        "duplicate_info": "\n".join(duplicated_functions) if duplicated_functions else ""
    }

def retrieve_context_from_sqlite(user_query, user_id, top_k=1):
    """Tìm kiếm đoạn code trùng khớp nhất dựa trên độ tương đồng Cosine"""
    conn = sqlite3.connect("code_rag_repository.db")
    cursor = conn.cursor()
    cursor.execute("SELECT file_name, class_name, line_number, embedding_text, vector_data FROM code_vectors WHERE user_id=?", (user_id,))
    rows = cursor.fetchall()
    conn.close()
    
    if not rows:
        return []
        
    query_vector = model.encode(user_query).tolist()
    matched_results = []
    
    for file_name, class_name, line_number, embedding_text, vector_data_str in rows:
        try:
            db_vector = [float(x) for x in vector_data_str.split(",")]
        except:
            continue
        
        dot_product = sum(a * b for a, b in zip(query_vector, db_vector))
        norm_a = sum(a * a for a in query_vector) ** 0.5
        norm_b = sum(b * b for b in db_vector) ** 0.5
        score = dot_product / (norm_a * norm_b) if (norm_a * norm_b) > 0 else 0
        
        matched_results.append({
            "file_name": file_name,
            "class_name": class_name,
            "line_number": line_number,
            "embedding_text": embedding_text,
            "score": score
        })
        
    matched_results.sort(key=lambda x: x['score'], reverse=True)
    return matched_results[:top_k]

def delete_file_from_sqlite(file_name, user_id):
    """Xóa toàn bộ các hàm đã lưu của một file cụ thể khỏi SQLite"""
    try:
        conn = sqlite3.connect("code_rag_repository.db")
        cursor = conn.cursor()
        cursor.execute("DELETE FROM code_vectors WHERE file_name = ? AND user_id = ?", (file_name, user_id))
        conn.commit()
        conn.close()
        return f"🗑️ Đã xóa sạch dữ liệu của file '{file_name}' khỏi hệ thống!"
    except Exception as e:
        return f"❌ Lỗi khi xóa file: {e}"

def get_uploaded_files(user_id):
    """Lấy danh sách tên các file độc nhất đã được lưu trong SQLite của user"""
    try:
        conn = sqlite3.connect("code_rag_repository.db")
        cursor = conn.cursor()
        cursor.execute("SELECT DISTINCT file_name FROM code_vectors WHERE user_id = ?", (user_id,))
        rows = cursor.fetchall()
        conn.close()
        return [row[0] for row in rows]
    except Exception as e:
        return []