dify-1.8.1/api/services/dataset_version_redis_service.py

"""
Document Segment Version Control Service based on Redis
Provides lightweight version management without database schema changes
"""

import hashlib
import json
from datetime import datetime
from typing import Any, Optional

from flask import request
from flask_login import current_user

from extensions.ext_redis import redis_client
from models.dataset import Dataset, Document, DocumentSegment


class RedisVersionService:
    """Redis-based version control service for document segments"""

    # Redis Key prefixes
    VERSION_PREFIX = "dataset:version"
    VERSION_LIST_PREFIX = "dataset:version:list"
    VERSION_COUNTER_PREFIX = "dataset:version:counter"
    VERSION_CURRENT_PREFIX = "dataset:version:current"

    # Default retention settings
    DEFAULT_TTL_DAYS = 30
    MAX_VERSIONS = 50

    @classmethod
    def _get_version_key(cls, segment_id: str, version_number: int) -> str:
        """Get Redis key for version data"""
        return f"{cls.VERSION_PREFIX}:{segment_id}:{version_number}"

    @classmethod
    def _get_version_list_key(cls, segment_id: str) -> str:
        """Get Redis key for version list (Sorted Set)"""
        return f"{cls.VERSION_LIST_PREFIX}:{segment_id}"

    @classmethod
    def _get_version_counter_key(cls, segment_id: str) -> str:
        """Get Redis key for version counter"""
        return f"{cls.VERSION_COUNTER_PREFIX}:{segment_id}"

    @classmethod
    def _get_current_hash_key(cls, segment_id: str) -> str:
        """Get Redis key for current content hash"""
        return f"{cls.VERSION_CURRENT_PREFIX}:{segment_id}"

    @classmethod
    def _generate_content_hash(cls, segment: DocumentSegment) -> str:
        """Generate content hash to detect real changes"""
        content_str = json.dumps(
            {
                "content": segment.content,
                "answer": segment.answer,
                "keywords": segment.keywords,
                "enabled": segment.enabled,
            },
            sort_keys=True,
        )
        return hashlib.md5(content_str.encode()).hexdigest()

    @classmethod
    def create_version_snapshot(
        cls,
        segment: DocumentSegment,
        document: Document,
        dataset: Dataset,
        change_type: str = "update",
        change_summary: Optional[str] = None,
        changed_fields: Optional[list[str]] = None,
    ) -> Optional[dict[str, Any]]:
        """
        Create version snapshot for a segment

        Args:
            segment: Document segment object
            document: Document object
            dataset: Dataset object
            change_type: Type of change (create/update/delete/rollback)
            change_summary: Summary of changes
            changed_fields: List of changed fields

        Returns:
            Version info dict, or None if content unchanged
        """
        segment_id = segment.id

        # Check if content actually changed
        current_hash = cls._generate_content_hash(segment)
        stored_hash = redis_client.get(cls._get_current_hash_key(segment_id))

        if stored_hash and stored_hash.decode() == current_hash and change_type == "update":
            # Content unchanged, don't create new version
            return None

        # Update current content hash
        redis_client.set(
            cls._get_current_hash_key(segment_id), current_hash, ex=cls.DEFAULT_TTL_DAYS * 86400
        )

        # Get next version number (atomic increment)
        counter_key = cls._get_version_counter_key(segment_id)
        version_number = redis_client.incr(counter_key)
        redis_client.expire(counter_key, cls.DEFAULT_TTL_DAYS * 86400)

        # Build version data
        version_data = {
            "version_number": version_number,
            "segment_id": segment_id,
            "document_id": document.id,
            "document_name": document.name,
            "dataset_id": dataset.id,
            "dataset_name": dataset.name,
            # Snapshot content
            "content": segment.content,
            "answer": segment.answer,
            "keywords": segment.keywords,
            "enabled": segment.enabled,
            "position": segment.position,
            "word_count": segment.word_count,
            # Change info
            "change_type": change_type,
            "change_summary": change_summary or f"{change_type.capitalize()} segment",
            "changed_fields": changed_fields or [],
            # Audit info
            "created_by": current_user.id if current_user and hasattr(current_user, "id") else None,
            "created_by_name": current_user.name if current_user and hasattr(current_user, "name") else "System",
            "created_at": datetime.utcnow().isoformat(),
            "created_ip": request.remote_addr if request else None,
            # Content hash
            "content_hash": current_hash,
        }

        # Store version data (Hash)
        version_key = cls._get_version_key(segment_id, version_number)
        redis_client.hset(
            version_key,
            mapping={
                k: json.dumps(v) if isinstance(v, (dict, list)) else (str(v) if v is not None else "")
                for k, v in version_data.items()
            },
        )
        redis_client.expire(version_key, cls.DEFAULT_TTL_DAYS * 86400)

        # Add to version list (Sorted Set, sorted by version number)
        list_key = cls._get_version_list_key(segment_id)
        redis_client.zadd(list_key, {str(version_number): version_number})
        redis_client.expire(list_key, cls.DEFAULT_TTL_DAYS * 86400)

        # Clean up old versions
        cls._cleanup_old_versions(segment_id)

        return version_data

    @classmethod
    def _cleanup_old_versions(cls, segment_id: str):
        """Clean up old versions, keep only recent ones"""
        list_key = cls._get_version_list_key(segment_id)

        # Get total count
        total_count = redis_client.zcard(list_key)

        if total_count > cls.MAX_VERSIONS:
            # Get versions to delete (oldest)
            to_delete_count = total_count - cls.MAX_VERSIONS
            old_versions = redis_client.zrange(list_key, 0, to_delete_count - 1)

            # Delete old version data
            for version_number in old_versions:
                version_key = cls._get_version_key(segment_id, int(version_number))
                redis_client.delete(version_key)

            # Remove from list
            redis_client.zremrangebyrank(list_key, 0, to_delete_count - 1)

    @classmethod
    def get_segment_versions(
        cls, segment_id: str, page: int = 1, per_page: int = 20, include_content: bool = False
    ) -> dict[str, Any]:
        """
        Get version history list for a segment

        Args:
            segment_id: Segment ID
            page: Page number
            per_page: Items per page
            include_content: Include full content (default False to save bandwidth)

        Returns:
            Dict with versions list and pagination info
        """
        list_key = cls._get_version_list_key(segment_id)

        # Get total count
        total = redis_client.zcard(list_key)

        if total == 0:
            return {"versions": [], "total": 0, "page": page, "per_page": per_page, "pages": 0}

        # Calculate pagination (newest first)
        start = (page - 1) * per_page
        end = start + per_page - 1

        # Get version numbers (reverse order)
        version_numbers = redis_client.zrevrange(list_key, start, end)

        # Get version details
        versions = []
        for version_number in version_numbers:
            version_key = cls._get_version_key(segment_id, int(version_number))
            version_data = redis_client.hgetall(version_key)

            if version_data:
                version_dict = {}
                for k, v in version_data.items():
                    key = k.decode() if isinstance(k, bytes) else k
                    value = v.decode() if isinstance(v, bytes) else v

                    # Try to parse JSON
                    try:
                        if value and value.startswith(("[", "{")):
                            version_dict[key] = json.loads(value)
                        else:
                            version_dict[key] = value
                    except json.JSONDecodeError:
                        version_dict[key] = value

                # Remove large fields if not needed
                if not include_content:
                    version_dict.pop("content", None)
                    version_dict.pop("answer", None)

                versions.append(version_dict)

        return {
            "versions": versions,
            "total": total,
            "page": page,
            "per_page": per_page,
            "pages": (total + per_page - 1) // per_page,
        }

    @classmethod
    def get_version_detail(cls, segment_id: str, version_number: int) -> Optional[dict[str, Any]]:
        """Get detailed info for a specific version"""
        version_key = cls._get_version_key(segment_id, version_number)
        version_data = redis_client.hgetall(version_key)

        if not version_data:
            return None

        version_dict = {}
        for k, v in version_data.items():
            key = k.decode() if isinstance(k, bytes) else k
            value = v.decode() if isinstance(v, bytes) else v

            try:
                if value and value.startswith(("[", "{")):
                    version_dict[key] = json.loads(value)
                else:
                    version_dict[key] = value
            except json.JSONDecodeError:
                version_dict[key] = value

        return version_dict

    @classmethod
    def compare_versions(
        cls, segment_id: str, version_number_1: int, version_number_2: int
    ) -> Optional[dict[str, Any]]:
        """Compare two versions"""
        v1 = cls.get_version_detail(segment_id, version_number_1)
        v2 = cls.get_version_detail(segment_id, version_number_2)

        if not v1 or not v2:
            return None

        # Calculate differences
        differences = {}
        compare_fields = ["content", "answer", "keywords", "enabled", "word_count"]

        for field in compare_fields:
            val1 = v1.get(field)
            val2 = v2.get(field)
            if val1 != val2:
                differences[field] = {"changed": True, "old_value": val1, "new_value": val2}

        return {
            "version_1": {
                "version_number": v1.get("version_number"),
                "created_at": v1.get("created_at"),
                "created_by_name": v1.get("created_by_name"),
                "change_type": v1.get("change_type"),
            },
            "version_2": {
                "version_number": v2.get("version_number"),
                "created_at": v2.get("created_at"),
                "created_by_name": v2.get("created_by_name"),
                "change_type": v2.get("change_type"),
            },
            "differences": differences,
            "has_changes": len(differences) > 0,
        }

    @classmethod
    def get_latest_version(cls, segment_id: str) -> Optional[dict[str, Any]]:
        """Get the latest version"""
        list_key = cls._get_version_list_key(segment_id)
        latest_versions = redis_client.zrevrange(list_key, 0, 0)

        if not latest_versions:
            return None

        return cls.get_version_detail(segment_id, int(latest_versions[0]))

    @classmethod
    def delete_all_versions(cls, segment_id: str):
        """Delete all versions for a segment (called when segment is deleted)"""
        list_key = cls._get_version_list_key(segment_id)

        # Get all version numbers
        version_numbers = redis_client.zrange(list_key, 0, -1)

        # Delete all version data
        for version_number in version_numbers:
            version_key = cls._get_version_key(segment_id, int(version_number))
            redis_client.delete(version_key)

        # Delete list, counter, and hash
        redis_client.delete(list_key)
        redis_client.delete(cls._get_version_counter_key(segment_id))
        redis_client.delete(cls._get_current_hash_key(segment_id))

    @classmethod
    def set_ttl_days(cls, days: int):
        """Set version retention days (global config)"""
        cls.DEFAULT_TTL_DAYS = days

    @classmethod
    def set_max_versions(cls, max_versions: int):
        """Set maximum version count (global config)"""
        cls.MAX_VERSIONS = max_versions

    @classmethod
    def get_segment_statistics(cls, segment_id: str) -> dict[str, Any]:
        """Get version statistics for a segment"""
        list_key = cls._get_version_list_key(segment_id)
        total_versions = redis_client.zcard(list_key)

        if total_versions == 0:
            return {
                "total_versions": 0,
                "first_version_at": None,
                "latest_version_at": None,
            }

        # Get first and last version
        first_version_number = redis_client.zrange(list_key, 0, 0)[0]
        latest_version_number = redis_client.zrevrange(list_key, 0, 0)[0]

        first_version = cls.get_version_detail(segment_id, int(first_version_number))
        latest_version = cls.get_version_detail(segment_id, int(latest_version_number))

        return {
            "total_versions": total_versions,
            "first_version_at": first_version.get("created_at") if first_version else None,
            "latest_version_at": latest_version.get("created_at") if latest_version else None,
            "latest_version_number": int(latest_version_number),
            "ttl_days": cls.DEFAULT_TTL_DAYS,
            "max_versions": cls.MAX_VERSIONS,
        }