ragflow_sync_cli/main.py - master - dictator/ragflow_sync

ragflow_sync_cli/main.py
503 строки15 KB
Developer
feat: add sync files command to upload Markdown files from local directory
05 июн 2026, 18:05
05 июн 2026, 18:0569cced8
#!/usr/bin/env python3
"""
RAGFlow Sync - Custom Datasource CLI

Command-line interface for syncing data from external sources (Confluence, TaskTracker)
to RAGFlow datasets.

Usage:
    # List available datasources
    python -m ragflow_sync_cli list-sources

    # List datasets in RAGFlow
    python -m ragflow_sync_cli list-datasets

    # Sync Confluence to RAGFlow
    python -m ragflow_sync_cli sync confluence --dataset "Confluence KB" --spaces PROJ,HACK

    # Sync TaskTracker to RAGFlow
    python -m ragflow_sync_cli sync tasktracker --dataset "TaskTracker KB" --project HACKATON

    # Sync all configured datasources
    python -m ragflow_sync_cli sync all
"""

import argparse
import logging
import sys
from typing import Optional

from ragflow_sync_cli.config import Config, load_config
from ragflow_uploader.client import RAGFlowClient, RAGFlowError
from ragflow_uploader.config import RAGFlowConfig
from tt_retriever.config import TaskTrackerConfig
from jira_confluence_retriever.config import ConfluenceConfig


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


def setup_logging(verbose: bool = False) -> None:
    """Configure logging level."""
    level = logging.DEBUG if verbose else logging.INFO
    logging.getLogger("ragflow_sync").setLevel(level)
    logging.getLogger("ragflow_uploader").setLevel(level)
    logging.getLogger("tt_retriever").setLevel(level)
    logging.getLogger("jira_confluence_retriever").setLevel(level)


def list_sources(config: Config) -> None:
    """List available datasources."""
    from ragflow_sync_cli.commands.list_sources import list_sources as _list_sources

    _list_sources(config)


def list_datasets(client: RAGFlowClient) -> None:
    """List all datasets in RAGFlow."""
    from ragflow_sync_cli.commands.list_datasets import list_datasets as _list_datasets

    _list_datasets(client)


def sync_confluence(
    client: RAGFlowClient,
    config: ConfluenceConfig,
    dataset_name: str,
    spaces: Optional[list[str]] = None,
    query: Optional[str] = None,
    page_ids: Optional[list[str]] = None,
    limit: int = 100,
    create_dataset: bool = True,
    recursive: bool = True,
    cleanup: bool = True,
    parse: bool = False,
    wait_for_parsing: bool = False,
) -> None:
    """Sync Confluence pages to RAGFlow."""
    from ragflow_sync_cli.commands.sync_confluence import sync_confluence as _sync_confluence

    _sync_confluence(
        client=client,
        config=config,
        dataset_name=dataset_name,
        spaces=spaces,
        query=query,
        page_ids=page_ids,
        limit=limit,
        create_dataset=create_dataset,
        recursive=recursive,
        cleanup=cleanup,
        parse=parse,
        wait_for_parsing=wait_for_parsing,
    )


def sync_tasktracker(
    client: RAGFlowClient,
    config: TaskTrackerConfig,
    dataset_name: str,
    tql_query: Optional[str] = None,
    project: Optional[str] = None,
    status: Optional[str] = None,
    ticket_ids: Optional[list[str]] = None,
    create_dataset: bool = True,
    save_json: Optional[str] = None,
    save_md: Optional[str] = None,
    recursive: bool = True,
    cleanup: bool = True,
    parse: bool = False,
    wait_for_parsing: bool = False,
) -> None:
    """Sync TaskTracker tickets to RAGFlow."""
    from ragflow_sync_cli.commands.sync_tasktracker import sync_tasktracker as _sync_tasktracker

    _sync_tasktracker(
        client=client,
        config=config,
        dataset_name=dataset_name,
        tql_query=tql_query,
        project=project,
        status=status,
        ticket_ids=ticket_ids,
        create_dataset=create_dataset,
        save_json=save_json,
        save_md=save_md,
        recursive=recursive,
        cleanup=cleanup,
        parse=parse,
        wait_for_parsing=wait_for_parsing,
    )


def sync_files(
    client: RAGFlowClient,
    directory: str,
    dataset_name: str,
    glob: str = "*.md",
    exclude: Optional[list[str]] = None,
    create_dataset: bool = True,
    cleanup: bool = True,
    parse: bool = True,
    wait_for_parsing: bool = False,
) -> None:
    """Sync Markdown files from a local directory to RAGFlow."""
    from ragflow_sync_cli.commands.sync_files import sync_files as _sync_files

    _sync_files(
        client=client,
        directory=directory,
        dataset_name=dataset_name,
        glob=glob,
        exclude=exclude,
        create_dataset=create_dataset,
        cleanup=cleanup,
        parse=parse,
        wait_for_parsing=wait_for_parsing,
    )


def sync_all(
    client: RAGFlowClient,
    config: Config,
    recursive: bool = True,
) -> None:
    """Sync all configured datasources."""
    from ragflow_sync_cli.commands.sync_all import sync_all as _sync_all

    _sync_all(client, config, recursive=recursive)


def main() -> None:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="RAGFlow Custom Datasource CLI",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )

    parser.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="Enable verbose output",
    )

    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # List sources command
    subparsers.add_parser(
        "list-sources",
        help="List available datasources",
    )

    # List datasets command
    subparsers.add_parser(
        "list-datasets",
        help="List datasets in RAGFlow",
    )

    # Sync command
    sync_parser = subparsers.add_parser(
        "sync",
        help="Sync data from external sources to RAGFlow",
    )
    sync_subparsers = sync_parser.add_subparsers(
        dest="source",
        help="Data source to sync",
        required=True,
    )

    # Sync Confluence
    confluence_parser = sync_subparsers.add_parser(
        "confluence",
        help="Sync Confluence pages",
    )
    confluence_parser.add_argument(
        "--dataset",
        "-d",
        default="Confluence Knowledge Base",
        help="Dataset name (default: 'Confluence Knowledge Base')",
    )
    confluence_parser.add_argument(
        "--spaces",
        "-s",
        help="Comma-separated list of space keys (e.g., 'PROJ,HACK')",
    )
    confluence_parser.add_argument(
        "--query",
        "-q",
        help="CQL search query",
    )
    confluence_parser.add_argument(
        "--page-ids",
        help="Comma-separated list of page IDs",
    )
    confluence_parser.add_argument(
        "--limit",
        "-l",
        type=int,
        default=100,
        help="Maximum number of pages to sync (default: 100)",
    )
    confluence_parser.add_argument(
        "--no-create",
        action="store_true",
        help="Don't create dataset if it doesn't exist",
    )
    confluence_parser.add_argument(
        "--no-recursive",
        action="store_true",
        help="Disable recursive child page fetching (sync only direct children of parent page)",
    )
    confluence_parser.add_argument(
        "--no-cleanup",
        action="store_true",
        help="Disable text cleanup before upload (preserves special characters, images, long URLs)",
    )
    confluence_parser.add_argument(
        "--no-parse",
        action="store_true",
        help="Disable document parsing after upload (default: enabled)",
    )
    confluence_parser.add_argument(
        "--wait-for-parsing",
        action="store_true",
        help="Wait for parsing to complete before returning",
    )

    # Sync TaskTracker
    tt_parser = sync_subparsers.add_parser(
        "tasktracker",
        help="Sync TaskTracker tickets",
    )
    tt_parser.add_argument(
        "--dataset",
        "-d",
        default="TaskTracker Knowledge Base",
        help="Dataset name (default: 'TaskTracker Knowledge Base')",
    )
    tt_parser.add_argument(
        "--tql",
        "-t",
        help='Custom TQL query (e.g., \'space ="CLDSPRCH" AND suit IN ("bug", "task")\')',
    )
    tt_parser.add_argument(
        "--project",
        "-p",
        help="Project/space code to filter (e.g., 'HACKATON') - used if --tql not provided",
    )
    tt_parser.add_argument(
        "--status",
        help="Status to filter (e.g., 'Open', 'In Progress')",
    )
    tt_parser.add_argument(
        "--ticket-ids",
        help="Comma-separated list of ticket IDs",
    )
    tt_parser.add_argument(
        "--no-create",
        action="store_true",
        help="Don't create dataset if it doesn't exist",
    )
    tt_parser.add_argument(
        "--save-json",
        help="Save original JSON data to specified directory before upload",
    )
    tt_parser.add_argument(
        "--save-md",
        help="Save converted Markdown files to specified directory before upload",
    )
    tt_parser.add_argument(
        "--no-recursive",
        action="store_true",
        help="Disable automatic wiki page tree expansion (sync only the root wiki page)",
    )
    tt_parser.add_argument(
        "--no-cleanup",
        action="store_true",
        help="Disable text cleanup before upload (preserves special characters, images, long URLs)",
    )
    tt_parser.add_argument(
        "--no-parse",
        action="store_true",
        help="Disable document parsing after upload (default: enabled)",
    )
    tt_parser.add_argument(
        "--wait-for-parsing",
        action="store_true",
        help="Wait for parsing to complete before returning",
    )

    # Sync files
    files_parser = sync_subparsers.add_parser(
        "files",
        help="Sync Markdown files from a local directory",
    )
    files_parser.add_argument(
        "--path",
        "-p",
        required=True,
        help="Local directory path to scan for Markdown files",
    )
    files_parser.add_argument(
        "--dataset",
        "-d",
        default="Files Knowledge Base",
        help="Dataset name (default: 'Files Knowledge Base')",
    )
    files_parser.add_argument(
        "--glob",
        default="*.md",
        help="Glob pattern for files to include (default: '*.md')",
    )
    files_parser.add_argument(
        "--exclude",
        action="append",
        default=[],
        help="Patterns to exclude (can be repeated, e.g. --exclude '*draft*' --exclude 'vendor/*')",
    )
    files_parser.add_argument(
        "--no-create",
        action="store_true",
        help="Don't create dataset if it doesn't exist",
    )
    files_parser.add_argument(
        "--no-cleanup",
        action="store_true",
        help="Disable text cleanup before upload (preserves special characters, images, long URLs)",
    )
    files_parser.add_argument(
        "--no-parse",
        action="store_true",
        help="Disable document parsing after upload (default: enabled)",
    )
    files_parser.add_argument(
        "--wait-for-parsing",
        action="store_true",
        help="Wait for parsing to complete before returning",
    )

    # Sync all
    sync_all_parser = sync_subparsers.add_parser(
        "all",
        help="Sync all configured datasources",
    )
    sync_all_parser.add_argument(
        "--no-recursive",
        action="store_true",
        help="Disable recursive fetching (wiki tree expansion / child page traversal)",
    )

    args = parser.parse_args()

    # Setup logging
    setup_logging(args.verbose)

    # Load configuration
    try:
        config = load_config()
    except ValueError as e:
        logger.error(f"Configuration error: {e}")
        sys.exit(1)

    # Initialize RAGFlow client
    try:
        client = RAGFlowClient(config.ragflow, verify_ssl=config.ragflow.verify_ssl)
    except Exception as e:
        logger.error(f"Failed to initialize RAGFlow client: {e}")
        sys.exit(1)

    # Execute command
    try:
        if args.command == "list-sources":
            list_sources(config)

        elif args.command == "list-datasets":
            list_datasets(client)

        elif args.command == "sync":
            if args.source == "confluence":
                spaces = None
                if args.spaces:
                    spaces = [s.strip() for s in args.spaces.split(",")]

                page_ids = None
                if args.page_ids:
                    page_ids = [p.strip() for p in args.page_ids.split(",")]

                sync_confluence(
                    client=client,
                    config=config.confluence
                    or ConfluenceConfig(
                        url="https://placeholder",
                        username="placeholder",
                        api_token="placeholder",
                    ),
                    dataset_name=args.dataset,
                    spaces=spaces,
                    query=args.query,
                    page_ids=page_ids,
                    limit=args.limit,
                    create_dataset=not args.no_create,
                    recursive=not args.no_recursive,
                    cleanup=not args.no_cleanup,
                    parse=not args.no_parse,
                    wait_for_parsing=args.wait_for_parsing,
                )

            elif args.source == "tasktracker":
                ticket_ids = None
                if args.ticket_ids:
                    ticket_ids = [t.strip() for t in args.ticket_ids.split(",")]

                sync_tasktracker(
                    client=client,
                    config=config.tasktracker
                    or TaskTrackerConfig(
                        url="https://placeholder",
                        api_token="placeholder",
                    ),
                    dataset_name=args.dataset,
                    tql_query=args.tql,
                    project=args.project,
                    status=args.status,
                    ticket_ids=ticket_ids,
                    create_dataset=not args.no_create,
                    save_json=args.save_json,
                    save_md=args.save_md,
                    recursive=not args.no_recursive,
                    cleanup=not args.no_cleanup,
                    parse=not args.no_parse,
                    wait_for_parsing=args.wait_for_parsing,
                )

            elif args.source == "files":
                sync_files(
                    client=client,
                    directory=args.path,
                    dataset_name=args.dataset,
                    glob=args.glob,
                    exclude=args.exclude,
                    create_dataset=not args.no_create,
                    cleanup=not args.no_cleanup,
                    parse=not args.no_parse,
                    wait_for_parsing=args.wait_for_parsing,
                )

            elif args.source == "all":
                sync_all(client, config, recursive=not args.no_recursive)

        else:
            parser.print_help()

    finally:
        client.close()


if __name__ == "__main__":
    main()