main.py

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from pydantic import SecretStr

from agents.ProjectState import ProjectState
from agents.ExpertAgent import ExpertAgent
from agents.ProjectManager import ProjectManager
from agents.TeamState import TeamState
from config import SILICONFLOW_API_KEY, MODEL_BASE_URL, MODEL_NAME

llm = ChatOpenAI(
    base_url=MODEL_BASE_URL, 
    api_key=SecretStr(SILICONFLOW_API_KEY), 
    model=MODEL_NAME, 
    temperature=0.2
)
manager_instance = ProjectManager(llm)


def executor_node(state: TeamState):
    active_agent_name = state["next_active_agent"]
    agent_instance: ExpertAgent = state["agents_pool"][active_agent_name]

    result = agent_instance.execute(
        task_description=state["current_task"],
        project_state=state["project_state"]
    )

    print(f"\n[{active_agent_name}'s Projection Log] \n{result.action_log}")
    print(f"\n[{active_agent_name}'s New Contribution] \n{result.contribution}\n" + "=" * 50)

    # Append new contribution to progress list
    new_project_state = state["project_state"].copy()
    new_progress = list(new_project_state.get("current_progress", []))
    new_progress.append({
        "author": active_agent_name,
        "content": result.contribution
    })
    new_project_state["current_progress"] = new_progress

    return {
        "project_state": new_project_state,
        "next_active_agent": "manager"
    }


workflow = StateGraph(TeamState)

workflow.add_node("manager", manager_instance.analyze_and_route)
workflow.add_node("executor", executor_node)

workflow.set_entry_point("manager")


def route_logic(state: TeamState):
    if state.get("is_finished"):
        return END
    return "executor"


workflow.add_conditional_edges("manager", route_logic)
workflow.add_edge("executor", "manager")


app = workflow.compile()

if __name__ == "__main__":
    print("\n" + "=" * 50)
    print("🚀 Welcome to ProjectAI Sandbox & Projection System 🚀")
    print("=" * 50 + "\n")

    print("✨ Please start your projection setup ✨")
    idea = input("💡 Enter your Project Idea (e.g., An AI-based automatic fish feeder): ")
    team_background = input("👥 Enter Team Background (e.g., High school students, physics knowledge, beginner coding): ")

    initial_project_state: ProjectState = {
        "idea": idea,
        "team_background": team_background,
        "current_progress": []  # Initialize as empty list
    }

    initial_team_state: TeamState = {
        "messages": [],
        "agents_pool": {},
        "next_active_agent": "manager",
        "project_state": initial_project_state,
        "is_finished": False,
        "current_task": "",
        "final_talents": "",
        "final_result": ""
    }

    for event in app.stream(initial_team_state, {"recursion_limit": 500}):
        for node_name, state_update in event.items():
            if node_name == "manager":
                if state_update.get("is_finished"):
                    print("\n" + "🎉" * 20)
                    print("🎉 [System Broadcast] Manager announced project closure!")
                    print("=" * 40)
                    talents = state_update.get("final_talents") or "No talent summary provided"
                    result = state_update.get("final_result") or "No result report provided"
                    print(f"👤 【Required Talent Summary】:\n{talents}")
                    print("-" * 40)
                    print(f"📦 【Final Project Archive】:\n{result}")
                    print("=" * 40)
                    print("🎉" * 20 + "\n")
                else:
                    print(
                        f"\n🎯 [Manager Instruction] -> Next step assigned to [{state_update['next_active_agent']}]: {state_update['current_task']}")
            elif node_name == "executor":
                print("\n🔄 [Handover] Expert finished projection, control returned to Manager...")