wx-agent/workflows/workflow_utils.py

"""
工作流公共工具函数
提取自 workflow_fault_diagnosis.py 和 workflow_operate.py 的共享代码
"""

from typing import TypedDict, Optional, Dict, Any, List
from prompts import COMMON_PROMPTS
from modelsAPI.model_api import OpenaiAPI
from utils.function_tracker import get_all_callbacks, get_current_stream_handler
from utils.image_utils import extract_image_paths, normalize_markdown_images
import json
import re
import random


def safe_json_extract(text: str) -> Optional[Any]:
    if not text:
        return None
    cleaned = text.strip()
    cleaned = re.sub(r"```json\s*", "", cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r"```\s*", "", cleaned)
    m = re.search(r"\[[\s\S]*\]", cleaned)
    if m:
        try:
            return json.loads(m.group(0))
        except Exception:
            pass
    m = re.search(r"\{[\s\S]*\}", cleaned)
    if m:
        try:
            return json.loads(m.group(0))
        except Exception:
            pass
    try:
        return json.loads(cleaned)
    except Exception:
        return None


from workflows.history_manager import parse_history, build_history_str, filter_image_urls, init_history  # noqa: F401


def normalize_latex_spacing(text: str) -> str:
    placeholder = "\x00DOUBLE_DOLLAR\x00"
    text = text.replace("$$", placeholder)
    text = re.sub(r'(?<!\s)\$', r' $', text)
    text = re.sub(r'\$(?!\s)', r'$ ', text)
    text = text.replace(placeholder, "$$")
    text = re.sub(r'(?<!\s)\$\$', r' $$', text)
    text = re.sub(r'\$\$(?!\s)', r'$$ ', text)
    return text


def remove_duplicate_content(text: str) -> str:
    if not text or len(text) < 10:
        return text

    lines = text.split('\n')
    result = []
    seen_segments = set()

    for line in lines:
        line = line.rstrip()
        if not line:
            result.append(line)
            continue

        normalized_line = line.strip().lower()
        if normalized_line in seen_segments:
            continue

        seen_segments.add(normalized_line)
        result.append(line)

    cleaned_text = '\n'.join(result)
    return cleaned_text


def convert_rag_result(rag_result):
    if not rag_result or not rag_result.get("success", False):
        return []
    source_citation = rag_result.get("sourceCitation", {})
    flat_results = []
    for resource, items in source_citation.items():
        for item in items:
            flat_results.append(item)
    return flat_results


def calculate_info_completeness(weights: Dict[str, float], values: Dict[str, str]) -> float:
    score = 0.0
    for field, weight in weights.items():
        value = values.get(field, "") or ""
        if value and str(value).strip():
            score += weight
    return round(score, 2)


def emit_callback_event(title_options: List[str], details: str):
    start_event = {"type": "function_execution", "title": random.choice(title_options), "details": details}
    for callback in get_all_callbacks():
        try:
            callback(start_event)
        except Exception:
            pass


def build_detail_info(fields: Dict[str, str]) -> str:
    detail_info = ""
    for label, value in fields.items():
        if value:
            detail_info += f"- {label}：{value}\n"
    return detail_info


def append_atlas_section(answer: str, atlas_results: Dict[str, Any]) -> str:
    if not atlas_results:
        return answer
    atlas_section = "\n\n---\n\n**参考图册：**\n"
    for node_name, node_data in atlas_results.items():
        if isinstance(node_data, dict):
            for title, urls in node_data.items():
                if isinstance(urls, list) and urls:
                    for url in urls:
                        atlas_section += f"- {title}: [ `{title}` ]({url})\n"
    return answer + atlas_section




async def stream_format_and_postprocess(
    raw_content: str,
    original_image_paths: List[str],
    temperature: float = 0.3,
    content_label: str = None,
) -> str:
    if content_label:
        format_user_content = COMMON_PROMPTS["format_user_simple"].format(
            content_label=content_label,
            raw_content=raw_content,
        )
    else:
        format_user_content = COMMON_PROMPTS["format_user_with_image_examples"].format(raw_content=raw_content)

    format_system_prompt = COMMON_PROMPTS["format_system"]

    stream_handler = get_current_stream_handler()
    answer = ""
    accumulated_content = ""

    async for chunk in OpenaiAPI.open_api_chat_stream(
        model=None,
        system_prompt=format_system_prompt,
        messages=[{"role": "user", "content": format_user_content}],
        temperature=temperature,
    ):
        answer += chunk
        accumulated_content += chunk
        if stream_handler:
            stream_handler.send_stream_content(accumulated_content)

    answer = answer.strip() if answer else raw_content
    answer = remove_duplicate_content(answer)
    answer = normalize_markdown_images(answer, original_paths=original_image_paths)
    answer = normalize_latex_spacing(answer)

    return answer