convoviz 0.1.6__py3-none-any.whl → 0.2.0__py3-none-any.whl

convoviz/interactive.py

@@ -0,0 +1,178 @@
+"""Interactive configuration prompts using questionary."""
+
+from pathlib import Path
+
+from questionary import Choice, Style, checkbox, select
+from questionary import path as qst_path
+from questionary import text as qst_text
+
+from convoviz.config import ConvovizConfig, get_default_config
+from convoviz.io.loaders import find_latest_zip, validate_zip
+from convoviz.utils import colormaps, default_font_path, font_names, font_path, validate_header
+
+CUSTOM_STYLE = Style(
+    [
+        ("qmark", "fg:#34eb9b bold"),
+        ("question", "bold fg:#e0e0e0"),
+        ("answer", "fg:#34ebeb bold"),
+        ("pointer", "fg:#e834eb bold"),
+        ("highlighted", "fg:#349ceb bold"),
+        ("selected", "fg:#34ebeb"),
+        ("separator", "fg:#eb3434"),
+        ("instruction", "fg:#eb9434"),
+        ("text", "fg:#b2eb34"),
+        ("disabled", "fg:#858585 italic"),
+    ]
+)
+
+
+def run_interactive_config(initial_config: ConvovizConfig | None = None) -> ConvovizConfig:
+    """Run interactive prompts to configure convoviz.
+
+    Args:
+        initial_config: Optional starting configuration (uses defaults if None)
+
+    Returns:
+        Updated configuration based on user input
+    """
+    config = initial_config or get_default_config()
+
+    # Set sensible defaults if not already set
+    if not config.zip_filepath:
+        latest = find_latest_zip()
+        if latest:
+            config.zip_filepath = latest
+
+    if not config.wordcloud.font_path:
+        config.wordcloud.font_path = default_font_path()
+
+    # Prompt for zip file path
+    zip_default = str(config.zip_filepath) if config.zip_filepath else ""
+    zip_result = qst_path(
+        "Enter the path to the zip file:",
+        default=zip_default,
+        validate=lambda p: validate_zip(Path(p))
+        or "Invalid zip file (must contain conversations.json)",
+        style=CUSTOM_STYLE,
+    ).ask()
+
+    if zip_result:
+        config.zip_filepath = Path(zip_result)
+
+    # Prompt for output folder
+    output_result = qst_path(
+        "Enter the path to the output folder:",
+        default=str(config.output_folder),
+        style=CUSTOM_STYLE,
+    ).ask()
+
+    if output_result:
+        config.output_folder = Path(output_result)
+
+    # Prompt for author headers
+    headers = config.message.author_headers
+    for role in ["system", "user", "assistant", "tool"]:
+        current = getattr(headers, role)
+        result = qst_text(
+            f"Enter the message header for '{role}':",
+            default=current,
+            validate=lambda t: validate_header(t)
+            or "Must be a valid markdown header (e.g., # Title)",
+            style=CUSTOM_STYLE,
+        ).ask()
+        if result:
+            setattr(headers, role, result)
+
+    # Prompt for LaTeX delimiters
+    latex_result = select(
+        "Select the LaTeX math delimiters:",
+        choices=["default", "dollars"],
+        default=config.conversation.markdown.latex_delimiters,
+        style=CUSTOM_STYLE,
+    ).ask()
+
+    if latex_result:
+        config.conversation.markdown.latex_delimiters = latex_result
+
+    # Prompt for YAML headers
+    yaml_config = config.conversation.yaml
+    yaml_choices = [
+        Choice(title=field, checked=getattr(yaml_config, field))
+        for field in [
+            "title",
+            "tags",
+            "chat_link",
+            "create_time",
+            "update_time",
+            "model",
+            "used_plugins",
+            "message_count",
+            "content_types",
+            "custom_instructions",
+        ]
+    ]
+
+    selected = checkbox(
+        "Select YAML metadata headers to include:",
+        choices=yaml_choices,
+        style=CUSTOM_STYLE,
+    ).ask()
+
+    if selected is not None:
+        selected_set = set(selected)
+        for field_name in [
+            "title",
+            "tags",
+            "chat_link",
+            "create_time",
+            "update_time",
+            "model",
+            "used_plugins",
+            "message_count",
+            "content_types",
+            "custom_instructions",
+        ]:
+            setattr(yaml_config, field_name, field_name in selected_set)
+
+    # Prompt for font
+    available_fonts = font_names()
+    if available_fonts:
+        current_font = (
+            config.wordcloud.font_path.stem if config.wordcloud.font_path else available_fonts[0]
+        )
+        font_result = select(
+            "Select the font for word clouds:",
+            choices=available_fonts,
+            default=current_font if current_font in available_fonts else available_fonts[0],
+            style=CUSTOM_STYLE,
+        ).ask()
+
+        if font_result:
+            config.wordcloud.font_path = font_path(font_result)
+
+    # Prompt for colormap
+    available_colormaps = colormaps()
+    if available_colormaps:
+        colormap_result = select(
+            "Select the color theme for word clouds:",
+            choices=available_colormaps,
+            default=config.wordcloud.colormap
+            if config.wordcloud.colormap in available_colormaps
+            else available_colormaps[0],
+            style=CUSTOM_STYLE,
+        ).ask()
+
+        if colormap_result:
+            config.wordcloud.colormap = colormap_result
+
+    # Prompt for custom stopwords
+    stopwords_result = qst_text(
+        "Enter custom stopwords (comma-separated):",
+        default=config.wordcloud.custom_stopwords,
+        style=CUSTOM_STYLE,
+    ).ask()
+
+    if stopwords_result is not None:
+        config.wordcloud.custom_stopwords = stopwords_result
+
+    return config

convoviz/io/__init__.py

@@ -0,0 +1,21 @@
+"""I/O operations for convoviz."""
+
+from convoviz.io.loaders import (
+    load_collection_from_json,
+    load_collection_from_zip,
+    load_conversation_from_json,
+)
+from convoviz.io.writers import (
+    save_collection,
+    save_conversation,
+    save_custom_instructions,
+)
+
+__all__ = [
+    "load_collection_from_json",
+    "load_collection_from_zip",
+    "load_conversation_from_json",
+    "save_collection",
+    "save_conversation",
+    "save_custom_instructions",
+]

convoviz/io/loaders.py

@@ -0,0 +1,135 @@
+"""Loading functions for conversations and collections."""
+
+from pathlib import Path
+from zipfile import ZipFile
+
+from orjson import loads
+
+from convoviz.exceptions import InvalidZipError
+from convoviz.models import Conversation, ConversationCollection
+
+
+def extract_archive(filepath: Path) -> Path:
+    """Extract a ZIP file and return the extraction folder path.
+
+    Args:
+        filepath: Path to the ZIP file
+
+    Returns:
+        Path to the extracted folder
+    """
+    folder = filepath.with_suffix("")
+    with ZipFile(filepath) as zf:
+        zf.extractall(folder)
+    return folder
+
+
+def validate_zip(filepath: Path) -> bool:
+    """Check if a ZIP file contains conversations.json.
+
+    Args:
+        filepath: Path to the ZIP file
+
+    Returns:
+        True if valid, False otherwise
+    """
+    if not filepath.is_file() or filepath.suffix != ".zip":
+        return False
+    try:
+        with ZipFile(filepath) as zf:
+            return "conversations.json" in zf.namelist()
+    except Exception:
+        return False
+
+
+def load_conversation_from_json(filepath: Path | str) -> Conversation:
+    """Load a single conversation from a JSON file.
+
+    Args:
+        filepath: Path to the JSON file
+
+    Returns:
+        Loaded Conversation object
+    """
+    filepath = Path(filepath)
+    with filepath.open(encoding="utf-8") as f:
+        data = loads(f.read())
+    return Conversation(**data)
+
+
+def load_collection_from_json(filepath: Path | str) -> ConversationCollection:
+    """Load a conversation collection from a JSON file.
+
+    The JSON file should contain an array of conversation objects.
+
+    Args:
+        filepath: Path to the JSON file
+
+    Returns:
+        Loaded ConversationCollection object
+    """
+    filepath = Path(filepath)
+    with filepath.open(encoding="utf-8") as f:
+        data = loads(f.read())
+    return ConversationCollection(conversations=data)
+
+
+def load_collection_from_zip(filepath: Path | str) -> ConversationCollection:
+    """Load a conversation collection from a ChatGPT export ZIP file.
+
+    Args:
+        filepath: Path to the ZIP file
+
+    Returns:
+        Loaded ConversationCollection object
+
+    Raises:
+        InvalidZipError: If the ZIP file is invalid or missing conversations.json
+    """
+    filepath = Path(filepath)
+
+    if not validate_zip(filepath):
+        raise InvalidZipError(str(filepath))
+
+    extracted_folder = extract_archive(filepath)
+    conversations_path = extracted_folder / "conversations.json"
+
+    return load_collection_from_json(conversations_path)
+
+
+def find_latest_zip(directory: Path | None = None) -> Path | None:
+    """Find the most recently created ZIP file in a directory.
+
+    Args:
+        directory: Directory to search (defaults to ~/Downloads)
+
+    Returns:
+        Path to the most recent ZIP, or None if none found
+    """
+    if directory is None:
+        directory = Path.home() / "Downloads"
+
+    zip_files = list(directory.glob("*.zip"))
+    if not zip_files:
+        return None
+
+    return max(zip_files, key=lambda p: p.stat().st_ctime)
+
+
+def find_latest_bookmarklet_json(directory: Path | None = None) -> Path | None:
+    """Find the most recent bookmarklet JSON file in a directory.
+
+    Args:
+        directory: Directory to search (defaults to ~/Downloads)
+
+    Returns:
+        Path to the most recent bookmarklet JSON, or None if none found
+    """
+    if directory is None:
+        directory = Path.home() / "Downloads"
+
+    bookmarklet_files = [f for f in directory.glob("*.json") if "bookmarklet" in f.name.lower()]
+    if not bookmarklet_files:
+        return None
+
+    return max(bookmarklet_files, key=lambda p: p.stat().st_ctime)

convoviz/io/writers.py

@@ -0,0 +1,96 @@
+"""Writing functions for conversations and collections."""
+
+from os import utime as os_utime
+from pathlib import Path
+
+from orjson import OPT_INDENT_2, dumps
+from tqdm import tqdm
+
+from convoviz.config import AuthorHeaders, ConversationConfig
+from convoviz.models import Conversation, ConversationCollection
+from convoviz.renderers import render_conversation
+from convoviz.utils import sanitize
+
+
+def save_conversation(
+    conversation: Conversation,
+    filepath: Path,
+    config: ConversationConfig,
+    headers: AuthorHeaders,
+) -> Path:
+    """Save a conversation to a markdown file.
+
+    Handles filename conflicts by appending a counter. Sets the file's
+    modification time to match the conversation's update time.
+
+    Args:
+        conversation: The conversation to save
+        filepath: Target file path
+        config: Conversation rendering configuration
+        headers: Author header configuration
+
+    Returns:
+        The actual path the file was saved to (may differ if there was a conflict)
+    """
+    # Handle filename conflicts
+    base_name = sanitize(filepath.stem)
+    final_path = filepath
+    counter = 0
+
+    while final_path.exists():
+        counter += 1
+        final_path = filepath.with_name(f"{base_name} ({counter}){filepath.suffix}")
+
+    # Render and write
+    markdown = render_conversation(conversation, config, headers)
+    with final_path.open("w", encoding="utf-8") as f:
+        f.write(markdown)
+
+    # Set modification time
+    timestamp = conversation.update_time.timestamp()
+    os_utime(final_path, (timestamp, timestamp))
+
+    return final_path
+
+
+def save_collection(
+    collection: ConversationCollection,
+    directory: Path,
+    config: ConversationConfig,
+    headers: AuthorHeaders,
+    *,
+    progress_bar: bool = False,
+) -> None:
+    """Save all conversations in a collection to markdown files.
+
+    Args:
+        collection: The collection to save
+        directory: Target directory
+        config: Conversation rendering configuration
+        headers: Author header configuration
+        progress_bar: Whether to show a progress bar
+    """
+    directory.mkdir(parents=True, exist_ok=True)
+
+    for conv in tqdm(
+        collection.conversations,
+        desc="Writing Markdown 📄 files",
+        disable=not progress_bar,
+    ):
+        filepath = directory / f"{sanitize(conv.title)}.md"
+        save_conversation(conv, filepath, config, headers)
+
+
+def save_custom_instructions(
+    collection: ConversationCollection,
+    filepath: Path,
+) -> None:
+    """Save all custom instructions from a collection to a JSON file.
+
+    Args:
+        collection: The collection to extract instructions from
+        filepath: Target JSON file path
+    """
+    instructions = collection.custom_instructions
+    with filepath.open("w", encoding="utf-8") as f:
+        f.write(dumps(instructions, option=OPT_INDENT_2).decode())

convoviz/models/__init__.py

@@ -1,8 +1,28 @@
-"""Models for the application."""
+"""Data models for convoviz."""
 
-from ._conversation import Conversation
-from ._conversation_set import ConversationSet
-from ._message import Message
-from ._node import Node
+from convoviz.models.collection import ConversationCollection
+from convoviz.models.conversation import Conversation
+from convoviz.models.message import (
+    AuthorRole,
+    Message,
+    MessageAuthor,
+    MessageContent,
+    MessageMetadata,
+)
+from convoviz.models.node import Node, build_node_tree
 
-__all__ = ["Conversation", "ConversationSet", "Message", "Node"]
+# Backward compatibility alias
+ConversationSet = ConversationCollection
+
+__all__ = [
+    "AuthorRole",
+    "Conversation",
+    "ConversationCollection",
+    "ConversationSet",
+    "Message",
+    "MessageAuthor",
+    "MessageContent",
+    "MessageMetadata",
+    "Node",
+    "build_node_tree",
+]

convoviz/models/collection.py

@@ -0,0 +1,107 @@
+"""ConversationCollection model - manages a set of conversations.
+
+This is a pure data model - I/O and visualization logic are in separate modules.
+"""
+
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from convoviz.models.conversation import Conversation
+from convoviz.models.message import AuthorRole
+
+
+class ConversationCollection(BaseModel):
+    """A collection of ChatGPT conversations.
+
+    Provides grouping and aggregation operations over conversations.
+    """
+
+    conversations: list[Conversation] = Field(default_factory=list)
+
+    @property
+    def index(self) -> dict[str, Conversation]:
+        """Get conversations indexed by conversation_id."""
+        return {conv.conversation_id: conv for conv in self.conversations}
+
+    @property
+    def last_updated(self) -> datetime:
+        """Get the most recent update time across all conversations."""
+        if not self.conversations:
+            return datetime.min
+        return max(conv.update_time for conv in self.conversations)
+
+    def update(self, other: "ConversationCollection") -> None:
+        """Merge another collection into this one.
+
+        Only updates if the other collection has newer content.
+        """
+        if other.last_updated <= self.last_updated:
+            return
+
+        merged_index = self.index
+        merged_index.update(other.index)
+        self.conversations = list(merged_index.values())
+
+    def add(self, conversation: Conversation) -> None:
+        """Add a conversation to the collection."""
+        self.conversations.append(conversation)
+
+    @property
+    def custom_instructions(self) -> list[dict[str, Any]]:
+        """Get all custom instructions from all conversations."""
+        instructions: list[dict[str, Any]] = []
+        for conv in self.conversations:
+            if not conv.custom_instructions:
+                continue
+            instructions.append(
+                {
+                    "chat_title": conv.title,
+                    "chat_link": conv.url,
+                    "time": conv.create_time,
+                    "custom_instructions": conv.custom_instructions,
+                }
+            )
+        return instructions
+
+    def timestamps(self, *authors: AuthorRole) -> list[float]:
+        """Get all message timestamps from specified authors."""
+        result: list[float] = []
+        for conv in self.conversations:
+            result.extend(conv.timestamps(*authors))
+        return result
+
+    def plaintext(self, *authors: AuthorRole) -> str:
+        """Get concatenated plain text from all conversations."""
+        return "\n".join(conv.plaintext(*authors) for conv in self.conversations)
+
+    def group_by_week(self) -> dict[datetime, "ConversationCollection"]:
+        """Group conversations by the week they were created."""
+        groups: dict[datetime, ConversationCollection] = {}
+        for conv in self.conversations:
+            week_start = conv.week_start
+            if week_start not in groups:
+                groups[week_start] = ConversationCollection()
+            groups[week_start].add(conv)
+        return groups
+
+    def group_by_month(self) -> dict[datetime, "ConversationCollection"]:
+        """Group conversations by the month they were created."""
+        groups: dict[datetime, ConversationCollection] = {}
+        for conv in self.conversations:
+            month_start = conv.month_start
+            if month_start not in groups:
+                groups[month_start] = ConversationCollection()
+            groups[month_start].add(conv)
+        return groups
+
+    def group_by_year(self) -> dict[datetime, "ConversationCollection"]:
+        """Group conversations by the year they were created."""
+        groups: dict[datetime, ConversationCollection] = {}
+        for conv in self.conversations:
+            year_start = conv.year_start
+            if year_start not in groups:
+                groups[year_start] = ConversationCollection()
+            groups[year_start].add(conv)
+        return groups