sim_core_wrapper.py

import os, re, json, time, uuid
from dataclasses import dataclass, asdict
from typing import List, Dict

# Don't use a hardcoded fallback - if no API key is set, use demo mode
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY") or os.getenv("GENAI_API_KEY")

def ask_real(prompt: str) -> str:
    """Use Google AI REST API instead of google-genai SDK to avoid websockets dependency"""
    if not GOOGLE_API_KEY:
        raise RuntimeError("Real LLM unavailable - missing GOOGLE_API_KEY.")
    
    import requests
    # Use v1beta API with gemini-2.5-flash model
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key={GOOGLE_API_KEY}"
    
    payload = {
        "contents": [{
            "parts": [{"text": prompt}]
        }]
    }
    
    try:
        response = requests.post(url, json=payload, timeout=30)
        response.raise_for_status()
        data = response.json()
        
        # Extract text from response
        if "candidates" in data and len(data["candidates"]) > 0:
            candidate = data["candidates"][0]
            if "content" in candidate and "parts" in candidate["content"]:
                parts = candidate["content"]["parts"]
                if len(parts) > 0 and "text" in parts[0]:
                    return parts[0]["text"].strip()
        
        return "(no response)"
    except Exception as e:
        raise RuntimeError(f"API call failed: {e}")

def _stub_summary(text: str) -> str:
    return """- Summary: A public incident prompts officials, executives, and advocates to respond.\n" 
           "- Impact: Affected users demand remedies, oversight considers penalties.\n" 
           "- Company: Leadership apologizes and outlines next steps.\n" 
           "- Regulator: Signals investigation and potential sanctions.\n" 
           "- Outlook: Stakeholders coordinate on mitigation and communications."""

def _stub_extract(tldr: str, text: str) -> list:
    return [
        {"name": "Dana Kim", "role": "CEO of Acme Health", "relevance": "Wants to contain damage and reassure customers.", "quote": "We apologize and are notifying all affected users."},
        {"name": "Luis Ortega", "role": "State Attorney General", "relevance": "Wants accountability and compliance.", "quote": "Fines are possible if negligence is found."},
        {"name": "Pat Rivera", "role": "Patient Advocate", "relevance": "Wants immediate consumer protections.", "quote": "Offer credit monitoring now."},
        {"name": "Casey Tran", "role": "Cybersecurity Consultant", "relevance": "Wants root cause identified and fixed.", "quote": "Likely a misconfigured cloud bucket."}
    ]

def _stub_turn(name: str, role: str, world_facts: List[str], memory: List[str]) -> str:
    if name == "Dana Kim":
        return "We will notify all impacted users and rotate credentials. <action>Announce 24-month free credit monitoring</action>"
    if name == "Luis Ortega":
        return "My office will review compliance controls and require a corrective plan. <action>Open a formal investigation</action>"
    if name == "Pat Rivera":
        return "Consumers deserve rapid help and transparency. <action>Publish hotline and guidance for affected patients</action>"
    if name == "Casey Tran":
        return "Patch the misconfiguration, enable logging, and run tabletop drills. <action>Deliver a 48-hour remediation report</action>"
    return "Noted."

def _stub_reflect(last: str, others_text: str) -> dict:
    return {
        "knowledge_additions": ["Stakeholders coordinated first actions."],
        "stance_updates": {},
        "emotions": ["focused"],
        "objectives_update": [],
        "summary_memory": "Committed initial steps and oversight review."
    }

def ask(prompt: str) -> str:
    if GOOGLE_API_KEY:
        try:
            return ask_real(prompt)
        except Exception:
            pass
    if "Summarize the following news event" in prompt or "总结以下新闻事件" in prompt:
        return _stub_summary(prompt)
    if "You are an extraction system" in prompt or "你是一个提取系统" in prompt:
        return json.dumps({"actors": _stub_extract("", "")})
    if prompt.strip().startswith("System:") or prompt.strip().startswith("系统:"):
        import re
        m = re.search(r"You are ([^,]+), ([^\n]+).", prompt)
        if not m:
             m = re.search(r"你是 ([^,]+), ([^\n]+).", prompt)
        name = m.group(1) if m else "Actor"
        role = m.group(2) if m else "Role"
        return _stub_turn(name, role, [], [])
    if "You are reflecting as" in prompt or "你正在作为" in prompt:
        return json.dumps(_stub_reflect("", ""))
    return "Acknowledged."

def llm_json(prompt: str) -> dict:
    resp = ask(prompt)
    import re, json
    m = re.search(r'{.*}', resp, flags=re.S)
    text = m.group(0) if m else resp
    try:
        return json.loads(text)
    except Exception:
        return {"raw": resp}

# --- PROMPTS ---

PROMPTS = {
    "en": {
        "summarize": "Summarize the following news event in 5 bullet points, factual and neutral:\n\n{text}\n",
        "extract": """You are an extraction system.
Article TLDR:
{article_tldr}

Full text:
{article_text}

Return JSON with a list 'actors' of up to 5 items. Each item:
- name
- role
- relevance
- any quoted line from them if present
""",
        "persona": """Build a compact personality card for this actor as JSON.

Actor:
{actor_json}

Event TLDR:
{article_tldr}

Return JSON with:
- traits: list of 5 short adjectives
- objectives: list of 3 concise goals
- red_lines: list of 3 concise constraints
- knowledge: list of 4 short facts they believe
- stance: object with up to 3 key topics -> single sentence stance each
""",
        "turn": """System:
You are {name}, {role}.
Traits: {traits}
Objectives: {objectives}
Red lines: {red_lines}

Shared facts so far:
- {world_facts}

Recent conversation:
{recent_conversation}

Your private memory:
- {private_memory}

Speak concisely to advance your objectives.
If you commit to a concrete step, announce exactly one action as:
<action>your concise action line</action>
Do not invent facts not in the article or shared facts.
""",
        "reflect": """You are reflecting as {name}.
Given your last message:
{last_message}

And others' replies:
{others_replies}

Update in JSON:
- knowledge_additions: list of new factual nuggets you accept
- stance_updates: topic -> one sentence if changed
- emotions: up to 2 words
- objectives_update: if any
- summary_memory: one line summary of what you learned
"""
    },
    "zh": {
        "summarize": "请用5个要点总结以下新闻事件，保持客观中立：\n\n{text}\n",
        "extract": """你是一个提取系统。
文章摘要：
{article_tldr}

全文：
{article_text}

请返回包含最多5个项目的列表 'actors' 的JSON。每个项目包含：
- name (姓名)
- role (角色)
- relevance (相关性)
- any quoted line from them if present (引用的原话)
""",
        "persona": """请为该角色构建一个简洁的个性卡片（JSON格式）。

角色：
{actor_json}

事件摘要：
{article_tldr}

返回 JSON 包含：
- traits: 5个简短的形容词列表
- objectives: 3个简洁的目标列表
- red_lines: 3个简洁的限制列表
- knowledge: 4个他们相信的简短事实列表
- stance: 对象，包含最多3个关键话题 -> 每个话题一句话的立场
""",
        "turn": """系统:
你是 {name}, {role}。
特征: {traits}
目标: {objectives}
底线: {red_lines}

目前共享的事实:
- {world_facts}

最近的对话:
{recent_conversation}

你的私人记忆:
- {private_memory}

请简洁发言以推进你的目标。
如果你承诺采取具体步骤，请确切宣布一个行动，格式为：
<action>你的简洁行动行</action>
不要捏造文章或共享事实中不存在的事实。
""",
        "reflect": """你正在作为 {name} 进行反思。
鉴于你的上一条消息：
{last_message}

以及其他人的回复：
{others_replies}

请用 JSON 更新：
- knowledge_additions: 你接受的新事实信息列表
- stance_updates: 话题 -> 如果改变，用一句话描述
- emotions: 最多2个词
- objectives_update: 如果有更新
- summary_memory: 一句话总结你学到的内容
"""
    }
}

def summarize_article(text: str, lang: str = "en") -> str:
    prompt = PROMPTS[lang]["summarize"].format(text=text)
    return ask(prompt).strip()

def extract_actors(article_tldr: str, article_text: str, lang: str = "en") -> list:
    prompt = PROMPTS[lang]["extract"].format(article_tldr=article_tldr, article_text=article_text)
    data = llm_json(prompt)
    return data.get("actors", [])

def bootstrap_persona(actor: dict, article_tldr: str, lang: str = "en") -> dict:
    prompt = PROMPTS[lang]["persona"].format(actor_json=json.dumps(actor, ensure_ascii=False), article_tldr=article_tldr)
    data = llm_json(prompt)
    if "traits" not in data:
        data = {
            "traits": ["measured","curious","pragmatic","analytical","candid"],
            "objectives": ["seek clarity","protect interests","act credibly"],
            "red_lines": ["avoid speculation","no illegal acts","respect privacy"],
            "knowledge": [],
            "stance": {}
        }
    return data

from dataclasses import dataclass
@dataclass
class Personality:
    name: str
    role: str
    traits: list
    objectives: list
    red_lines: list
    knowledge: list
    stance: dict

class CharacterState:
    def __init__(self, persona: Personality):
        self.id = str(uuid.uuid4())[:8]
        self.card = persona
        self.memory_log: List[str] = []
        self.has_acted: bool = False
        self.actions: List[str] = []

class Conversation:
    def __init__(self, article, lang: str = "en"):
        self.article = article
        self.lang = lang
        self.characters: Dict[str, CharacterState] = {}
        self.transcript: List[Dict] = []
        self.world_facts: List[str] = [
            f"Title: {article['title']}",
            f"Source: {article['source']}",
            article['tldr']
        ]
        self.debug: List[str] = []

    def add_character(self, persona: Personality):
        self.characters[persona.name] = CharacterState(persona)

    def everyone_acted(self) -> bool:
        return all(ch.has_acted for ch in self.characters.values()) and len(self.characters) > 0

    def next_speaker(self) -> str:
        order = sorted(self.characters.keys(), key=lambda n: self.characters[n].has_acted)
        if not self.transcript:
            return order[0]
        last = self.transcript[-1]['speaker']
        idx = (order.index(last) + 1) % len(order)
        return order[idx]

    def step(self):
        import re
        if not self.characters:
            self.debug.append("No characters present, skipping step.")
            return

        speaker_name = self.next_speaker()
        ch = self.characters[speaker_name]
        self.debug.append(f"Step: next speaker -> {speaker_name}")

        tail_text = "\n".join([f"{t['speaker']}: {t['text']}" for t in self.transcript[-6:]])
        p = ch.card

        prompt = PROMPTS[self.lang]["turn"].format(
            name=p.name,
            role=p.role,
            traits=', '.join(p.traits),
            objectives=', '.join(p.objectives),
            red_lines=', '.join(p.red_lines),
            world_facts="\n- ".join(self.world_facts[-8:]),
            recent_conversation=tail_text,
            private_memory="\n- ".join(ch.memory_log[-5:])
        )

        reply = ask(prompt).strip()
        self.transcript.append({"speaker": speaker_name, "text": reply})
        self.debug.append(f"Received reply from {speaker_name} ({len(reply)} chars).")

        acts = re.findall(r"<action>(.*?)</action>", reply, flags=re.S)
        if acts:
            ch.has_acted = True
            ch.actions.extend([a.strip() for a in acts])
            last_action = acts[-1].strip()
            self.world_facts.append(f"{speaker_name} committed: {last_action}")
            self.debug.append(f"Action parsed for {speaker_name}: {last_action}")

        others = [t for t in self.transcript[-3:] if t['speaker'] != speaker_name]
        others_text = "\n".join([f"{o['speaker']}: {o['text']}" for o in others])

        reflect_prompt = PROMPTS[self.lang]["reflect"].format(
            name=ch.card.name,
            last_message=reply,
            others_replies=others_text
        )

        reflection = llm_json(reflect_prompt)
        if isinstance(reflection, dict):
            mem = reflection.get("summary_memory")
            if mem:
                ch.memory_log.append(mem)
                self.debug.append(f"Memory appended for {speaker_name}: {mem}")
            # Ensure knowledge_additions is a list
            knowledge_additions = reflection.get("knowledge_additions")
            if isinstance(knowledge_additions, list):
                for nugget in knowledge_additions:
                    self.world_facts.append(nugget)
                    self.debug.append(f"World fact added: {nugget}")
            # Ensure stance_updates is a dict
            stance_updates = reflection.get("stance_updates")
            if isinstance(stance_updates, dict):
                for k, v in stance_updates.items():
                    ch.card.stance[k] = v
                    self.debug.append(f"Stance update for {speaker_name} [{k}]: {v}")

    def run(self, max_rounds: int = 24):
        self.debug.append(f"Run start: max_rounds={max_rounds}, characters={list(self.characters.keys())}")
        for r in range(max_rounds):
            if self.everyone_acted():
                self.debug.append(f"All actors have taken an action by round {r}. Stopping.")
                break
            self.debug.append(f"Round {r + 1} begin.")
            self.step()
            self.debug.append(f"Round {r + 1} end. Transcript len={len(self.transcript)}.")
        else:
            self.debug.append(f"Hit max_rounds={max_rounds} without all actions taken. Stopping.")

    def report(self) -> Dict:
        return {
            "world_facts": self.world_facts,
            "characters": {
                name: {
                    "role": ch.card.role,
                    "traits": ch.card.traits,
                    "objectives": ch.card.objectives,
                    "stance": ch.card.stance,
                    "actions": ch.actions
                } for name, ch in self.characters.items()
            },
            "transcript": self.transcript,
            "debug": self.debug,
        }

def run_simulation_from_text(title: str, source: str, text: str, url: str = None, date: str = "", language: str = "en", additional_context: str = "") -> Dict:
    tldr = summarize_article(text, lang=language)
    article = {"title": title, "date": date, "source": source, "url": url, "text": text, "tldr": tldr}

    actors = extract_actors(tldr, text, lang=language)
    convo = Conversation(article, lang=language)

    for a in actors:
        card = bootstrap_persona(a, tldr, lang=language)
        persona = Personality(
            name=a.get("name", "Unknown"),
            role=a.get("role", "Actor"),
            traits=card.get("traits", ["measured","curious","pragmatic","analytical","candid"]),
            objectives=card.get("objectives", ["seek clarity","protect interests","act credibly"]),
            red_lines=card.get("red_lines", ["avoid speculation","no illegal acts","respect privacy"]),
            knowledge=card.get("knowledge", []),
            stance=card.get("stance", {})
        )
        convo.add_character(persona)

    for a in actors[:3]:
        if a.get("relevance"):
            convo.world_facts.append(f"{a['name']}: {a['relevance']}")
    
    # Add additional context if provided
    if additional_context and additional_context.strip():
        convo.world_facts.append(f"Additional context: {additional_context.strip()}")

    convo.run(max_rounds=24)
    return convo.report()