realtimex-deeptutor 0.5.0.post1__py3-none-any.whl

src/agents/research/data_structures.py

@@ -0,0 +1,451 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+DR-in-KG 2.0 Core Data Structures
+Includes: TopicBlock, ToolTrace, DynamicTopicQueue
+"""
+
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from enum import Enum
+import json
+from pathlib import Path
+from typing import Any
+
+
+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_NAIVE = "rag_naive"
+    RAG_HYBRID = "rag_hybrid"
+    QUERY_ITEM = "query_item"
+    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_naive, 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
+        try:
+            data = json.loads(raw_answer)
+
+            # If it's a dict with common RAG response fields, truncate content fields
+            if isinstance(data, dict):
+                # Truncate long content fields
+                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):
+                            # Keep only first few items
+                            data[field_name] = data[field_name][:3]
+                            if data[field_name]:
+                                data[field_name].append({"note": "... additional items truncated"})
+
+                truncated = json.dumps(data, ensure_ascii=False)
+                if len(truncated) <= max_size:
+                    return truncated
+        except (json.JSONDecodeError, TypeError):
+            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"""
+        # Handle backward compatibility - old data may not have new fields
+        if "raw_answer_truncated" not in data:
+            data["raw_answer_truncated"] = False
+        if "raw_answer_original_size" not in data:
+            data["raw_answer_original_size"] = len(data.get("raw_answer", ""))
+        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 (text or "").strip().lower()
+
+    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 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",
+]

src/agents/research/main.py

@@ -0,0 +1,188 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+DR-in-KG 2.0 - Main Entry
+Deep research system based on dynamic topic queue
+
+Usage:
+  python main.py --topic "Research Topic" [--preset quick/standard/deep]
+"""
+
+import argparse
+import asyncio
+from pathlib import Path
+import sys
+
+from dotenv import load_dotenv
+import yaml
+
+from src.agents.research.research_pipeline import ResearchPipeline
+from src.services.llm import get_llm_config
+
+
+def load_config(config_path: str = None, preset: str = None) -> dict:
+    """
+    Load configuration file (with main.yaml merge)
+
+    Args:
+        config_path: Configuration file path (default: config/research_config.yaml)
+        preset: Preset mode (quick/standard/deep)
+
+    Returns:
+        Configuration dictionary (merged with main.yaml)
+    """
+    if config_path is None:
+        project_root = Path(__file__).parent.parent.parent.parent
+        from src.services.config import load_config_with_main
+
+        config = load_config_with_main("research_config.yaml", project_root)
+    else:
+        # If custom config path provided, load it directly (for backward compatibility)
+        config_file = Path(config_path)
+        if not config_file.exists():
+            raise FileNotFoundError(f"Configuration file not found: {config_file}")
+        with open(config_file, encoding="utf-8") as f:
+            config = yaml.safe_load(f) or {}
+
+    # Apply preset
+    if preset and "presets" in config and preset in config["presets"]:
+        print(f"✓ Applied preset configuration: {preset}")
+        preset_config = config["presets"][preset]
+        for key, value in preset_config.items():
+            if key in config and isinstance(value, dict):
+                config[key].update(value)
+
+    return config
+
+
+def display_config(config: dict):
+    """Display current configuration"""
+    print("\n" + "=" * 70)
+    print("📋 Current Configuration")
+    print("=" * 70)
+
+    planning = config.get("planning", {})
+    researching = config.get("researching", {})
+    reporting = config.get("reporting", {})
+
+    print("【Planning Configuration】")
+    print(f"  Initial subtopics: {planning.get('decompose', {}).get('initial_subtopics', 5)}")
+    print(f"  Max subtopics: {planning.get('decompose', {}).get('max_subtopics', 10)}")
+
+    print("\n【Researching Configuration】")
+    print(f"  Max iterations: {researching.get('max_iterations', 5)}")
+    print(f"  Research mode: {researching.get('research_mode', 'deep')}")
+    print("  Enabled tools:")
+    print(f"    - RAG: {researching.get('enable_rag_hybrid', True)}")
+    print(f"    - Web Search: {researching.get('enable_web_search', True)}")
+    print(f"    - Paper Search: {researching.get('enable_paper_search', True)}")
+
+    print("\n【Reporting Configuration】")
+    print(f"  Min section length: {reporting.get('min_section_length', 500)} characters")
+    print(f"  Enable topic deduplication: {reporting.get('enable_deduplication', True)}")
+
+    print("=" * 70 + "\n")
+
+
+async def main():
+    """Main function"""
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(
+        description="DR-in-KG 2.0 - Deep research system based on dynamic topic queue",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Quick mode
+  python main.py --topic "Deep Learning Basics" --preset quick
+
+  # Standard mode
+  python main.py --topic "Transformer Architecture" --preset standard
+
+  # Deep mode
+  python main.py --topic "Graph Neural Networks" --preset deep
+        """,
+    )
+
+    parser.add_argument("--topic", type=str, required=True, help="Research topic")
+
+    parser.add_argument(
+        "--config",
+        type=str,
+        default="config.yaml",
+        help="Configuration file path (default: config.yaml)",
+    )
+
+    parser.add_argument(
+        "--preset",
+        type=str,
+        choices=["quick", "standard", "deep"],
+        help="Preset configuration (quick: fast, standard: standard, deep: deep)",
+    )
+
+    parser.add_argument("--output-dir", type=str, help="Output directory (overrides config file)")
+
+    args = parser.parse_args()
+
+    # Load environment variables
+    load_dotenv()
+
+    # Check API configuration
+    try:
+        llm_config = get_llm_config()
+    except ValueError as e:
+        print(f"✗ Error: {e}")
+        print("Please configure in .env or DeepTutor.env file:")
+        print("  LLM_MODEL=gpt-4o")
+        print("  LLM_API_KEY=your_api_key_here")
+        print("  LLM_HOST=https://api.openai.com/v1")
+        sys.exit(1)
+
+    # Load configuration
+    try:
+        config = load_config(args.config, args.preset)
+    except Exception as e:
+        print(f"✗ Failed to load configuration: {e!s}")
+        sys.exit(1)
+
+    # Override configuration (command line arguments take priority)
+    if args.output_dir:
+        config["system"]["output_base_dir"] = args.output_dir
+        config["system"]["reports_dir"] = args.output_dir
+
+    # Display configuration
+    display_config(config)
+
+    # Create research pipeline
+    pipeline = ResearchPipeline(
+        config=config, api_key=llm_config.api_key, base_url=llm_config.base_url
+    )
+
+    # Execute research
+    try:
+        result = await pipeline.run(topic=args.topic)
+
+        print("\n" + "=" * 70)
+        print("✓ Research completed!")
+        print("=" * 70)
+        print(f"Research ID: {result['research_id']}")
+        print(f"Topic: {result['topic']}")
+        print(f"Final Report: {result['final_report_path']}")
+        print("=" * 70 + "\n")
+
+    except KeyboardInterrupt:
+        print("\n\n⚠️  Research interrupted by user")
+        sys.exit(0)
+    except Exception as e:
+        print(f"\n\n✗ Research failed: {e!s}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    # Windows compatibility
+    if sys.platform == "win32":
+        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+    asyncio.run(main())