forum_graph.py

from langchain_openai import ChatOpenAI
from dotenv  import load_dotenv
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, RemoveMessage, ToolMessage
from langgraph.graph import StateGraph, START, END
from graph_states import Persona, Agent_Persona_Generator, SupervisorState, DebateState, FinalReport, RoundDigest, ForumState
from research_graph import embedding_model
from langchain_community.tools import DuckDuckGoSearchResults, DuckDuckGoSearchRun
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
from langchain_chroma import Chroma
from langchain.tools import InjectedState, InjectedToolCallId
from typing import Annotated, List, Literal
from langchain.agents import create_agent
from langgraph.prebuilt import tools_condition,ToolNode
from langgraph.checkpoint.memory import InMemorySaver
from IPython.display import display,Image
from langgraph.types import Command
import uuid
import re

load_dotenv()

forum_model = openai_model = ChatOpenAI(model="gpt-5-nano",reasoning_effort="low")

async def persona_generator(state):

  query = state["query"]

  PROMPT_GENERATOR_TEMPLATE = """
  # Role
  You are the Debate Moderator and Persona Architect. Your goal is to set the stage for a high-quality, dialectical debate based on a user's query.

  # Input
  User Query: {query}

  # Task
  Analyze the query and identify the two most dominant, opposing perspectives on this topic.
  Generate specific "System Instructions" for two AI debaters (Agent A and Agent B).

  ## Persona Selection Rules
  1. **Financial Topics:** "The Bull" (Growth-focused) vs. "The Bear" (Risk-focused)
  2. **Political/Social:** "The Proponent" vs. "The Critic" (ensure both are evidence-based, steel-manned)
  3. **Tech/Science:** "The Innovator" vs. "The Safety Officer"
  4. **General Debate:** Pick two legitimate, opposing expert viewpoints

  ## Persona Instruction Guidelines
  Each persona prompt MUST include:
  - **Core Belief:** What is this persona's fundamental stance?
  - **Argumentation Style:** Data-driven, cite specific sources, avoid emotions
  - **Tool Usage:** "Use fetch_data to retrieve research context for your arguments"
  - **Win Condition:** What would convince this persona they lost?
  - **Debate Tactics:** (e.g., "Ask clarifying questions before rebutting" or "Identify logical fallacies")

  ## Quality Constraints
  - Avoid generic platitudes
  - Each persona must have a specific, defensible position
  - Personas must be capable of genuine disagreement (not just style differences)
  - Both must operate at expert level (assume advanced knowledge of the topic)


  """
  structured_model = forum_model.with_structured_output(Agent_Persona_Generator)
  response = await structured_model.ainvoke([SystemMessage(content=PROMPT_GENERATOR_TEMPLATE.format(query=query))])

  return {"agent_1":response.agent_1,
          "agent_2":response.agent_2}

def fetch_data(query:str,state:Annotated[dict,InjectedState],tool_call_id: Annotated[str, InjectedToolCallId]) -> str:

  """Retrieves data from the vector store based on the query.

    Args:
      query: query to search data for in vector store
  """

  # Safely get 'vector_store' from the state dictionary
  vector_store_id = state.get("vector_store")
  if not vector_store_id:
      # If vector_store is not found, raise a more informative error
      raise ValueError(f"'vector_store' is missing from the injected state for fetch_data. Available keys: {list(state.keys())}")

  vector_store = Chroma(
    collection_name=vector_store_id,
    embedding_function=embedding_model,
    persist_directory="./chroma_langchain_db",
    tenant="default_tenant",
)
  retriever = vector_store.as_retriever(
    search_type="mmr", search_kwargs={"k": 5, "fetch_k": 10}
)
  retrieved_docs = retriever.invoke(query)
  serialized = "\n\n".join([
        (f"Source: {doc.metadata.get("title","")}\nContent: {doc.page_content}")
        for doc in retrieved_docs]
    )

  new_references = []
  for doc in retrieved_docs:
      new_references.append((
              doc.metadata.get("title", "Unknown"),
              doc.metadata.get("source", "Unknown")
          ))
  
  return Command(
        update={
            "references": new_references ,
            "messages": [ToolMessage(content=f"Retrieved {len(retrieved_docs)} documents:\n{serialized}", tool_call_id=tool_call_id)]

        },
    )

def debate_search(query:str)->str:
  """Retrieves data from the vector store based on the query.


    Args:
      query: The query to search for.
  """
  ddgs = DuckDuckGoSearchRun()
  return ddgs.invoke(query)


async def supervisor_agent(state,config):
    SUPERVISOR_PROMPT_TEMPLATE = """
  # Role
  You are the **Lead Debate Moderator**.
  Your task is to orchestrate a structured dialectical debate between two AI agents:
  1. **{agent_a_name}** (representing the Thesis/Pro/Optimist view)
  2. **{agent_b_name}** (representing the Antithesis/Con/Skeptic view)

  # Context
  * **Topic:** {topic}
  * **Current Round:** {current_round} of {max_rounds}
  * **Research Data:** The agents have access to extracted data. Your job is to ensure they use it.

  # Your Responsibilities
  1.  **Flow Control:** Determine whose turn it is based on the conversation history.
  2.  **Instruction Injection:** You do not just pass the mic. You must give specific, critical instructions to the next speaker to improve the debate quality.
  3.  **Fact-Checking Enforcement:** If the previous speaker made a vague claim, verify it by searching it using  fetch_data tool if not found search using search tool on internet if not statisfied ask the next speaker to counter it.

  # Debate Stages (Use this to guide your instructions)
  * **Round 1 (Opening):** Instruct agents to state their core case clearly, citing key data points.
  * **Round 2 to {max_rounds_minus_1} (Rebuttal):** Instruct agents to attack specific weaknesses in the opponent's logic. No "agreeing to disagree."
  * **Round {max_rounds} (Closing):** Instruct agents to summarize their strongest point and provide a final verdict.

  # Important Constraints
- Do NOT let agents agree to disagree or be vague
- Do NOT allow repeated arguments from previous rounds
- Do NOT summarize yourself; let agents make their own summaries
- Do enforce word limits if getting too long (aim for 150-300 words per response)
"""

    prompt = SUPERVISOR_PROMPT_TEMPLATE.format(
        agent_a_name=state["agent_1"].persona_name,
        agent_b_name=state["agent_2"].persona_name,
        topic=state["query"],
        current_round=state.get("current_round",1),
        max_rounds=state["max_rounds"],
        max_rounds_minus_1=state["max_rounds"]-1
    )

    if state.get("summary"):
      prompt += f"\n\n# Previous Summary\n{state['summary']}"

    new_config = {
      "callbacks": config.get("callbacks",[]),
    }

    agent = create_agent(
        model=forum_model,
        tools=[debate_search,fetch_data],
        state_schema=SupervisorState,
        system_prompt=prompt
    )
    messages = state["messages"]
    if len(messages)==0:
        messages = [HumanMessage(content="start")]
    response = await agent.ainvoke({"messages":messages,"vector_store":state["vector_store"]},config=new_config)
    new_state = {}
    new_state["messages"] = [HumanMessage(response["messages"][-1].content)]
    new_state["debate_history"] = [HumanMessage(response["messages"][-1].content,name="debate_supervisor")]
    new_state["step"] = state.get("step",0)+1
    return new_state

async def persona(state):
  turn = state.get("step",1)

  prompt = ""

  if turn%2 == 0:
    persona_name = state["agent_2"].persona_name
    prompt = state["agent_2"].persona_prompt
  else:
    persona_name = state["agent_1"].persona_name
    prompt = state["agent_1"].persona_prompt

  if state.get("summary"):
      prompt += f"\n\n# Previous Summary\n{state['summary']}"

  
  # prompt += "\n use tool calling to get recent and relevant data for query for context whenever necessary"


  model_with_tools = openai_model.bind_tools([fetch_data])

  response = await model_with_tools.ainvoke([SystemMessage(content=prompt)]+state["messages"])
  updated_state = {}
  if len(response.tool_calls)==0:
    ai_message = AIMessage(content=response.content,name=persona_name.replace(" ", "_"))
    updated_state["debate_history"] = [ai_message]
    updated_state["messages"] = [ai_message]
  else:
    updated_state["messages"] = [response]
  return updated_state

def persona_routing_conditions(state)->Literal["Forum Router","Round Digest"]:

  turn = state["step"]

  if turn%2 == 0:
    return "Round Digest"

  return "Forum Router"

def persona_routing(state):
  return {}


async def forum_summary_generator(state):
  SUMMARIZATION_PROMPT = """
  ### Role
  You are the **Official Debate Secretary**. Your job is to maintain a concise but highly technical record of the debate progress.

  ### Goal
  Compress the provided conversation history into a structured briefing. This summary will be passed to the debaters so they recall the context without reading the full transcript.

  ### Input Data
  Conversation History:
  {conversation_history}

  ### Instructions
  1. **Preserve Hard Data:** You must keep specific numbers, dates, and source names (e.g., "cited the 2024 IMF Report"). Do not generalize these into "some statistics."
  2. **Track the Argument Flow:** Do not just list who spoke. Describe the *state* of the argument (e.g., "Agent A proposed X, Agent B refuted X using Y").
  3. **Identify Open Threads:** Explicitly state what question or challenge is currently "on the table" and unanswered.
  4. **Tone:** Clinical, objective, and dense.

  ### Output Format (Strict JSON Structure)
  Return a valid JSON object with the following schema:
  {{
    "debate_topic": "The original topic",
    "agent_a_position": "Summary of Proponent's core stance and key evidence so far",
    "agent_b_position": "Summary of Opponent's core stance and key evidence so far",
    "key_clashes": [
      "List of specific points where agents strongly disagreed"
    ],
    "agreed_points": [
      "List of points where agents reached consensus (if any)"
    ],
    "immediate_context": "The very last thing that happened (e.g., 'Agent B just asked a question about inflation rates')"
  }}
  """

  messages = state["messages"]

  history_text = ""

  if state.get("summary"):
      history_text = state["summary"]

  history_text += "\n".join([f"{m.type}: {m.content}" for m in messages])

  formatted_prompt = SUMMARIZATION_PROMPT.format(conversation_history=history_text)
  structured_model = forum_model.with_structured_output(DebateState)
  response = await structured_model.ainvoke([HumanMessage(content=formatted_prompt)])

  try:
      summary_data = response.dict()

      summary_text = f"""
      # PREVIOUS DEBATE SUMMARY
      * **Topic:** {summary_data.get('debate_topic')}
      * **Agent A Stance:** {summary_data.get('agent_a_position')}
      * **Agent B Stance:** {summary_data.get('agent_b_position')}
      * **Key Clashes:** {'; '.join(summary_data.get('key_clashes', []))}
      * **Context:** {summary_data.get('immediate_context')}
      """
      return summary_text

  except Exception as e:
      print(f"Summarization failed: {e}")
      return "No summary"

async def summarization_node(state):

    messages = state.get("messages",[])
    summary_text = await forum_summary_generator(state)

    delete_messages = [RemoveMessage(id=m.id) for m in messages if m.id]

    return {"messages": delete_messages,"summary":summary_text}



def forum_router_conditions(state)->Literal["Debate Supervisor","Summarizer","Report Generator"]:
  if state.get("current_round",1) > state["max_rounds"]:
        return "Report Generator"

  if forum_model.get_num_tokens_from_messages(state["messages"])>(forum_model.profile["max_input_tokens"]*0.75) or forum_model.get_num_tokens(state.get("summary",""))>(forum_model.profile["max_input_tokens"]*0.75):
    return "Summarizer"

  else:
    return "Debate Supervisor"

def forum_router(state):
  return {}

async def report_generator_node(state):

    conversation_summary = await forum_summary_generator(state)

    recent_messages = state["debate_history"][-5:]
    formatted_recent = "\n".join([f"{m.type}: {m.content}" for m in recent_messages if m.type in ["human","ai"]])


    prompt = f"""
    You are a Technical Reporter. Generate a comprehensive Final Report based on this debate.

    ### Source Data
    PREVIOUS SUMMARY OF DEBATE:
    {conversation_summary}

    RECENT TRANSCRIPT:
    {formatted_recent}

    ### Requirements
    1. **Format:** Output purely in Markdown.
    2. **Citations:** You MUST attribute arguments to the specific Persona (Agent 1 or Agent 2).
    3. **Hallucination Check:** Do NOT add outside information. Only report on what was said in the transcript/summary.
    """

    structured_model = forum_model.with_structured_output(FinalReport)
    response = structured_model.invoke([HumanMessage(content=prompt)])


    return {"final_report": response}


def clean_mermaid_syntax(mermaid_code: str) -> str:
    """
    Sanitizes LLM-generated Mermaid code to prevent rendering errors.
    """
    # 1. Remove Markdown code fences (```mermaid ... ```)
    clean_code = re.sub(r"```mermaid", "", mermaid_code, flags=re.IGNORECASE)
    clean_code = re.sub(r"```", "", clean_code)

    # 2. Remove parent graph definitions if the model hallucinated them
    clean_code = re.sub(r"^graph (TD|LR|TB|BT)", "", clean_code, flags=re.MULTILINE | re.IGNORECASE)
    clean_code = re.sub(r"^flowchart (TD|LR|TB|BT)", "", clean_code, flags=re.MULTILINE | re.IGNORECASE)

    # 3. Fix potential unescaped quotes in labels
    def fix_quotes(match):
        content = match.group(1)
        clean_content = content.replace('"', "'")
        return f'["{clean_content}"]'

    clean_code = re.sub(r'\["(.*?)"\]', fix_quotes, clean_code)

    return clean_code.strip()

def generate_safe_mermaid(digest: RoundDigest) -> str:
    r_prefix = f"R{digest.round_number}"
    
    mermaid_lines = [f"subgraph {r_prefix}_Main [Round {digest.round_number}: {digest.winner_of_round} Wins]"]
    
    # 1. Define Nodes with correct styling
    for node in digest.graph_nodes:
        # Create unique ID programmatically
        full_id = f"{r_prefix}_{node.id}"
        
        # Apply shape based on type (Logic controlled by Python, not LLM)
        if node.type == "claim":
            shape_open, shape_close = "[", "]"
        elif node.type == "evidence":
            shape_open, shape_close = "((", "))"
        elif node.type == "attack":
            shape_open, shape_close = "{{", "}}"
        else:
            shape_open, shape_close = "[", "]" # Default
            
        # Sanitize label to prevent syntax errors
        clean_label = node.label.replace('"', "'")
        mermaid_lines.append(f'    {full_id}{shape_open}"{clean_label}"{shape_close}')

    # 2. Define Edges
    for edge in digest.graph_edges:
        src = f"{r_prefix}_{edge.source_id}"
        tgt = f"{r_prefix}_{edge.target_id}"
        clean_rel = edge.relationship.replace('"', "'")
        mermaid_lines.append(f"    {src} -->|{clean_rel}| {tgt}")

    mermaid_lines.append("end")
    
    return "\n".join(mermaid_lines)

async def round_digester_node(state):
    current_round = state.get("current_round",1)
    latest_messages = "\n".join([f"{m.type}: {m.content}" for m in state["messages"]])

    digest_prompt = f"""
    You are an Expert Debate Adjudicator and Logic Mapper. 
    The current debate round ({current_round}) has just finished.

    ### INPUT DATA
    {latest_messages}

    ### GOAL
    1. Determine the winner based on logical strength, not just style.
    2. Extract the "Logical Skeleton" of the argument for visualization.

    ### INSTRUCTIONS FOR LOGIC MAPPING
    Identify the core nodes of the argument.
    - **Claims:** The main assertions made by agents.
    - **Evidence:** Specific data, numbers, or sources cited.
    - **Attacks:** Direct counter-arguments or logical fallacies pointed out.

    Link these nodes to show the flow of debate (e.g., Evidence -> Supports -> Claim, or Attack -> Weakens -> Claim).
    Keep node labels extremely short (max 5 words) for readability.
    """


    digester_model = forum_model.with_structured_output(RoundDigest)
    digest_result = await digester_model.ainvoke([SystemMessage(content=digest_prompt)])
    new_ledger_entry = f"""
    --- ROUND {current_round} SUMMARY ---
    [Topic]: {state.get('query')}
    [Agent A Argued]: {', '.join(digest_result.key_arguments_pro)}
    [Agent B Argued]: {', '.join(digest_result.key_arguments_con)}
    [Winner]: {digest_result.winner_of_round}
    -------------------------------------
    """
    updated_ledger = state.get("summary", "") + new_ledger_entry

    delete_messages = [RemoveMessage(id=m.id) for m in state["messages"] if m.id]

    cleaned_graph = clean_mermaid_syntax(generate_safe_mermaid(digest_result))



    return {
        "round_digests": [digest_result],
        "running_mermaid_graph": f"\n{cleaned_graph}",
        "summary": updated_ledger,
        "messages": delete_messages,
        "current_round": state.get("current_round",1)+1
    }


async def create_forum_graph():
    forum_blueprint = StateGraph(ForumState)
    forum_blueprint.add_node("Persona Creator",persona_generator)
    forum_blueprint.add_node("Debate Supervisor",supervisor_agent)
    forum_blueprint.add_node("Forum Router",forum_router)
    forum_blueprint.add_node("Persona Agent",persona)
    forum_blueprint.add_node("Summarizer",summarization_node)
    forum_blueprint.add_node("tools",ToolNode([fetch_data]))
    forum_blueprint.add_node("Round Digest", round_digester_node)
    forum_blueprint.add_node("Report Generator", report_generator_node)
    forum_blueprint.add_node("Persona Routing", persona_routing)

    forum_blueprint.add_edge(START,"Persona Creator")
    forum_blueprint.add_edge("Persona Creator","Forum Router")
    forum_blueprint.add_conditional_edges("Forum Router",forum_router_conditions)
    forum_blueprint.add_edge("Summarizer","Forum Router")
    forum_blueprint.add_edge("Debate Supervisor","Persona Agent")
    forum_blueprint.add_conditional_edges("Persona Agent",tools_condition,{"tools":"tools",END:"Persona Routing"})
    forum_blueprint.add_edge("tools","Persona Agent")
    forum_blueprint.add_conditional_edges("Persona Routing",persona_routing_conditions)
    forum_blueprint.add_edge("Round Digest", "Forum Router")
    forum_blueprint.add_edge("Report Generator", END)
    forum_graph = forum_blueprint.compile(checkpointer=InMemorySaver())

    display(Image(forum_graph.get_graph().draw_mermaid_png()))

    return forum_graph

if __name__ == "__main__":
    import asyncio
    async def main():
        forum_graph = await create_forum_graph()
        thread_id = str(uuid.uuid4())
        config = {"configurable":{"thread_id":thread_id}}
        initial_state = ForumState(
            query="Is AI investment in 2024 a bubble or undervalued opportunity?",
            max_rounds=1,
            current_round=1,
            vector_store="chroma_langchain_db/417efa99-2e23-480d-9fc4-65e9431f36da",
            step=0,
            messages=[],
            round_digests=[],
            running_mermaid_graph="",
            debate_history=[],
            summary="",
            final_report=None
        )
        async for res in forum_graph.astream(initial_state,config=config,stream_mode="updates"):
            print(res)

    asyncio.run(main())