graph_states.py

from pydantic import BaseModel, Field
from typing import TypedDict, List, Annotated
from langchain_core.messages import BaseMessage
from langgraph.graph import MessagesState
from langchain.agents import AgentState
import operator
from enum import Enum

# ------------ RESEARCH GRAPH --------------

class Queries(BaseModel):
    queries: list[str] = Field(..., max_length=5,description="The list of queries generated by user query to search on different sources")

class Extractor(BaseModel):
  query_understanding: str
  extracted_information: list[str]
  missing_information: list[str]

class Analyzer(BaseModel):
  query_statisfied: bool
  reasoning: str
  query: str
  missing_or_uncertain_information: list[str]

class DataExtractionState(TypedDict):
    query:str
    subqueries:list[str]
    subquery_extracted_data:list[str] | None
    subquery_reviews:list[Analyzer]
    results:list[str]
    vector_store:str
    query_limit:int 

# ------------ END RESEARCH GRAPH --------------

# ------------ FORUM MODEL --------------

class Persona(BaseModel):
  persona_name:str = Field(..., description="The name of the persona")
  persona_prompt:str = Field(..., description="The prompt of the persona")

class Agent_Persona_Generator(BaseModel):
  agent_1:Persona = Field(..., description="The first persona")
  agent_2:Persona = Field(..., description="The second persona")

class SupervisorState(AgentState):
      vector_store:str

class DebateState(BaseModel):
    debate_topic: str = Field(
        description="The original debate topic"
    )
    agent_a_position: str = Field(
        description="Summary of Proponent's core stance and key evidence so far"
    )
    agent_b_position: str = Field(
        description="Summary of Opponent's core stance and key evidence so far"
    )
    key_clashes: List[str] = Field(
        description="Specific points where agents strongly disagreed"
    )
    agreed_points: List[str] = Field(
        description="Points where agents reached consensus, if any"
    )
    immediate_context: str = Field(
        description="The very last thing that happened in the debate"
  )

class NodeType(str, Enum):
    CLAIM = "claim"       # The main point
    EVIDENCE = "evidence" # Data/Stats supporting a claim
    ATTACK = "attack"     # A direct rebuttal

class LogicNode(BaseModel):
    id: str = Field(description="A short, unique variable name (e.g., 'cost_argument'). Do NOT use prefixes like R1_.")
    label: str = Field(description="Max 5 words summary of the point.")
    type: NodeType

class LogicEdge(BaseModel):
    source_id: str
    target_id: str
    relationship: str = Field(description="Short verb describing the link, e.g., 'supports', 'refutes', 'weakens'.")

class RoundDigest(BaseModel):
    round_number: int
    winner_of_round: str = Field(description="Who had the stronger argument? 'Agent A', 'Agent B', or 'Draw'")
    winner_rationale: str = Field(description="One sentence explaining WHY they won.")
    key_arguments_pro: List[str]
    key_arguments_con: List[str]
    
    # Instead of raw string, we ask for structured graph data
    graph_nodes: List[LogicNode] = Field(description="The key logical entities in this round.")
    graph_edges: List[LogicEdge] = Field(description="How the entities connect.")

class FinalReport(BaseModel):
  title: str
  executive_summary: str = Field(description="High-level overview of the debate outcome")
  key_arguments_agent_1: list[str]
  key_arguments_agent_2: list[str]
  consensus_points: list[str]
  unique_perspectives: list[str]
  

class ForumState(MessagesState):
    query:str
    summary:str
    max_rounds:int
    current_round:int
    agent_1:Persona
    agent_2:Persona
    vector_store:str
    round_digests: Annotated[list[RoundDigest],operator.add]
    running_mermaid_graph: Annotated[str,operator.add]
    debate_history: Annotated[list[BaseMessage],operator.add]
    final_report: FinalReport | None 
    references: Annotated[list[list[str]],operator.add]
    step:int

# ------------ END FORUM MODEL --------------



# ------------ Final Graph -------------------

class FinalState(MessagesState):
  query:str
  vector_store:str
  query_limit:int
  max_rounds:int
  current_round:int
  step:int
  round_digests: Annotated[list[RoundDigest],operator.add]
  running_mermaid_graph: Annotated[str,operator.add]
  debate_history: Annotated[list[BaseMessage],operator.add]
  references: Annotated[list[list[str]],operator.add]
  final_report: FinalReport | None 

# ------------ END Final Graph --------------