#!/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", ]