Supervisor.py

import asyncio
import json
import os
import logging
import time
from typing import List, Literal, Annotated, Union
from abc import ABC, abstractmethod
from Rsearcher import researcher_subgraph
from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain_core.messages import (
    AIMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
    MessageLikeRepresentation,
)
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langgraph.graph import END, START, StateGraph
from langgraph.types import Command
from langgraph.graph.state import CompiledStateGraph
from pydantic import BaseModel, Field
from typing_extensions import TypedDict
from tavily import AsyncTavilyClient

from prompt import research_system_prompt
from config import Configuration

load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

configurable_model = init_chat_model(
    configurable_fields=("model", "api_key", "base_url", "model_provider"),
)


def override_reducer(current_value, new_value):
    """Reducer که هم override و هم append را پشتیبانی می‌کند."""
    if isinstance(new_value, dict) and new_value.get("type") == "override":
        return new_value.get("value", new_value)
    if isinstance(current_value, list) and isinstance(new_value, list):
        return current_value + new_value
    return new_value


class SupervisorState(TypedDict):
    supervisor_messages: Annotated[List[MessageLikeRepresentation], override_reducer]
    research_brief: str
    notes: Annotated[List[str], override_reducer]
    research_iterations: int
    raw_notes: Annotated[List[str], override_reducer]
    pending_research_topics: Annotated[List[str], override_reducer]


class ConductResearch(BaseModel):
    research_topic: str = Field(description="The topic to research")


class ResearchComplete(BaseModel):
    pass


def think_tool(reflection: str) -> str:
    logger.info(f"think_tool | Reflection: {reflection[:100]}...")
    return f"Reflection recorded: {reflection}"


@tool
async def tavily_search_tool(
    search_queries: Union[List[str], str],
    max_results: int = 5,
    topic: Literal["general", "news", "finance"] = "general",
) -> str:
    """
    ابزار جستجو برای researchers.
    """
    logger.info(
        f"tavily_search_tool | Searching {len(search_queries) if isinstance(search_queries, list) else 1} queries"
    )

    if isinstance(search_queries, str):
        try:
            search_queries = json.loads(search_queries)
        except json.JSONDecodeError:
            search_queries = [q.strip().strip("'\"") for q in search_queries.split(",")]

    if not isinstance(search_queries, list):
        return "Error: search_queries must be a list"

    tavily_api_key = os.getenv("TAVILY_API_KEY")
    if not tavily_api_key:
        raise ValueError("TAVILY_API_KEY is required")

    client = AsyncTavilyClient(api_key=tavily_api_key)

    try:
        tasks = [
            client.search(
                query=query,
                max_results=max_results,
                include_raw_content=True,
                topic=topic,
            )
            for query in search_queries
        ]
        results = await asyncio.gather(*tasks)
        return json.dumps(results, indent=2)
    except Exception as e:
        logger.error(f"Search failed: {str(e)}")
        return f"Error: {str(e)}"


class BaseAgentNode(ABC):
    """Base class for all research nodes."""

    def __init__(self, name: str):
        self.name = name
        self.logger = logging.getLogger(f"{__name__}.{name}")

    @abstractmethod
    async def execute(self, state: dict, config: RunnableConfig):
        ...

    async def __call__(self, state: dict, config: RunnableConfig):
        self.logger.info(f"{self.name} | Starting")
        start_time = time.monotonic()
        try:
            result = await self.execute(state, config)
            elapsed = time.monotonic() - start_time
            self.logger.info(f"{self.name} | Complete in {elapsed:.2f}s ✓")
            return result
        except Exception as e:
            self.logger.error(f"{self.name} | Failed: {e}", exc_info=True)
            raise

    def get_model_config(self, config: RunnableConfig, model_name: str) -> dict:
        configurable = Configuration.from_runnable_config(config)
        return {
            "model": model_name,
            "api_key": configurable.NVIDIA_api_key,
            "base_url": configurable.research_base_url,
            "model_provider": "openai",
            "tags": ["langsmith:nostream"],
        }


class Supervisor(BaseAgentNode):
    def __init__(self):
        super().__init__("SupervisorAgent")

    async def execute(
        self, state: SupervisorState, config: RunnableConfig
    ) -> Command[Literal["supervisor_tools"]]:
        configurable = Configuration.from_runnable_config(config)
        model_config = self.get_model_config(config, configurable.research_model)

        supervisor_tools = [ConductResearch, ResearchComplete, think_tool]

        model = (
            configurable_model.bind_tools(supervisor_tools)
            .with_retry(stop_after_attempt=3)
            .with_config(model_config)
        )

        messages = state.get("supervisor_messages", [])
        prompt = research_system_prompt.format(
            configurable.max_researcher_iterations,
            configurable.max_concurrent_research_units,
        )
        messages = [SystemMessage(content=prompt)] + messages

        self.logger.info(f"Invoking supervisor with {len(messages)} messages")
        response: AIMessage = await model.ainvoke(messages)
        self.logger.info(f"Supervisor Agent response: {response}")

        if getattr(response, "tool_calls", None):
            self.logger.info(
                f"Tool calls: {[tc['name'] for tc in response.tool_calls]}"
            )
            return Command(
                goto="supervisor_tools",
                update={
                    "supervisor_messages": [response],
                    "research_iterations": state.get("research_iterations", 0) + 1,
                },
            )

        self.logger.info("No tool calls, ending workflow")
        return Command(goto="__end__")


class SupervisorTool(BaseAgentNode):
    def __init__(self):
        super().__init__("SupervisorTools")

    async def execute(
        self, state: SupervisorState, config: RunnableConfig
    ) -> Command[Literal["supervisor", "conduct_research", "__end__"]]:
        configurable = Configuration.from_runnable_config(config)

        messages = state.get("supervisor_messages", [])
        if not messages:
            self.logger.warning("No supervisor messages, ending")
            return Command(goto="__end__")

        last_message = messages[-1]
        has_tool_calls = bool(getattr(last_message, "tool_calls", []))
        exceeded_iterations = state.get("research_iterations", 0) >= configurable.max_researcher_iterations

        if not has_tool_calls:
            self.logger.info("No tool calls, ending")
            return Command(goto="__end__")

        if exceeded_iterations:
            self.logger.info("Max iterations exceeded, ending workflow")
            return Command(goto="__end__")

        tool_outputs: List[ToolMessage] = []
        research_topics: List[str] = []
        next_node: Literal["supervisor", "conduct_research", "__end__"] = "supervisor"

        for tool_call in last_message.tool_calls:
            tool_name = tool_call["name"]
            tool_args = tool_call.get("args", {})
            self.logger.info(f"Processing tool: {tool_name}")

            if tool_name == "ConductResearch":
                topic = tool_args.get("research_topic", "")
                research_topics.append(topic)
                tool_outputs.append(
                    ToolMessage(
                        content=f"Research initiated for: {topic}",
                        name=tool_name,
                        tool_call_id=tool_call["id"],
                    )
                )
                next_node = "conduct_research"

            elif tool_name == "ResearchComplete":
                tool_outputs.append(
                    ToolMessage(
                        content="Research marked as complete",
                        name=tool_name,
                        tool_call_id=tool_call["id"],
                    )
                )
                next_node = "__end__"

            elif tool_name == "think_tool":
                reflection = tool_args.get("reflection", "")
                result = think_tool(reflection)
                tool_outputs.append(
                    ToolMessage(
                        content=result,
                        name=tool_name,
                        tool_call_id=tool_call["id"],
                    )
                )
                next_node = "supervisor"

        update: dict = {"supervisor_messages": tool_outputs}
        if research_topics:
            update["pending_research_topics"] = research_topics

        self.logger.info(f"Routing to: {next_node}")
        return Command(goto=next_node, update=update)


class ConductResearchNode(BaseAgentNode):
    def __init__(self):
        super().__init__("ConductResearch")

    async def execute(
        self, state: SupervisorState, config: RunnableConfig
    ) -> Command[Literal["supervisor"]]:
        topics = state.get("pending_research_topics", [])
        if not topics:
            self.logger.warning("No research topics found")
            return Command(
                goto="supervisor",
                update={
                    "notes": ["No research topics provided"],
                    "supervisor_messages": [
                        HumanMessage(content="No research topics provided.")
                    ],
                },
            )

        self.logger.info(f"Starting research on {len(topics)} topics")

        research_tasks = [
            researcher_subgraph.ainvoke({
                "researcher_messages": [
                    HumanMessage(content=topic)
                ],
                "research_topic": topic,
                "tool_call_iterations": 0,
                "raw_notes": []
            }, config)
            for topic in topics
        ]


        # TODO: اینجا جایی است که باید subgraph واقعی researcher را صدا بزنی
        research_results = await asyncio.gather(*research_tasks)

        
        # جمع‌آوری نتایج
        compressed_research = [
            result.get("compressed_research", "") 
            for result in research_results
        ]
        raw_notes = [
            note 
            for result in research_results 
            for note in result.get("raw_notes", [])
        ]

        research_brief = "\n\n".join(compressed_research)

        
        return Command(
            goto="supervisor",
            update={
                "research_brief": research_brief,
                "notes": compressed_research,
                "raw_notes": raw_notes,
                "supervisor_messages": [
                    HumanMessage(content=f"Research completed:\n{research_brief}")
                ],
                "pending_research_topics": {"type": "override", "value": []}
            }
        )



def create_supervisor_graph() -> CompiledStateGraph:
    builder = StateGraph(SupervisorState)

    builder.add_node("supervisor", Supervisor())
    builder.add_node("supervisor_tools", SupervisorTool())
    builder.add_node("conduct_research", ConductResearchNode())

    builder.add_edge(START, "supervisor")
    # سایر انتقال‌ها داینامیک هستند (با Command)

    return builder.compile()


async def main():
    print("\n" + "=" * 60)
    print("🚀 Testing Supervisor-Researcher System")
    print("=" * 60)

    nvidia_api_key = os.getenv("NVIDIA_API_KEY")
    tavily_api_key = os.getenv("TAVILY_API_KEY")

    if not nvidia_api_key or not tavily_api_key:
        print("❌ API Keys missing!")
        return

    config = {
        "configurable": {
            "NVIDIA_api_key": nvidia_api_key,
            "research_base_url": "https://integrate.api.nvidia.com/v1",
            "research_model": "meta/llama-3.1-8b-instruct",
            "max_researcher_iterations": 5,
            "max_concurrent_research_units": 3,
        }
    }

    initial_state: SupervisorState = {
        "supervisor_messages": [
            HumanMessage(
                content="Research: What are the latest developments in AI agents for 2025?"
            )
        ],
        "research_brief": "",
        "notes": [],
        "research_iterations": 0,
        "raw_notes": [],
        "pending_research_topics": [],
    }

    graph = create_supervisor_graph()

    try:
        print("\n📡 Starting supervisor workflow...\n")
        result = await graph.ainvoke(initial_state, config)

        print("\n" + "=" * 60)
        print("📊 Results")
        print("=" * 60)
        print(f"✓ Research iterations: {result.get('research_iterations', 0)}")
        print(f"✓ Notes count: {len(result.get('notes', []) )}")
        print(
            f"\n📝 Research Brief:\n{result.get('research_brief', 'N/A')[:500]}..."
        )
    except Exception as e:
        print(f"\n❌ Error: {e}")
        logger.error("Execution failed", exc_info=True)


if __name__ == "__main__":
    asyncio.run(main())