124 lines
5.1 KiB
Python
124 lines
5.1 KiB
Python
import json
|
|
import os
|
|
from core import utils
|
|
from ai import models as ai_models
|
|
|
|
|
|
def _empty_state():
|
|
return {"active_threads": [], "immediate_handoff": "", "resolved_threads": [], "chapter": 0}
|
|
|
|
|
|
def load_story_state(folder, project_id=None):
|
|
"""Load structured story state from DB (if project_id given) or story_state.json fallback."""
|
|
if project_id is not None:
|
|
try:
|
|
from web.db import StoryState
|
|
record = StoryState.query.filter_by(project_id=project_id).first()
|
|
if record and record.state_json:
|
|
return json.loads(record.state_json) or _empty_state()
|
|
except Exception:
|
|
pass # Fall through to file-based load if DB unavailable (e.g. CLI context)
|
|
|
|
path = os.path.join(folder, "story_state.json")
|
|
if os.path.exists(path):
|
|
return utils.load_json(path) or _empty_state()
|
|
return _empty_state()
|
|
|
|
|
|
def update_story_state(chapter_text, chapter_num, current_state, folder, project_id=None):
|
|
"""Use model_logic to extract structured story threads from the new chapter
|
|
and save the updated state to the StoryState DB table and/or story_state.json.
|
|
Returns the new state."""
|
|
utils.log("STATE", f"Updating story state after Ch {chapter_num}...")
|
|
prompt = f"""
|
|
ROLE: Story State Tracker
|
|
TASK: Update the structured story state based on the new chapter.
|
|
|
|
CURRENT_STATE:
|
|
{json.dumps(current_state)}
|
|
|
|
NEW_CHAPTER (Chapter {chapter_num}):
|
|
{utils.truncate_to_tokens(chapter_text, 4000)}
|
|
|
|
INSTRUCTIONS:
|
|
1. ACTIVE_THREADS: 2-5 concise strings, each describing what a key character is currently trying to achieve.
|
|
- Carry forward unresolved threads from CURRENT_STATE.
|
|
- Add new threads introduced in this chapter.
|
|
- Remove threads that are now resolved.
|
|
2. IMMEDIATE_HANDOFF: Write exactly 3 sentences describing how this chapter ended:
|
|
- Sentence 1: Where are the key characters physically right now?
|
|
- Sentence 2: What emotional state are they in at the very end of this chapter?
|
|
- Sentence 3: What immediate unresolved threat, question, or decision is hanging in the air?
|
|
3. RESOLVED_THREADS: Carry forward from CURRENT_STATE + add threads explicitly resolved in this chapter.
|
|
|
|
OUTPUT_FORMAT (JSON):
|
|
{{
|
|
"active_threads": ["Thread 1", "Thread 2"],
|
|
"immediate_handoff": "Sentence 1. Sentence 2. Sentence 3.",
|
|
"resolved_threads": ["Resolved thread 1"],
|
|
"chapter": {chapter_num}
|
|
}}
|
|
"""
|
|
try:
|
|
response = ai_models.model_logic.generate_content(prompt)
|
|
utils.log_usage(folder, ai_models.model_logic.name, response.usage_metadata)
|
|
new_state = json.loads(utils.clean_json(response.text))
|
|
new_state['chapter'] = chapter_num
|
|
|
|
# Write to DB if project_id is available
|
|
if project_id is not None:
|
|
try:
|
|
from web.db import db, StoryState
|
|
from datetime import datetime
|
|
record = StoryState.query.filter_by(project_id=project_id).first()
|
|
if record:
|
|
record.state_json = json.dumps(new_state)
|
|
record.updated_at = datetime.utcnow()
|
|
else:
|
|
record = StoryState(project_id=project_id, state_json=json.dumps(new_state))
|
|
db.session.add(record)
|
|
db.session.commit()
|
|
except Exception as db_err:
|
|
utils.log("STATE", f" -> DB write failed: {db_err}. Falling back to file.")
|
|
|
|
# Always write to file for backward compat with CLI
|
|
path = os.path.join(folder, "story_state.json")
|
|
with open(path, 'w') as f:
|
|
json.dump(new_state, f, indent=2)
|
|
|
|
utils.log("STATE", f" -> Story state saved. Active threads: {len(new_state.get('active_threads', []))}")
|
|
return new_state
|
|
except Exception as e:
|
|
utils.log("STATE", f" -> Story state update failed: {e}. Keeping previous state.")
|
|
return current_state
|
|
|
|
|
|
def format_for_prompt(state, chapter_beats=None):
|
|
"""Format the story state into a prompt-ready string.
|
|
Active threads and immediate handoff are always included.
|
|
Resolved threads are only included if referenced in the chapter's beats."""
|
|
if not state or (not state.get('immediate_handoff') and not state.get('active_threads')):
|
|
return None
|
|
|
|
beats_text = " ".join(str(b) for b in (chapter_beats or [])).lower()
|
|
lines = []
|
|
|
|
if state.get('immediate_handoff'):
|
|
lines.append(f"IMMEDIATE STORY HANDOFF (exactly how the previous chapter ended):\n{state['immediate_handoff']}")
|
|
|
|
if state.get('active_threads'):
|
|
lines.append("ACTIVE PLOT THREADS:")
|
|
for t in state['active_threads']:
|
|
lines.append(f" - {t}")
|
|
|
|
relevant_resolved = [
|
|
t for t in state.get('resolved_threads', [])
|
|
if any(w in beats_text for w in t.lower().split() if len(w) > 4)
|
|
]
|
|
if relevant_resolved:
|
|
lines.append("RESOLVED THREADS (context only — do not re-introduce):")
|
|
for t in relevant_resolved:
|
|
lines.append(f" - {t}")
|
|
|
|
return "\n".join(lines)
|