hkuds--deeptutor/deeptutor/agents/research/data_structures.py

#!/usr/bin/env python
"""
DR-in-KG 2.0 Core Data Structures
Includes: TopicBlock, ToolTrace, DynamicTopicQueue
"""

from dataclasses import asdict, dataclass, field
from datetime import datetime
import difflib
from enum import Enum
import json
from pathlib import Path
import re
from typing import Any

from deeptutor.utils.json_parser import parse_json_response

# Default fuzzy-match threshold used by :meth:`DynamicTopicQueue.find_similar`.
# 0.85 reliably catches near-duplicate titles (case / punctuation / one-word
# reorderings) while letting genuinely distinct sub-topics through.
DEFAULT_TOPIC_SIMILARITY_THRESHOLD = 0.85
_TOPIC_TOKEN_RE = re.compile(r"[^\W_]+", re.UNICODE)
_TOPIC_STOPWORDS = {
    "a",
    "an",
    "and",
    "as",
    "at",
    "by",
    "for",
    "from",
    "in",
    "into",
    "of",
    "on",
    "or",
    "the",
    "to",
    "vs",
    "with",
}


class TopicStatus(Enum):
    """Topic block status enumeration"""

    PENDING = "pending"  # Pending research
    RESEARCHING = "researching"  # Researching
    COMPLETED = "completed"  # Completed
    FAILED = "failed"  # Failed


class ToolType(Enum):
    """Tool type enumeration"""

    RAG = "rag"
    PAPER_SEARCH = "paper_search"
    RUN_CODE = "run_code"
    WEB_SEARCH = "web_search"


# Default max size for raw_answer (50KB)
DEFAULT_RAW_ANSWER_MAX_SIZE = 50 * 1024


@dataclass
class ToolTrace:
    """
    Tool trace - Records complete loop of a single tool call
    """

    tool_id: str  # Unique identifier (e.g., "tool_1", "tool_2")
    citation_id: str  # Citation ID (for report citations and anchors, e.g., CIT-1-01)
    tool_type: str  # Tool type (rag, web_search, etc.)
    query: str  # Query statement issued
    raw_answer: str  # Raw detailed result returned by tool (may be truncated)
    summary: str  # Core summary generated by Note Agent
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
    raw_answer_truncated: bool = field(default=False)  # Whether raw_answer was truncated
    raw_answer_original_size: int = field(default=0)  # Original size before truncation

    def __post_init__(self):
        """Post-initialization to handle raw_answer size limit"""
        if self.raw_answer_original_size == 0:
            self.raw_answer_original_size = len(self.raw_answer)

        # Truncate if needed
        if len(self.raw_answer) > DEFAULT_RAW_ANSWER_MAX_SIZE:
            self.raw_answer = self._truncate_raw_answer(
                self.raw_answer, DEFAULT_RAW_ANSWER_MAX_SIZE
            )
            self.raw_answer_truncated = True

    @staticmethod
    def _truncate_raw_answer(raw_answer: str, max_size: int) -> str:
        """
        Truncate raw_answer while trying to preserve valid JSON structure

        Args:
            raw_answer: Original raw answer string
            max_size: Maximum size in bytes

        Returns:
            Truncated string
        """
        if len(raw_answer) <= max_size:
            return raw_answer

        # Try to parse as JSON and truncate intelligently
        data = parse_json_response(raw_answer, fallback=None)
        if isinstance(data, dict):
            content_fields = ["answer", "content", "text", "chunks", "documents"]
            for field_name in content_fields:
                if field_name in data:
                    if isinstance(data[field_name], str) and len(data[field_name]) > max_size // 2:
                        data[field_name] = data[field_name][: max_size // 2] + "... [truncated]"
                    elif isinstance(data[field_name], list):
                        data[field_name] = data[field_name][:3]
                        if data[field_name]:
                            data[field_name].append({"note": "... additional items truncated"})

            try:
                truncated = json.dumps(data, ensure_ascii=False)
                if len(truncated) <= max_size:
                    return truncated
            except (TypeError, ValueError):
                pass

        # Fallback: simple truncation with marker
        truncation_marker = "\n... [content truncated, original size: {} bytes]".format(
            len(raw_answer)
        )
        return raw_answer[: max_size - len(truncation_marker)] + truncation_marker

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary"""
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ToolTrace":
        """Create from dictionary"""
        return cls(**data)

    @classmethod
    def create_with_size_limit(
        cls,
        tool_id: str,
        citation_id: str,
        tool_type: str,
        query: str,
        raw_answer: str,
        summary: str,
        max_size: int = DEFAULT_RAW_ANSWER_MAX_SIZE,
    ) -> "ToolTrace":
        """
        Create a ToolTrace with explicit size limit

        Args:
            tool_id: Tool ID
            citation_id: Citation ID
            tool_type: Tool type
            query: Query string
            raw_answer: Raw answer (will be truncated if needed)
            summary: Summary
            max_size: Maximum size for raw_answer

        Returns:
            ToolTrace instance
        """
        original_size = len(raw_answer)
        truncated = len(raw_answer) > max_size

        if truncated:
            raw_answer = cls._truncate_raw_answer(raw_answer, max_size)

        return cls(
            tool_id=tool_id,
            citation_id=citation_id,
            tool_type=tool_type,
            query=query,
            raw_answer=raw_answer,
            summary=summary,
            raw_answer_truncated=truncated,
            raw_answer_original_size=original_size,
        )


@dataclass
class TopicBlock:
    """
    Topic block - Minimum scheduling unit in queue
    """

    block_id: str  # Unique identifier (e.g., "block_1", "block_2")
    sub_topic: str  # Sub-topic name
    overview: str  # Topic overview/background
    status: TopicStatus = TopicStatus.PENDING  # Topic status
    tool_traces: list[ToolTrace] = field(default_factory=list)  # Tool call trace list
    iteration_count: int = 0  # Current iteration count
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
    updated_at: str = field(default_factory=lambda: datetime.now().isoformat())
    metadata: dict[str, Any] = field(default_factory=dict)  # Additional metadata

    def add_tool_trace(self, trace: ToolTrace) -> None:
        """Add tool trace"""
        self.tool_traces.append(trace)
        self.updated_at = datetime.now().isoformat()

    def get_latest_trace(self) -> ToolTrace | None:
        """Get latest tool trace"""
        return self.tool_traces[-1] if self.tool_traces else None

    def get_all_summaries(self) -> str:
        """Get concatenated summaries of all tool traces"""
        if not self.tool_traces:
            return ""
        return "\n".join([f"[{trace.tool_type}] {trace.summary}" for trace in self.tool_traces])

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary"""
        data = asdict(self)
        data["status"] = self.status.value
        data["tool_traces"] = [trace.to_dict() for trace in self.tool_traces]
        return data

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "TopicBlock":
        """Create from dictionary"""
        data_copy = data.copy()
        if isinstance(data_copy.get("status"), str):
            data_copy["status"] = TopicStatus(data_copy["status"])
        if "tool_traces" in data_copy:
            data_copy["tool_traces"] = [
                ToolTrace.from_dict(t) if isinstance(t, dict) else t
                for t in data_copy["tool_traces"]
            ]
        return cls(**data_copy)


class DynamicTopicQueue:
    """
    Dynamic topic queue - Core memory and scheduling center of the system
    """

    def __init__(
        self, research_id: str, max_length: int | None = None, state_file: str | None = None
    ):
        """
        Initialize queue

        Args:
            research_id: Research task ID
            max_length: Maximum queue length (None means unlimited)
            state_file: Auto-persistence file path
        """
        self.research_id = research_id
        self.blocks: list[TopicBlock] = []
        self.block_counter = 0
        self.created_at = datetime.now().isoformat()
        self.max_length = max_length if isinstance(max_length, int) and max_length > 0 else None
        self.state_file = state_file

    def set_state_file(self, filepath: str | None) -> None:
        """Set queue auto-persistence file"""
        self.state_file = filepath
        self._auto_save()

    @staticmethod
    def _normalize_topic(text: str) -> str:
        return re.sub(r"\s+", " ", (text or "").strip().lower())

    @classmethod
    def _topic_tokens(cls, text: str) -> set[str]:
        tokens: set[str] = set()
        for raw in _TOPIC_TOKEN_RE.findall(cls._normalize_topic(text)):
            token = raw.strip()
            if not token or token in _TOPIC_STOPWORDS:
                continue
            # Tiny English stemmer: enough to align "basics" and "basic"
            # without adding a heavyweight NLP dependency.
            if len(token) > 4 and token.endswith("ies"):
                token = token[:-3] + "y"
            elif len(token) > 3 and token.endswith("s"):
                token = token[:-1]
            tokens.add(token)
        return tokens

    @classmethod
    def _topic_similarity(cls, left: str, right: str) -> float:
        left_norm = cls._normalize_topic(left)
        right_norm = cls._normalize_topic(right)
        if not left_norm or not right_norm:
            return 0.0
        if left_norm == right_norm:
            return 1.0

        sequence_score = difflib.SequenceMatcher(None, left_norm, right_norm).ratio()
        left_tokens = cls._topic_tokens(left_norm)
        right_tokens = cls._topic_tokens(right_norm)
        if not left_tokens or not right_tokens:
            return sequence_score

        overlap = left_tokens & right_tokens
        jaccard = len(overlap) / max(1, len(left_tokens | right_tokens))
        containment = len(overlap) / max(1, min(len(left_tokens), len(right_tokens)))
        token_score = jaccard
        if len(left_tokens) >= 2 and len(right_tokens) >= 2 and jaccard >= 0.5:
            token_score = max(token_score, containment * 0.95)
        return max(sequence_score, token_score)

    def add_block(self, sub_topic: str, overview: str) -> TopicBlock:
        """
        Add new topic block to the end of queue

        Args:
            sub_topic: Sub-topic name
            overview: Topic overview

        Returns:
            Created TopicBlock
        """
        if self.max_length and len(self.blocks) >= self.max_length:
            raise RuntimeError(
                f"Queue has reached maximum capacity ({self.max_length}), cannot add new topic."
            )
        self.block_counter += 1
        block_id = f"block_{self.block_counter}"
        block = TopicBlock(block_id=block_id, sub_topic=sub_topic, overview=overview)
        self.blocks.append(block)
        self._auto_save()
        return block

    def has_topic(self, sub_topic: str) -> bool:
        """Check if topic already exists (case-insensitive, ignoring leading/trailing spaces)"""
        target = self._normalize_topic(sub_topic)
        if not target:
            return False
        return any(self._normalize_topic(b.sub_topic) == target for b in self.blocks)

    def is_full(self) -> bool:
        """Return ``True`` when the queue has reached its configured cap."""
        return self.max_length is not None and len(self.blocks) >= self.max_length

    def find_similar(
        self,
        sub_topic: str,
        *,
        threshold: float = DEFAULT_TOPIC_SIMILARITY_THRESHOLD,
    ) -> TopicBlock | None:
        """Return an existing block whose title is fuzzily similar to
        ``sub_topic``, or ``None`` when no match exceeds ``threshold``.

        Used to dedup ``APPEND`` requests so the LLM can't reliably keep
        proposing the same topic in slightly different words. Exact
        normalised matches always win; otherwise the highest-scoring
        block above ``threshold`` is returned.
        """
        target = self._normalize_topic(sub_topic)
        if not target:
            return None

        best: tuple[float, TopicBlock] | None = None
        for block in self.blocks:
            candidate = self._normalize_topic(block.sub_topic)
            if not candidate:
                continue
            if candidate == target:
                return block
            score = self._topic_similarity(target, candidate)
            if score >= threshold and (best is None or score > best[0]):
                best = (score, block)
        return best[1] if best else None

    def append_child(
        self,
        *,
        parent: TopicBlock | None,
        sub_topic: str,
        overview: str = "",
    ) -> TopicBlock | None:
        """Append a new block to the queue tail, optionally tagging the
        parent block's id in metadata so reporting can reconstruct the
        topic tree.

        Returns the new block on success, or ``None`` when the queue is
        already full. Duplicate detection is the caller's responsibility
        (use :meth:`find_similar` first when needed).
        """
        if self.is_full():
            return None
        self.block_counter += 1
        block_id = f"block_{self.block_counter}"
        metadata: dict[str, Any] = {}
        if parent is not None:
            metadata["parent_block_id"] = parent.block_id
        block = TopicBlock(
            block_id=block_id,
            sub_topic=sub_topic,
            overview=overview,
            metadata=metadata,
        )
        self.blocks.append(block)
        self._auto_save()
        return block

    def list_topics(self) -> list[str]:
        """List all current topic titles"""
        return [b.sub_topic for b in self.blocks]

    def get_pending_block(self) -> TopicBlock | None:
        """
        Get first pending topic block

        Returns:
            First TopicBlock with PENDING status, or None if not found
        """
        for block in self.blocks:
            if block.status == TopicStatus.PENDING:
                return block
        return None

    def get_block_by_id(self, block_id: str) -> TopicBlock | None:
        """
        Get topic block by ID

        Args:
            block_id: Topic block ID

        Returns:
            Corresponding TopicBlock, or None if not found
        """
        for block in self.blocks:
            if block.block_id == block_id:
                return block
        return None

    def mark_researching(self, block_id: str) -> bool:
        """
        Mark topic block as researching

        Args:
            block_id: Topic block ID

        Returns:
            Whether marking was successful
        """
        block = self.get_block_by_id(block_id)
        if block:
            block.status = TopicStatus.RESEARCHING
            block.updated_at = datetime.now().isoformat()
            self._auto_save()
            return True
        return False

    def mark_completed(self, block_id: str) -> bool:
        """
        Mark topic block as completed

        Args:
            block_id: Topic block ID

        Returns:
            Whether marking was successful
        """
        block = self.get_block_by_id(block_id)
        if block:
            block.status = TopicStatus.COMPLETED
            block.updated_at = datetime.now().isoformat()
            self._auto_save()
            return True
        return False

    def mark_failed(self, block_id: str) -> bool:
        """
        Mark topic block as failed

        Args:
            block_id: Topic block ID

        Returns:
            Whether marking was successful
        """
        block = self.get_block_by_id(block_id)
        if block:
            block.status = TopicStatus.FAILED
            block.updated_at = datetime.now().isoformat()
            self._auto_save()
            return True
        return False

    def get_all_completed_blocks(self) -> list[TopicBlock]:
        """Get all completed topic blocks"""
        return [b for b in self.blocks if b.status == TopicStatus.COMPLETED]

    def get_all_pending_blocks(self) -> list[TopicBlock]:
        """Get all pending topic blocks"""
        return [b for b in self.blocks if b.status == TopicStatus.PENDING]

    def is_all_completed(self) -> bool:
        """Check if all topic blocks are completed"""
        if not self.blocks:
            return False
        return all(b.status == TopicStatus.COMPLETED for b in self.blocks)

    def get_statistics(self) -> dict[str, Any]:
        """Get queue statistics"""
        return {
            "total_blocks": len(self.blocks),
            "pending": len(self.get_all_pending_blocks()),
            "researching": len([b for b in self.blocks if b.status == TopicStatus.RESEARCHING]),
            "completed": len(self.get_all_completed_blocks()),
            "failed": len([b for b in self.blocks if b.status == TopicStatus.FAILED]),
            "total_tool_calls": sum(len(b.tool_traces) for b in self.blocks),
        }

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary"""
        return {
            "research_id": self.research_id,
            "created_at": self.created_at,
            "blocks": [b.to_dict() for b in self.blocks],
            "statistics": self.get_statistics(),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "DynamicTopicQueue":
        """Create from dictionary"""
        queue = cls(data["research_id"])
        queue.created_at = data.get("created_at", queue.created_at)
        for block_data in data.get("blocks", []):
            block = TopicBlock.from_dict(block_data)
            queue.blocks.append(block)
            # Update counter
            if block.block_id.startswith("block_"):
                try:
                    block_num = int(block.block_id.split("_")[1])
                    queue.block_counter = max(queue.block_counter, block_num)
                except (ValueError, IndexError):
                    pass
        return queue

    def save_to_json(self, filepath: str) -> None:
        """Save queue to JSON file"""
        Path(filepath).parent.mkdir(parents=True, exist_ok=True)
        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, ensure_ascii=False, indent=2)

    def _auto_save(self) -> None:
        """Auto-save if state_file is set"""
        if self.state_file:
            try:
                self.save_to_json(self.state_file)
            except Exception as exc:
                print(f"⚠️ Failed to save queue progress: {exc}")

    @classmethod
    def load_from_json(cls, filepath: str) -> "DynamicTopicQueue":
        """Load queue from JSON file"""
        with open(filepath, encoding="utf-8") as f:
            data = json.load(f)
        return cls.from_dict(data)


__all__ = [
    "DynamicTopicQueue",
    "ToolTrace",
    "ToolType",
    "TopicBlock",
    "TopicStatus",
]