smart_writer_kit/agent_for_outline_update.py

import gradio as gr
import pandas as pd
import json
from model_handler import ModelHandler
from config import LING_FLASH_2_0

def _format_outline_for_prompt(df: pd.DataFrame) -> str:
    """Formats the outline DataFrame into a simple numbered list for the prompt."""
    if df is None or df.empty:
        return "无任务。"
    
    tasks = [f"{i+1}. {row['Task']}" for i, row in df.iterrows()]
    return "\n".join(tasks)

def update_outline_status_agent(short_outline_df: pd.DataFrame, editor_content: str):
    """
    Agent to analyze text and update the outline's completion status using a real LLM.
    """

    if editor_content is None or len(editor_content.strip()) < 20:
        return short_outline_df # Return original df

    try:
        # 1. Prepare Prompts
        system_prompt = (
            "你是一个任务分析机器人。请仔细阅读用户提供的'已完成大纲'和'当前文本'，判断大纲中的每项任务是否已经在文本中被完成。\n"
            "你的回答必须是一个遵循以下规则的 JSON 对象：\n"
            "1. JSON 的 key 是大纲中的任务原文。\n"
            "2. JSON 的 value 是一个布尔值 (`true` 或 `false`)，`true` 代表任务已完成，`false` 代表未完成。\n"
            "3. 不要返回除了这个 JSON 对象之外的任何其他文本、解释或代码块标记。"
        )
        
        outline_str = _format_outline_for_prompt(short_outline_df)
        user_prompt = (
            f"### 已有大纲\n{outline_str}\n\n"
            f"### 当前文本\n{editor_content[-4000:]}\n\n"
            "### 指令\n请根据上述'当前文本'，分析'已有大纲'中的任务完成情况，并返回 JSON 对象。"
        )

        # 2. Call LLM
        model_handler = ModelHandler()
        response_generator = model_handler.generate_code(
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            model_choice=LING_FLASH_2_0
        )
        full_response = "".join(chunk for chunk in response_generator)
        
        # 3. Parse JSON and Update DataFrame
        print("【收到的完整上下文】")
        print("full_response:", repr(full_response))
        
        # Clean up potential markdown code block
        if full_response.strip().startswith("```json"):
            full_response = full_response.strip()[7:-3].strip()

        completion_status = json.loads(full_response)
        
        # Create a copy to avoid modifying the original df in place
        updated_df = short_outline_df.copy()
        
        for i, row in updated_df.iterrows():
            task_text = row['Task']
            if task_text in completion_status:
                updated_df.at[i, 'Done'] = bool(completion_status[task_text])

        print("【收到的完整上下文】")
        print("updated_df:\n", updated_df.to_string())
        return updated_df

    except json.JSONDecodeError:
        print(f"[Agent] Error: Failed to decode JSON from LLM response: {full_response}")
        # On JSON error, we don't want to change anything.
        return short_outline_df
    except Exception as e:
        print(f"[Agent] Error updating outline status: {e}")
        # On other errors, also return the original dataframe.
        return short_outline_df