Merge pull request somogyijanos#5 from somogyijanos/pr-noemptyimgs

Pr noemptyimgs

Author: somogyijanos

.gitignore

@@ -1,3 +1,4 @@
 .venv
 **/__pycache__
-output
+output
+**/.DS_Store

README.md

@@ -8,8 +8,8 @@ Also see [this](https://forum.cursor.com/t/guide-5-steps-exporting-chats-prompts
 
 ## Features
 
-- **Discover Chats**: Discover all `state.vscdb` files in a directory and print a few lines of dialogue so one can identify which is the workspace (chat) one is searching for. It's also possible to filter by text.
-- **Export Chats**: Export chats data for a certain workspace to Markdown files or print it to the command line.
+- **Discover Chats**: Discover all chats from all workspaces and print a few lines of dialogue so one can identify which is the workspace (or chat) one is searching for. It's also possible to filter by text.
+- **Export Chats**: Export chats for the most recent (or a specific) workspace to Markdown files or print it to the command line.
 
 ## Installation
 
@@ -26,20 +26,44 @@ Also see [this](https://forum.cursor.com/t/guide-5-steps-exporting-chats-prompts
 
 ## Usage
 
-Find, where the `state.vscdb` is located in your computer. The table below may help:
+First, find, where the `state.vscdb` files are located on your computer. Confirm that corresponding to your system, the right path is set in the [config.yml](./config.yml) file. Update it if not set correctly.
 
-| OS               | Path                                                      |
-|------------------|-----------------------------------------------------------|
-| Windows          | `%APPDATA%\Cursor\User\workspaceStorage`                  |
-| macOS            | `/Users/YOUR_USERNAME/Library/Application Support/Cursor/User/workspaceStorage` |
-| Linux            | `/home/YOUR_USERNAME/.config/Cursor/User/workspaceStorage` |
+Both the `discover` and `export` commands will work with this path by default, but you can also provide a custom path any time.
 
-### Discover Chats of all Workspaces
+---
+
+### Discover Chats
 ```sh
-./chat.py discover --search-text "matplotlib" "/Users/myuser/Library/Application Support/Cursor/User/workspaceStorage"
+# Help on usage
+./chat.py discover --help
+
+# Discover all chats from all workspaces
+./chat.py discover
+
+# Apply text filter
+./chat.py discover --search-text "matplotlib"
+
+# Discover all chats from all workspaces at a custom path
+./chat.py discover "/path/to/workspaces"
 ```
 
-### Export Chats of a Workspace
+---
+
+### Export Chats
+See `./chat.py export --help` for general help. Examples:
 ```sh
-./chat.py export --output-dir output "/Users/myuser/Library/Application Support/Cursor/User/workspaceStorage/b989572f2e2186b48b808da2da437416/state.vscdb"
-```
+# Help on usage
+./chat.py export --help
+
+# Print all chats of the most recent workspace to the command line
+./chat.py export
+
+# Export all chats of the most recent workspace as Markdown
+./chat.py export --output-dir "/path/to/output"
+
+# Export only the latest chat of the most recent workspace
+./chat.py export --latest-tab --output-dir "/path/to/output"
+
+# Export all chats of a specifc workspace
+./chat.py export --output-dir "/path/to/output" "/path/to/workspaces/workspace-dir/state.vscdb"
+```

chat.py

@@ -9,6 +9,9 @@
 from rich.markdown import Markdown
 from loguru import logger
 import json
+import yaml
+import platform
+from pathlib import Path
 
 logger.remove()
 logger.add(sys.stderr, level="INFO")
@@ -18,15 +21,18 @@
 console = Console()
 
 @app.command()
-def export(db_path: str, output_dir: str = None):
+def export(
+    db_path: str = typer.Argument(None, help="The path to the SQLite database file. If not provided, the latest workspace will be used."),
+    output_dir: str = typer.Option(None, help="The directory where the output markdown files will be saved. If not provided, prints to command line."),
+    latest_tab: bool = typer.Option(False, "--latest-tab", help="Export only the latest tab. If not set, all tabs will be exported.")
+):
     """
     Export chat data from the database to markdown files or print it to the command line.
-
-    Args:
-        db_path (str): The path to the SQLite database file.
-        output_dir (str): The directory where the output markdown files will be saved. If not provided, prints to command line.
     """
-    image_dir = os.path.join(output_dir, 'images') if output_dir else None
+    if not db_path:
+        db_path = get_latest_workspace_db_path()
+
+    image_dir = None
 
     try:
         # Query the AI chat data from the database
@@ -41,6 +47,17 @@ def export(db_path: str, output_dir: str = None):
         # Convert the chat data from JSON string to dictionary
         chat_data_dict = json.loads(chat_data[0])
 
+        if latest_tab:
+            # Get the latest tab by timestamp
+            latest_tab = max(chat_data_dict['tabs'], key=lambda tab: tab.get('timestamp', 0))
+            chat_data_dict['tabs'] = [latest_tab]
+
+        # Check if there are any images in the chat data
+        has_images = any('image' in bubble for tab in chat_data_dict['tabs'] for bubble in tab.get('bubbles', []))
+
+        if has_images and output_dir:
+            image_dir = os.path.join(output_dir, 'images')
+
         # Format the chat data
         formatter = MarkdownChatFormatter()
         formatted_chats = formatter.format(chat_data_dict, image_dir)
@@ -57,6 +74,10 @@ def export(db_path: str, output_dir: str = None):
             for formatted_data in formatted_chats:
                 console.print(Markdown(formatted_data))
             logger.info("Chat data has been successfully printed to the command line")
+    except KeyError as e:
+        error_message = f"KeyError: {e}. The chat data structure is not as expected. Please check the database content."
+        logger.error(error_message)
+        raise typer.Exit(code=1)
     except json.JSONDecodeError as e:
         error_message = f"JSON decode error: {e}"
         logger.error(error_message)
@@ -70,16 +91,60 @@ def export(db_path: str, output_dir: str = None):
         logger.error(error_message)
         raise typer.Exit(code=1)
 
+def get_cursor_workspace_path() -> Path:
+    config_path = Path("config.yml")
+    logger.debug(f"Looking for configuration file at: {config_path}")
+    
+    if not config_path.exists():
+        error_message = f"Configuration file not found: {config_path}"
+        logger.error(error_message)
+        raise FileNotFoundError(error_message)
+
+    with open(config_path, 'r') as f:
+        config = yaml.safe_load(f)
+    logger.debug("Configuration file loaded successfully")
+
+    system = platform.system()
+    logger.debug(f"Detected operating system: {system}")
+
+    if system not in config["default_vscdb_dir_paths"]:
+        error_message = f"Unsupported operating system: {system}"
+        logger.error(error_message)
+        raise ValueError(error_message)
+
+    base_path = Path(os.path.expandvars(config["default_vscdb_dir_paths"][system])).expanduser()
+    logger.debug(f"Resolved base path: {base_path}")
+
+    if not base_path.exists():
+        error_message = f"Cursor workspace storage directory not found: {base_path}"
+        logger.error(error_message)
+        raise FileNotFoundError(error_message)
+
+    logger.info(f"Cursor workspace storage directory found: {base_path}")
+    return base_path
+
+def get_latest_workspace_db_path() -> str:
+    base_path = get_cursor_workspace_path()
+    workspace_folder = max(base_path.glob("*"), key=os.path.getmtime)
+    db_path = workspace_folder / "state.vscdb"
+    
+    if not db_path.exists():
+        raise FileNotFoundError(f"state.vscdb not found in {workspace_folder}")
+
+    return str(db_path)
+
 @app.command()
-def discover(directory: str, limit: int = None, search_text: str = None):
+def discover(
+    directory: str = typer.Argument(None, help="The directory to search for state.vscdb files. If not provided, the default Cursor workspace storage directory will be used."),
+    limit: int = typer.Option(None, help="The maximum number of state.vscdb files to process. Defaults to 10 if search_text is not provided, else -1."),
+    search_text: str = typer.Option(None, help="The text to search for in the chat history.")
+):
     """
     Discover all state.vscdb files in a directory and its subdirectories, and print a few lines of dialogue.
-
-    Args:
-        directory (str): The directory to search for state.vscdb files.
-        limit (int): The maximum number of state.vscdb files to process, sorted by most recent edits. Defaults to 10 if search_text is not provided, else -1.
-        search_text (str): The text to search for in the chat history. If provided, only chat entries containing this text will be printed.
     """
+    if not directory:
+        directory = str(get_cursor_workspace_path())
+    
     if limit is None:
         limit = -1 if search_text else 10
 
@@ -119,7 +184,6 @@ def discover(directory: str, limit: int = None, search_text: str = None):
                     for formatted_data in formatted_chats:
                         filtered_lines = [line for line in formatted_data.splitlines() if search_text.lower() in line.lower()]
                         if filtered_lines:
-                            # results.append((db_path, "[...]" + "  \n[...]  \n".join(filtered_lines[:10]) + "[...]"))
                             results.append((db_path, "\n".join(formatted_data.splitlines()[:10]) + "\n..."))
                     if not filtered_lines:
                         logger.debug(f"No chat entries containing '{search_text}' found in {db_path}")
@@ -153,4 +217,4 @@ def discover(directory: str, limit: int = None, search_text: str = None):
         raise typer.Exit(code=1)
 
 if __name__ == "__main__":
-    app()
+    app()

config.yml

@@ -1 +1,5 @@
+default_vscdb_dir_paths:
+  Windows: "%APPDATA%/Cursor/User/workspaceStorage"
+  Darwin: "~/Library/Application Support/Cursor/User/workspaceStorage"
+  Linux: "~/.config/Cursor/User/workspaceStorage"
 aichat_query: "SELECT value FROM ItemTable WHERE [key] IN ('workbench.panel.aichat.view.aichat.chatdata');"

requirements.txt

@@ -1,3 +1,4 @@
 loguru
 pyyaml
 typer
+rich

src/export.py

@@ -1,4 +1,3 @@
-
 import re
 import shutil
 import os
@@ -50,10 +49,12 @@ def format(self, chat_data: dict[str, Any], image_dir: str | None = 'images') ->
                             shutil.copy(image_path, new_image_path)
                             formatted_chat.append(f"![User Image]({new_image_path})\n")
                     elif bubble['type'] == 'ai':
+                        model_type = bubble.get('modelType', 'Unknown')
                         raw_text = re.sub(r'```python:[^\n]+', '```python', bubble['rawText'])
-                        formatted_chat.append(f"## AI:\n\n{raw_text}\n")
+                        formatted_chat.append(f"## AI ({model_type}):\n\n{raw_text}\n")
 
                 formatted_chats.append("\n".join(formatted_chat))
+
             return formatted_chats
         except KeyError as e:
             logger.error(f"KeyError: {e}")
@@ -126,5 +127,4 @@ def export(self, chat_data: dict[str, Any], output_dir: str, image_dir: str) ->
 # formatter = MarkdownChatFormatter()
 # saver = MarkdownFileSaver()
 # exporter = ChatExporter(formatter, saver)
-# exporter.export(chat_data, 'output_folder', 'images')
-
+# exporter.export(chat_data, 'output_folder', 'images')