按分组条件同步知识库

sync-knowledge-to-dify-group.py

+import requests
+import pymysql
+import sys
+from datetime import datetime, timedelta
+import logging
+import argparse
+import re
+import os
+import time
+
+# Configure logging
+log_path = os.path.join(os.path.dirname(__file__), 'sync_log_group.log')
+logging.basicConfig(
+    filename=log_path,
+    level=logging.INFO,
+    format='%(asctime)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+
+
+# Define a function to replace print statements with logging
+def log_message(message):
+    logging.info(message)
+    print(message)
+
+
+# Configuration parameters (modify according to your environment)
+DIFY_API_BASE_URL = 'http://192.168.141.145/v1/datasets'
+DATASET_ID = 'ad4b7b26-e12b-4e8f-9d40-e927d57baa08'
+API_KEY = 'dataset-eETdI20aVtQwr9w89Fnsu9uJ'
+CATEGORY_TITLE = 'was'
+GROUP_TITLE = '监控小组'
+DB_CONFIG = {
+    'host': 'localhost',
+    'user': 'root',
+    'password': '123456',
+    'db': 'knowledge_base',
+    'charset': 'utf8mb4',
+    'cursorclass': pymysql.cursors.DictCursor,
+}
+RERANK_MAX_LENGTH = 510  # Max segment length (including alarm_name)
+OVERLAP_LENGTH = 0  # Segment overlap length
+SYNC_STATE_FILE = os.path.join(os.path.dirname(__file__), 'sync_state_group.txt')
+DOCUMENT_READY_TIMEOUT = 10  # Document ready timeout (seconds)
+POLLING_INTERVAL = 1  # Polling interval (seconds)
+
+
+# Enhanced timestamp management
+def read_sync_state():
+    """Initialize sync timestamp strategy"""
+    try:
+        with open(SYNC_STATE_FILE, 'r') as f:
+            content = f.read().strip()
+            if not content:
+                raise ValueError("Trigger full sync initialization")
+            return datetime.strptime(content, '%Y-%m-%d %H:%M:%S')
+    except (FileNotFoundError, ValueError):
+        initial_time = datetime.now() - timedelta(days=365 * 30)
+        with open(SYNC_STATE_FILE, 'w') as f:
+            f.write(initial_time.strftime('%Y-%m-%d %H:%M:%S'))
+        return initial_time
+
+
+def save_sync_state(timestamp):
+    """Atomic timestamp storage"""
+    with open(SYNC_STATE_FILE, 'w') as f:
+        f.write(timestamp.strftime('%Y-%m-%d %H:%M:%S'))
+
+
+# Core data synchronization module
+def fetch_records(sync_type, last_sync):
+    """Dynamic SQL generator"""
+    conn = pymysql.connect(**DB_CONFIG)
+    try:
+        with conn.cursor() as cursor:
+            base_sql = """
+            SELECT b.id, b.alarm_name, b.solution_html, b.update_datetime, b.is_delete
+            FROM api_knowledge b
+            WHERE b.is_audit=1 AND b.is_publish=1 and b.alarm_name is not null and b.solution_html is not null
+            AND b.group_category_id IN (SELECT a.id FROM api_groupcategory a where a.category_id in 
+            (SELECT t.id from api_category t where t.title = %s )
+            AND a.group_id IN (SELECT t.id FROM api_group t WHERE t.title = %s))
+            """
+            if sync_type == 'full':
+                sql = base_sql
+                log_message(f"{sql}")
+                cursor.execute(sql,(CATEGORY_TITLE,GROUP_TITLE))
+            else:
+                sql = base_sql + " AND (update_datetime > %s or create_datetime > %s) ORDER BY create_datetime ASC"
+                log_message(f"{sql}")
+                cursor.execute(sql, (CATEGORY_TITLE,GROUP_TITLE,last_sync, last_sync))
+            return cursor.fetchall()
+    finally:
+        conn.close()
+
+
+def clean_solution_html(solution_html):
+    """Clean solution_html: keep img tags, remove all other HTML tags"""
+    return re.sub(r'&nbsp;|<(?!img\b)[^>]*>', '', solution_html)
+
+
+def split_content(alarm_name, content, labels=None):
+    """Split content into segments with overlap control. Each segment starts with alarm_name followed by a newline. Total length does not exceed RERANK_MAX_LENGTH."""
+    segments = []
+    start = 0
+    content_length = len(content)
+
+    # Handle empty alarm name
+    if not alarm_name:
+        alarm_name = ""
+        
+    label_str = ",".join(labels) if labels else "无"
+    
+    # Context template
+    context_prefix = f"{alarm_name}\n标签：{label_str}\n"
+
+    # Calculate the length occupied by the context prefix
+    header_length = len(context_prefix)
+
+    # Calculate the max length for the content portion
+    max_content_length = RERANK_MAX_LENGTH - header_length
+
+    # Ensure overlap length does not exceed max content length
+    effective_overlap = min(OVERLAP_LENGTH, max_content_length)
+
+    # If context prefix is too long, return it as a single segment
+    if header_length > RERANK_MAX_LENGTH:
+        segments.append(context_prefix[:RERANK_MAX_LENGTH])
+        return segments
+
+    while start < content_length:
+        # Calculate the end position of the current segment
+        end = start + max_content_length
+
+        # If this is the last segment, take content to the end
+        if end >= content_length:
+            segment_content = content[start:]
+            # Build segment prefixed with context
+            formatted_segment = f"{context_prefix}{segment_content}"
+            segments.append(formatted_segment)
+            break
+
+        # Extract current segment
+        segment_content = content[start:end]
+        # Build segment prefixed with context
+        formatted_segment = f"{context_prefix}{segment_content}"
+        segments.append(formatted_segment)
+
+        # Calculate next segment start position (considering overlap)
+        next_start = end - effective_overlap
+        # Prevent start position from going backwards after overlap
+        if next_start <= start:
+            next_start = start + 1  # Advance at least 1 char to avoid infinite loop
+        start = next_start
+
+    # Filter out empty segments
+    return [s for s in segments if s.strip()]
+
+
+def get_labels_by_knowledge_id(knowledge_id):
+    """Get label list by knowledge_id"""
+    conn = pymysql.connect(**DB_CONFIG)
+    try:
+        with conn.cursor() as cursor:
+            sql = """
+            SELECT ul.title
+            FROM api_knowledge_open_label kol
+            JOIN api_userlabel ul ON kol.userlabel_id = ul.id
+            WHERE kol.knowledge_id = %s
+            """
+            cursor.execute(sql, (knowledge_id,))
+            results = cursor.fetchall()
+            return [row['title'] for row in results]
+    finally:
+        conn.close()
+
+
+def get_document_id_by_name(alarm_name):
+    """Query document ID by document name"""
+    headers = {
+        'Authorization': f'Bearer {API_KEY}',
+        'Content-Type': 'application/json'
+    }
+    url = f"{DIFY_API_BASE_URL}/{DATASET_ID}/documents?keyword={alarm_name}"
+
+    try:
+        response = requests.get(url, headers=headers)
+        log_message(f"[get_document_id_by_name] GET {url} - Status: {response.status_code}")
+        response.raise_for_status()
+
+        try:
+            data = response.json()
+        except requests.exceptions.JSONDecodeError as e:
+            log_message(f"[ERROR] JSON decode error in get_document_id_by_name: {e}")
+            log_message(f"[ERROR] Status Code: {response.status_code}")
+            log_message(f"[ERROR] Response Headers: {dict(response.headers)}")
+            log_message(f"[ERROR] Response Content: {response.text[:1000]}")
+            raise
+
+        # Iterate results to find exact match by document name
+        for doc in data.get('data', []):
+            if doc.get('name') == alarm_name:
+                return doc.get('id')
+
+        return None
+    except requests.exceptions.RequestException as e:
+        log_message(f"[ERROR] Request failed in get_document_id_by_name: {e}")
+        if hasattr(e, 'response') and e.response is not None:
+            log_message(f"[ERROR] Status Code: {e.response.status_code}")
+            log_message(f"[ERROR] Response Content: {e.response.text[:1000]}")
+        raise
+
+
+def create_document(alarm_name):
+    """Create a document"""
+    headers = {
+        'Authorization': f'Bearer {API_KEY}',
+        'Content-Type': 'application/json'
+    }
+    create_url = f"{DIFY_API_BASE_URL}/{DATASET_ID}/document/create-by-text"
+    payload = {
+        "name": alarm_name,
+        "text": '',
+        "indexing_technique": "high_quality",
+        "process_rule": {
+            "mode": "custom",
+            "rules": {
+                "pre_processing_rules": [
+                    {"id": "remove_extra_spaces", "enabled": True},
+                    {"id": "remove_urls_emails", "enabled": False}
+                ],
+                "segmentation": {
+                    "separator": "\n",
+                    "max_tokens": 1024
+                }
+            }
+        },
+    }
+
+    try:
+        response = requests.post(create_url, headers=headers, json=payload)
+        log_message(f"[create_document] POST {create_url} - Status: {response.status_code}")
+        response.raise_for_status()
+
+        try:
+            result = response.json()
+        except requests.exceptions.JSONDecodeError as e:
+            log_message(f"[ERROR] JSON decode error in create_document: {e}")
+            log_message(f"[ERROR] Status Code: {response.status_code}")
+            log_message(f"[ERROR] Response Headers: {dict(response.headers)}")
+            log_message(f"[ERROR] Response Content: {response.text[:1000]}")
+            raise
+
+        return result['document']['id']
+    except requests.exceptions.RequestException as e:
+        log_message(f"[ERROR] Request failed in create_document: {e}")
+        if hasattr(e, 'response') and e.response is not None:
+            log_message(f"[ERROR] Status Code: {e.response.status_code}")
+            log_message(f"[ERROR] Response Content: {e.response.text[:1000]}")
+        raise
+
+
+def is_document_ready(document_id, segments, alarm_name, labels=None):
+    """Check if document is ready by attempting to create a segment, and return the test segment ID"""
+    headers = {
+        'Authorization': f'Bearer {API_KEY}',
+        'Content-Type': 'application/json'
+    }
+    create_segment_url = f"{DIFY_API_BASE_URL}/{DATASET_ID}/documents/{document_id}/segments"
+
+    valid_segments = [s for s in segments if s.strip()]
+    if not valid_segments:
+        log_message(f"No valid segments to check readiness for document {document_id}")
+        return False, None
+
+    test_segment = valid_segments[0]
+    # Merge labels and alarm_name as keywords
+    keywords = labels if labels else []
+    if alarm_name not in keywords:
+        keywords.insert(0, alarm_name)
+    payload = {
+        "segments": [{"content": test_segment, "keywords": keywords}]
+    }
+
+    try:
+        response = requests.post(create_segment_url, headers=headers, json=payload)
+        log_message(f"[is_document_ready] POST {create_segment_url} - Status: {response.status_code}")
+        response.raise_for_status()
+
+        try:
+            result = response.json()
+        except requests.exceptions.JSONDecodeError as e:
+            log_message(f"[ERROR] JSON decode error in is_document_ready: {e}")
+            log_message(f"[ERROR] Status Code: {response.status_code}")
+            log_message(f"[ERROR] Response Headers: {dict(response.headers)}")
+            log_message(f"[ERROR] Response Content: {response.text[:1000]}")
+            raise
+
+        segment_id = result['data'][0]['id']
+        log_message(f"Created test segment ID: {segment_id} for document {document_id}")
+        return True, segment_id
+    except requests.exceptions.HTTPError as e:
+        if e.response.status_code == 404:
+            log_message(f"Document {document_id} not ready yet - 404 error")
+            return False, None
+        log_message(f"[ERROR] HTTP error in is_document_ready: {e}")
+        log_message(f"[ERROR] Status Code: {e.response.status_code}")
+        log_message(f"[ERROR] Response Content: {e.response.text[:1000]}")
+        raise
+    except requests.exceptions.RequestException as e:
+        log_message(f"[ERROR] Request failed in is_document_ready: {e}")
+        if hasattr(e, 'response') and e.response is not None:
+            log_message(f"[ERROR] Status Code: {e.response.status_code}")
+            log_message(f"[ERROR] Response Content: {e.response.text[:1000]}")
+        raise
+
+
+def wait_for_document_ready(document_id, segments, alarm_name, labels=None):
+    """Wait for document to be ready and return test segment ID"""
+    start_time = time.time()
+    log_message(f"Waiting for document {document_id} to be ready...")
+
+    while time.time() - start_time < DOCUMENT_READY_TIMEOUT:
+        try:
+            ready, segment_id = is_document_ready(document_id, segments, alarm_name, labels)
+            if ready:
+                test_segment_id = segment_id
+                log_message(f"Document {document_id} is ready, test segment ID: {test_segment_id}")
+                return True, test_segment_id
+        except Exception as e:
+            log_message(f"Temporary error while checking document readiness: {str(e)}")
+
+        log_message(f"Document not ready yet, retrying in {POLLING_INTERVAL} seconds...")
+        time.sleep(POLLING_INTERVAL)
+
+    log_message(f"Document {document_id} did not become ready within timeout")
+    return False, None
+
+
+def create_segments(document_id, segments, alarm_name, labels=None):
+    """Create document segments, excluding the first test segment"""
+    headers = {
+        'Authorization': f'Bearer {API_KEY}',
+        'Content-Type': 'application/json'
+    }
+    create_segment_url = f"{DIFY_API_BASE_URL}/{DATASET_ID}/documents/{document_id}/segments"
+
+    valid_segments = [s for s in segments if s.strip()]
+    if not valid_segments:
+        log_message(f"No valid segments to create for document {document_id}")
+        return {"message": "No valid segments provided"}
+
+    # Wait for document to be ready and get test segment ID
+    ready, test_segment_id = wait_for_document_ready(document_id, valid_segments, alarm_name, labels)
+    if not ready:
+        raise Exception(f"Document {document_id} not ready after waiting")
+
+    # Insert all segments (excluding the first test segment)
+    segments_to_create = valid_segments[1:]
+    if not segments_to_create:
+        log_message(f"Only 1 segment available, no new segments to create for document {document_id}")
+        return {"message": "No new segments to create"}
+
+    # Merge labels and alarm_name as keywords
+    keywords = labels if labels else []
+    if alarm_name not in keywords:
+        keywords.insert(0, alarm_name)
+
+    payload = {
+        "segments": [{"content": s, "keywords": keywords} for s in segments_to_create]
+    }
+
+    # Log API request parameters
+    log_message(f"[create_segments] POST {create_segment_url}")
+    log_message(f"[create_segments] Creating {len(segments_to_create)} segments")
+
+    try:
+        response = requests.post(create_segment_url, headers=headers, json=payload)
+        log_message(f"[create_segments] Status: {response.status_code}")
+        response.raise_for_status()
+
+        try:
+            result = response.json()
+        except requests.exceptions.JSONDecodeError as e:
+            log_message(f"[ERROR] JSON decode error in create_segments: {e}")
+            log_message(f"[ERROR] Status Code: {response.status_code}")
+            log_message(f"[ERROR] Response Headers: {dict(response.headers)}")
+            log_message(f"[ERROR] Response Content: {response.text[:1000]}")
+            raise
+
+        log_message(f"Created {len(segments_to_create)} segments for document {document_id}")
+        return result
+    except requests.exceptions.RequestException as e:
+        log_message(f"[ERROR] Request failed in create_segments: {e}")
+        if hasattr(e, 'response') and e.response is not None:
+            log_message(f"[ERROR] Status Code: {e.response.status_code}")
+            log_message(f"[ERROR] Response Content: {e.response.text[:1000]}")
+        raise
+
+
+def delete_document(document_id):
+    """Delete a document"""
+    headers = {
+        'Authorization': f'Bearer {API_KEY}',
+        'Content-Type': 'application/json'
+    }
+    delete_url = f"{DIFY_API_BASE_URL}/{DATASET_ID}/documents/{document_id}"
+
+    try:
+        response = requests.delete(delete_url, headers=headers)
+        log_message(f"[delete_document] DELETE {delete_url} - Status: {response.status_code}")
+        response.raise_for_status()
+
+        # Dify API returns 204 No Content on successful deletion (no response body)
+        if response.status_code == 204:
+            log_message(f"Deleted document with ID: {document_id}")
+            return {"result": "success", "message": "Document deleted"}
+
+        # Handle other success codes with potential JSON response
+        if response.text.strip():
+            try:
+                return response.json()
+            except requests.exceptions.JSONDecodeError:
+                log_message(f"[WARNING] Non-JSON response: {response.text[:500]}")
+                return {"result": "success", "message": "Document deleted"}
+
+        return {"result": "success", "message": "Document deleted"}
+    except requests.exceptions.RequestException as e:
+        log_message(f"[ERROR] Request failed in delete_document: {e}")
+        if hasattr(e, 'response') and e.response is not None:
+            log_message(f"[ERROR] Status Code: {e.response.status_code}")
+            log_message(f"[ERROR] Response Content: {e.response.text[:1000]}")
+        raise
+
+
+def sync_record(record):
+    """Sync a single record to Dify"""
+    alarm_name = record['alarm_name']
+    knowledge_id = record['id']
+    solution_html = clean_solution_html(record['solution_html'])
+    log_message(f"Processing record: {alarm_name}")
+
+    # Get knowledge base labels
+    labels = get_labels_by_knowledge_id(knowledge_id)
+    if labels:
+        log_message(f"Labels for '{alarm_name}': {labels}")
+        
+    segments = split_content(alarm_name, solution_html, labels)
+
+    # Handle deletion logic
+    if record['is_delete'] == 1:
+        document_id = get_document_id_by_name(alarm_name)
+        if document_id:
+            delete_document(document_id)
+            log_message(f"Deleted document: {alarm_name}")
+            return "deleted"
+        log_message(f"Document not found for deletion: {alarm_name}")
+        return "not_found"
+
+    # Handle update or create logic
+    document_id = get_document_id_by_name(alarm_name)
+    if document_id:
+        delete_document(document_id)
+        log_message(f"Deleted existing document: {alarm_name}, ID: {document_id}")
+
+    valid_segments = [s for s in segments if s.strip()]
+    if not valid_segments:
+        log_message(f"Skipping document creation for '{alarm_name}' - no valid content")
+        return "skipped"
+
+    # Create new document and segments
+    dify_id = create_document(alarm_name)
+    create_segments(dify_id, valid_segments, alarm_name, labels)
+    log_message(f"Created new document: {alarm_name}, ID: {dify_id} with {len(valid_segments) - 1} segments")
+    return "recreated"
+
+
+# Progress visualization module
+def print_progress(idx, total, prefix=""):
+    """Terminal progress bar visualization"""
+    bar_length = 40
+    filled = int(bar_length * idx / total)
+    bar = '#' * filled + '-' * (bar_length - filled)
+    message = f"\r{prefix} |{bar}| {idx}/{total} ({idx / total:.0%}) "
+    sys.stdout.write(message)
+    sys.stdout.flush()
+
+
+def main():
+    """Main control flow"""
+    parser = argparse.ArgumentParser(description='Sync data with Dify API')
+    parser.add_argument('--sync-type', choices=['full', 'incremental'], default=None,
+                        help='Specify sync type: full or incremental')
+    args = parser.parse_args()
+    last_sync = read_sync_state()
+    if args.sync_type:
+        sync_type = args.sync_type
+    else:
+        sync_type = 'full' if last_sync.year < 2010 else 'incremental'
+    log_message(f"{sync_type.upper()} SYNC MODE | Base time: {last_sync}")
+    records = fetch_records(sync_type, last_sync)
+    log_message(f"{sync_type.upper()} SYNC MODE | Records: {len(records)}")
+    if not records:
+        log_message("NO INCREMENTAL DATA NEED TO SYNC")
+        return
+    total = len(records)
+    max_update_datetime = last_sync
+    error_log = []
+    deleted_count = 0
+    recreated_count = 0
+    skipped_count = 0
+    for idx, record in enumerate(records, 1):
+        try:
+            result = sync_record(record)
+            if result == "deleted":
+                deleted_count += 1
+            elif result == "recreated":
+                recreated_count += 1
+            elif result == "skipped":
+                skipped_count += 1
+            current_update = record['update_datetime']
+            if current_update and current_update > max_update_datetime:
+                max_update_datetime = current_update
+            print_progress(idx, total, prefix=f"SYNCING")
+        except Exception as e:
+            error_msg = f"ID:{record['id']} NAME:{record['alarm_name']}"
+            error_log.append(f"{error_msg} ERROR:{str(e)}")
+            log_message(f"[ERROR] {error_msg}")
+            log_message(f"[ERROR] Exception Type: {type(e).__name__}")
+            log_message(f"[ERROR] Exception Details: {str(e)}")
+            import traceback
+            log_message(f"[ERROR] Traceback:\n{traceback.format_exc()}")
+    log_message("\n" + "=" * 50)
+    save_sync_state(datetime.now())
+    log_message(
+        f"SYNC COMPLETED | SUCCESS: {deleted_count + recreated_count} (Deleted: {deleted_count}, Recreated: {recreated_count}, Skipped: {skipped_count}) | FAILURE: {len(error_log)}")
+    log_message(f"LATEST TIMESTAMP: {max_update_datetime}")
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file