agent.py

# agent.py
import logging
from typing import Literal
from pydantic import BaseModel, Field
from langchain.chat_models import init_chat_model
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.runnables import RunnableConfig
from langgraph.types import Command
from langgraph.graph import END

from config import Configuration
from state import AgentState, ResearchQuestion
from prompt import lead_researcher_prompt

logger = logging.getLogger(__name__)

# Global configurable model
configurable_model = init_chat_model(
    configurable_fields=("model", "max_tokens", "api_key", "base_url", "model_provider"),
)

# Structured output models
class ClarifyWithUser(BaseModel):
    """Model for user clarification requests."""
    need_clarification: bool = Field(
        description="Whether the user needs to be asked a clarifying question."
    )
    question: str = Field(
        description="Question to ask the user. Empty if no clarification needed."
    )
    verification: str = Field(
        description="Verification message that research will start."
    )


class ClarifyAgent:
    """Agent for checking if user query needs clarification."""
    
    def __init__(self):
        """Initialize ClarifyAgent."""
        logger.info("ClarifyAgent initialized")
    
    async def invoke(
        self, 
        state: AgentState, 
        config: RunnableConfig
    ) -> Command[Literal["write_research_brief", "__end__"]]:
        """Check if user query needs clarification."""
        logger.info(f"Clarify agent started. Query: '{state.get('query', '')}'")
        
        try:
            configurable = Configuration.from_runnable_config(config)
            
            # تنظیم model configuration - بدون max_completion_tokens
            model_config = {
                "model": configurable.clarify_model,
                "api_key": configurable.NVIDIA_api_key,
                "base_url": configurable.research_base_url,
                "model_provider": "openai",
                # حذف max_tokens - NVIDIA از max_completion_tokens پشتیبانی نمی‌کند
                # "max_tokens": configurable.clarify_max_tokens,
                "tags": ["langsmith:nostream"]
            }
            
            # استفاده از method="json_mode" به جای json_schema
            model = (
                configurable_model
                .with_structured_output(
                    ClarifyWithUser,
                    method="json_mode",  # کلیدی! به جای json_schema
                    include_raw=False
                )
                .with_retry(stop_after_attempt=configurable.max_structured_output_retries)
                .with_config(model_config)
            )
            
            # Format messages - اضافه کردن دستورالعمل JSON
            messages_str = "\n".join([str(msg.content) for msg in state.get('messages', [])])
            
            prompt = f"""You must respond with valid JSON matching this schema:
{{
    "need_clarification": boolean,
    "question": "string (question to ask user, empty if no clarification needed)",
    "verification": "string (verification message that research will start)"
}}

{configurable.clarify_prompt.format(query=state.get('query', ''), messages=messages_str)}

Respond ONLY with valid JSON, no other text."""
            
            response = await model.ainvoke(prompt)
            
            logger.info(f"Clarification needed: {response.need_clarification}")
            
            if response.need_clarification:
                logger.info(f"Asking clarification: {response.question}")
                return Command(
                    goto="__end__",
                    update={"messages": [AIMessage(content=response.question)]}
                )
            else:
                logger.info(f"Proceeding to research: {response.verification}")
                return Command(
                    goto="write_research_brief",
                    update={"messages": [AIMessage(content=response.verification)]}
                )
        
        except Exception as e:
            logger.error(f"Error in clarify_agent: {e}", exc_info=True)
            return Command(
                goto="write_research_brief",
                update={"messages": [AIMessage(content="Processing your request...")]}
            )


class WriteResearchBrief:
    """Agent for transforming user messages into structured research brief."""
    
    def __init__(self):
        """Initialize WriteResearchBrief."""
        logger.info("WriteResearchBrief initialized")
    
    async def invoke(
        self, 
        state: AgentState, 
        config: RunnableConfig
    ) -> Command[Literal["__end__"]]:
        """Transform user messages into research brief."""
        logger.info("Starting research brief generation...")
        
        try:
            configurable = Configuration.from_runnable_config(config)
            
            # تنظیم model configuration - بدون max_completion_tokens
            model_config = {
                "model": configurable.research_model,
                "api_key": configurable.NVIDIA_api_key,
                "base_url": configurable.research_base_url,
                "model_provider": "openai",
                # حذف max_tokens
                # "max_tokens": configurable.research_model_max_tokens,
                "tags": ["langsmith:nostream"]
            }
            
            # استفاده از method="json_mode"
            model = (
                configurable_model
                .with_structured_output(
                    ResearchQuestion,
                    method="json_mode",  # کلیدی!
                    include_raw=False
                )
                .with_retry(stop_after_attempt=configurable.max_structured_output_retries)
                .with_config(model_config)
            )
            
            # آماده‌سازی messages
            messages_str = "\n".join([
                str(msg.content) for msg in state.get("messages", [])
            ])
            
            if not messages_str.strip():
                messages_str = state.get("query", "")
            
            # اضافه کردن دستورالعمل JSON
            prompt = f"""You must respond with valid JSON matching this schema:
{{
    "research_brief": "string (detailed research brief)"
}}

{configurable.research_topic_prompt.format(messages=messages_str)}

Respond ONLY with valid JSON, no other text."""
            
            logger.info(f"Calling model for research brief...")
            
            response = await model.ainvoke(prompt)
            
            logger.info(f"Research brief created: {response.research_brief[:100]}...")
            
            # ساخت supervisor prompt
            supervisor_system_prompt = lead_researcher_prompt.format(
                max_concurrent_research_units=configurable.max_concurrent_research_units,
                max_researcher_iterations=configurable.max_researcher_iterations
            )
            
            return Command(
                goto="__end__",
                update={
                    "research_brief": response.research_brief,
                    "supervisor_messages": {
                        "type": "override",
                        "value": [
                            SystemMessage(content=supervisor_system_prompt),
                            HumanMessage(content=response.research_brief)
                        ]
                    },
                    "messages": [
                        AIMessage(content=f"Research brief created:\n\n{response.research_brief}")
                    ]
                }
            )
            
        except Exception as e:
            logger.error(f"Error in write_research_brief: {e}", exc_info=True)
            return Command(
                goto="__end__",
                update={
                    "messages": [AIMessage(content=f"Error: {str(e)}")]
                }
            )