import gradio as gr
import time
from smart_writer_kit.agent_for_streaming_completion import fetch_flow_suggestion_agent, accept_flow_suggestion_agent
from smart_writer_kit.agent_for_inspiration_expansion import fetch_inspiration_agent, apply_inspiration_agent
from smart_writer_kit.agent_for_outline_update import update_outline_status_agent
from smart_writer_kit.agent_for_kb_update import suggest_new_kb_terms_agent
# --- Mock Data (for UI population only) ---
MOCK_STYLE = """风格：赛博朋克 / 黑色电影
视角：第三人称限制视角（主角：凯）
基调：阴郁、压抑、霓虹闪烁的高科技低生活
核心规则：
1. 强调感官描写，特别是光影和声音。
2. 避免过多的心理独白，通过行动展现心理。
"""
MOCK_KNOWLEDGE_BASE = [
["凯 (Kai)", "主角，前黑客，现在是义体医生。左臂是老式的军用义体。"],
["夜之城 (Night City)", "故事发生的舞台，一座永夜的巨型都市，被企业掌控。"],
["荒坂塔 (Arasaka Tower)", "市中心的最高建筑，象征着绝对的权力。"],
["赛博精神病 (Cyberpsychosis)", "过度改装义体导致的解离性精神障碍。"],
["网络监察 (NetWatch)", "负责维护网络安全的组织，被黑客们视为走狗。"]
]
MOCK_SHORT_TERM_OUTLINE = [
[True, "凯接到一个神秘电话，对方声称知道他失踪妹妹的下落。"],
[False, "凯前往'来生'酒吧与接头人见面。"],
[False, "在酒吧遇到旧识，引发一场关于过去的争执。"],
[False, "接头人出现，但似乎被跟踪了。"]
]
MOCK_LONG_TERM_OUTLINE = [
[False, "揭露夜之城背后的惊天阴谋。"],
[False, "凯找回妹妹，或者接受她已经改变的事实。"],
[False, "与荒坂公司的最终决战。"]
]
# --- UI Helper Functions ---
def get_stats(text):
"""Calculate word count and read time."""
if not text:
return "0 Words | 0 mins"
words = len(text.split())
read_time = max(1, words // 200) # Average reading speed
return f"{words} Words | ~{read_time} mins"
def dismiss_inspiration():
return gr.update(visible=False)
# --- UI Construction ---
def create_smart_writer_tab():
debounce_state = gr.State({"last_change": 0, "active": False, "style": "", "kb": [], "short_outline": [], "long_outline": []})
debounce_timer = gr.Timer(0.5, active=False)
with gr.Row(equal_height=False, elem_id="indicator-writing-tab"):
# --- Left Column: Entity Console ---
with gr.Column(scale=0, min_width=384) as left_panel:
gr.Markdown("### 🧠 核心实体控制台")
with gr.Accordion("整体章程 (Style)", open=True):
style_input = gr.Textbox(
label="整体章程",
lines=8,
value=MOCK_STYLE,
interactive=True
)
with gr.Accordion("知识库 (Knowledge Base)", open=True):
kb_input = gr.Dataframe(
headers=["Term", "Description"],
datatype=["str", "str"],
value=MOCK_KNOWLEDGE_BASE,
interactive=True,
label="知识库",
wrap=True
)
with gr.Row():
btn_suggest_kb = gr.Button("🔍 提取新词条", size="sm")
md_suggested_terms_header = gr.Markdown("#### 推荐词条", visible=False) # Placeholder for suggested terms
suggested_kb_dataframe = gr.Dataframe(
headers=["Term", "Description"],
datatype=["str", "str"],
visible=False, # Initially hidden
interactive=False,
label="推荐词条"
)
with gr.Accordion("当前章节大纲 (Short-Term)", open=True):
short_outline_input = gr.Dataframe(
headers=["Done", "Task"],
datatype=["bool", "str"],
value=MOCK_SHORT_TERM_OUTLINE,
interactive=True,
label="当前章节大纲",
col_count=(2, "fixed"),
)
with gr.Row():
btn_sync_outline = gr.Button("🔄 同步状态", size="sm")
with gr.Accordion("故事总纲 (Long-Term)", open=False):
long_outline_input = gr.Dataframe(
headers=["Done", "Task"],
datatype=["bool", "str"],
value=MOCK_LONG_TERM_OUTLINE,
interactive=True,
label="故事总纲",
col_count=(2, "fixed"),
)
# --- Right Column: Writing Canvas ---
with gr.Column(scale=1) as right_panel:
# Toolbar
with gr.Row(elem_classes=["toolbar"]):
stats_display = gr.Markdown("0 Words | 0 mins")
inspiration_btn = gr.Button("✨ 继续整段 (Cmd/Ctrl+Enter)", size="sm", variant="primary", elem_id="btn-action-create-paragraph")
# 主要编辑器区域
editor = gr.Textbox(
label="沉浸写作画布",
placeholder="开始你的创作...",
lines=30,
elem_classes=["writing-editor"],
elem_id="writing-editor",
show_label=False,
)
# Flow Suggestion
with gr.Row(variant="panel"):
flow_suggestion_display = gr.Textbox(
label="AI 实时续写建议 (按 Tab 采纳)",
value="（等待输入...）",
interactive=False,
scale=4,
elem_classes=["flow-suggestion-box"]
)
accept_flow_btn = gr.Button("采纳(Tab)", scale=1, elem_id='btn-action-accept-flow')
refresh_flow_btn = gr.Button("换一个(Shift+Tab)", scale=1, elem_id='btn-action-change-flow')
# Debounce Progress
debounce_progress = gr.HTML(value="", visible=False)
# Inspiration Modal
with gr.Group(visible=False) as inspiration_modal:
gr.Markdown("### 💡 灵感选项 (由 Ling 模型生成)")
inspiration_prompt_input = gr.Textbox(
label="设定脉络 (可选)",
placeholder="例如：写一段激烈的打斗 / 描写赛博朋克夜景...",
lines=1
)
refresh_inspiration_btn = gr.Button("生成选项(Shift+Enter)")
with gr.Row():
opt1_btn = gr.Button("...", elem_classes=["inspiration-card"])
opt2_btn = gr.Button("...", elem_classes=["inspiration-card"])
opt3_btn = gr.Button("...", elem_classes=["inspiration-card"])
cancel_insp_btn = gr.Button("取消")
# --- Interactions ---
# 1. Stats
editor.change(fn=get_stats, inputs=editor, outputs=stats_display)
# 2. Inspiration Workflow
# Open Modal (triggered by visible button or hidden trigger button for Cmd+Enter)
open_inspiration_modal_fn = lambda: (gr.update(visible=True), "")
inspiration_btn.click(fn=open_inspiration_modal_fn, outputs=[inspiration_modal, inspiration_prompt_input])
# Generate Options based on Prompt
refresh_inspiration_btn.click(
fn=fetch_inspiration_agent,
inputs=[inspiration_prompt_input, editor, style_input, kb_input, short_outline_input, long_outline_input],
outputs=[inspiration_modal, opt1_btn, opt2_btn, opt3_btn]
)
# Apply Option
for btn in [opt1_btn, opt2_btn, opt3_btn]:
btn.click(
fn=apply_inspiration_agent,
inputs=[editor, btn],
outputs=[editor, inspiration_modal, inspiration_prompt_input]
)
cancel_insp_btn.click(fn=dismiss_inspiration, outputs=inspiration_modal, show_progress="hidden")
# 3. Flow Suggestion with Debounce
def start_debounce(editor_content, style, kb, short_outline, long_outline):
return {"last_change": time.time(), "active": True, "style": style, "kb": kb, "short_outline": short_outline, "long_outline": long_outline}, gr.update(active=True), gr.update(visible=True, value=" 补全中... 3.0s")
def update_debounce(debounce_state, editor_content):
if not debounce_state["active"]:
return gr.update(), gr.update(), debounce_state, gr.update()
elapsed = time.time() - debounce_state["last_change"]
if elapsed >= 3:
suggestion = fetch_flow_suggestion_agent(editor_content, debounce_state["style"], debounce_state["kb"], debounce_state["short_outline"], debounce_state["long_outline"])
return gr.update(visible=False), suggestion, {"last_change": 0, "active": False, "style": "", "kb": [], "short_outline": [], "long_outline": []}, gr.update(active=False)
else:
progress = int((elapsed / 3) * 100)
remaining = 3 - elapsed
progress_html = f" 补全中... {remaining:.1f}s"
return gr.update(value=progress_html), gr.update(), debounce_state, gr.update()
editor.change(fn=start_debounce, inputs=[editor, style_input, kb_input, short_outline_input, long_outline_input], outputs=[debounce_state, debounce_timer, debounce_progress])
debounce_timer.tick(fn=update_debounce, inputs=[debounce_state, editor], outputs=[debounce_progress, flow_suggestion_display, debounce_state, debounce_timer])
refresh_flow_btn.click(fn=fetch_flow_suggestion_agent, inputs=[editor, style_input, kb_input, short_outline_input, long_outline_input], outputs=flow_suggestion_display)
# Accept Flow (Triggered by visible Button or hidden Tab Key trigger)
accept_flow_fn_inputs = [editor, flow_suggestion_display]
accept_flow_fn_outputs = [editor]
accept_flow_btn.click(fn=accept_flow_suggestion_agent, inputs=accept_flow_fn_inputs, outputs=accept_flow_fn_outputs, show_progress="hidden")
# 4. Agent-based Context Updates
btn_sync_outline.click(
fn=update_outline_status_agent,
inputs=[short_outline_input, editor],
outputs=[short_outline_input]
)
btn_suggest_kb.click(
fn=suggest_new_kb_terms_agent,
inputs=[kb_input, editor],
outputs=[suggested_kb_dataframe, md_suggested_terms_header]
)