ragflow_sync_cli/commands/sync_tasktracker.py - master - dictator/ragflow_sync

ragflow_sync_cli/commands/sync_tasktracker.py
188 строк6 KB
Developer
refactor: rename ttjson_to_markdown → ttjson2md
22 июн 2026, 14:00
22 июн 2026, 14:00babc303
"""
Sync TaskTracker to RAGFlow command.
"""

import json
import logging
import os
import sys
from typing import Optional

from ragflow_uploader.client import RAGFlowClient
from ragflow_uploader.uploader import Document, DocumentUploader
from tt_retriever.config import TaskTrackerConfig
from tt_retriever.retriever import TaskTrackerRetriever
from ttjson2md.core import convert_json_to_markdown
from cleanup.clean import preprocess_text


logger = logging.getLogger(__name__)


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 = True,
    wait_for_parsing: bool = False,
) -> None:
    """Sync TaskTracker tickets to RAGFlow."""
    logger.info(f"Syncing TaskTracker to dataset: {dataset_name}")

    try:
        # Initialize retriever
        retriever = TaskTrackerRetriever(config)

        # Get or create dataset
        uploader = DocumentUploader(client)
        if create_dataset:
            dataset_id = uploader.get_or_create_dataset(
                dataset_name,
                description="TaskTracker tickets",
            )
        else:
            # Find existing dataset
            datasets = client.list_datasets(name=dataset_name)
            if not datasets:
                logger.error(f"Dataset '{dataset_name}' not found")
                sys.exit(1)
            dataset_id = datasets[0]["id"]

        # Fetch tickets from TaskTracker
        tickets = retriever.fetch_data(
            tql_query=tql_query,
            project=project,
            status=status,
            ticket_ids=ticket_ids,
            recursive=recursive,
        )

        logger.info(f"Fetched {len(tickets)} tickets from TaskTracker")

        # Save raw JSON if requested
        if save_json:
            os.makedirs(save_json, exist_ok=True)
            json_file = os.path.join(save_json, f"{dataset_name.replace(' ', '_')}_raw.json")
            with open(json_file, "w", encoding="utf-8") as f:
                json.dump(tickets, f, ensure_ascii=False, indent=2)
            logger.info(f"Saved {len(tickets)} tickets to {json_file}")

        # Convert to documents
        documents = _convert_tickets_to_documents(tickets, config.url, cleanup=cleanup)

        # Save Markdown files if requested
        if save_md:
            os.makedirs(save_md, exist_ok=True)
            md_files_saved = 0
            for doc in documents:
                safe_filename = doc.file_name.replace(" ", "_")
                md_file = os.path.join(save_md, safe_filename)
                with open(md_file, "w", encoding="utf-8") as f:
                    f.write(doc.content)
                md_files_saved += 1
            logger.info(f"Saved {md_files_saved} markdown files to {save_md}")

        # Fetch existing documents for deduplication
        try:
            existing_docs = client.list_documents(dataset_id, page_size=1000)
        except Exception as e:
            logger.warning(f"Failed to list existing documents: {e}")
            existing_docs = []

        # Upload to RAGFlow
        result = uploader.upload_documents(
            dataset_id=dataset_id,
            documents=documents,
            existing_documents=existing_docs,
            dedup_key="ticket_id",
            parse=parse,
            wait_for_parsing=wait_for_parsing,
        )

        # Print results
        _print_sync_results(result)

        if not result.success:
            sys.exit(1)

    except Exception as e:
        logger.error(f"TaskTracker sync failed: {e}")
        sys.exit(1)


def _convert_tickets_to_documents(
    tickets: list[dict],
    base_url: str,
    cleanup: bool = True,
) -> list[Document]:
    """Convert TaskTracker tickets to Document objects."""
    documents = []

    for ticket in tickets:
        try:
            # Convert ticket to markdown using converter
            markdown_content = convert_json_to_markdown(ticket)

            # Apply cleanup if enabled
            if cleanup:
                markdown_content = preprocess_text(markdown_content)

            # Extract metadata
            ticket_id = ticket.get("code", ticket.get("id", ""))
            title = ticket.get("summary", "Untitled")
            status = (
                ticket.get("status", {}).get("name", "")
                if isinstance(ticket.get("status"), dict)
                else str(ticket.get("status", ""))
            )

            # Build document
            doc = Document(
                title=f"[{ticket_id}] {title}",
                content=markdown_content,
                metadata={
                    "source": "tasktracker",
                    "ticket_id": ticket_id,
                    "status": status,
                    "url": f"{base_url.rstrip('/')}/swtr/units/all/unit/{ticket_id}",
                },
            )

            documents.append(doc)

        except Exception as e:
            logger.warning(f"Failed to convert ticket {ticket.get('code', 'unknown')}: {e}")

    return documents


def _print_sync_results(result) -> None:
    """Print sync results to stdout."""
    print("\n" + "=" * 50)
    print("Sync Results:")
    print("=" * 50)
    print(f"Status: {'✓ Success' if result.success else '✗ Failed'}")
    print(f"Documents created: {result.documents_created}")
    print(f"Documents updated: {result.documents_updated}")
    print(f"Documents failed: {result.documents_failed}")

    if result.documents_parsed > 0:
        print(f"Documents parsed: {result.documents_parsed}")
    if result.documents_parse_failed > 0:
        print(f"Documents parse failed: {result.documents_parse_failed}")

    if result.errors:
        print("\nErrors:")
        for error in result.errors:
            print(f"  - {error}")

    print()