Rsearcher.py

import asyncio
import json
import os
import operator
from typing import List, Literal, Annotated, Union
from abc import ABC, abstractmethod
from langchain.chat_models import init_chat_model
from langchain_core.messages import (
    AIMessage, 
    HumanMessage, 
    SystemMessage, 
    ToolMessage, 
    filter_messages,
    BaseMessage
)
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
from typing_extensions import TypedDict
from tavily import AsyncTavilyClient
import logging
import time

# ✅ Load environment variables
from dotenv import load_dotenv
load_dotenv()

from prompt import research_system_prompt, compress_research_system_prompt
from config import Configuration

# Configure logging
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(existing: list, new: list) -> list:
    """Override reducer that replaces old values with new ones"""
    return new


class ResearcherState(TypedDict):
    """State for individual researchers conducting research."""
    
    researcher_messages: Annotated[list[Union[BaseMessage, dict]], operator.add]
    tool_call_iterations: int
    research_topic: str
    compressed_research: str
    raw_notes: Annotated[list[str], override_reducer]


class ResearcherOutputState(BaseModel):
    """Output state from individual researchers."""
    
    compressed_research: str
    raw_notes: Annotated[list[str], override_reducer] = []


# ===========================
# Tools Definition
# ===========================

@tool
async def tavily_search_tool(
    search_queries: Union[List[str], str],
    max_results: int = 5,
    topic: Literal["general", "news", "finance"] = "general"
) -> str:
    """
    Perform parallel searches using Tavily API
    
    Args:
        search_queries: List of search queries (or JSON string of list)
        max_results: Maximum results per query
        topic: Search topic category
        
    Returns:
        JSON string of search results
    """
    logger.info(f"tavily_search_tool | Raw input type: {type(search_queries)}")
    
    # Convert string to list if needed
    if isinstance(search_queries, str):
        try:
            search_queries = json.loads(search_queries)
            logger.info(f"Parsed JSON string to list: {search_queries}")
        except json.JSONDecodeError:
            # Fallback: split by comma
            search_queries = [q.strip().strip("'\"") for q in search_queries.split(',')]
            logger.info(f"Split string to list: {search_queries}")
    
    if not isinstance(search_queries, list):
        return f"Error: search_queries must be a list, got {type(search_queries)}"
    
    logger.info(f"tavily_search_tool | Searching {len(search_queries)} queries")
    
    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:
        search_tasks = [
            client.search(
                query=query,
                max_results=max_results,
                include_raw_content=True,
                topic=topic
            )
            for query in search_queries
        ]
        
        search_results = await asyncio.gather(*search_tasks)
        
        total_results = sum(len(result.get('results', [])) for result in search_results)
        logger.info(f"tavily_search_tool | Found {total_results} total results")
        
        return json.dumps(search_results, indent=2)
        
    except Exception as e:
        logger.error(f"tavily_search_tool | Failed: {str(e)}")
        return f"Error performing search: {str(e)}"


@tool
def think_tool(reflection: str) -> str:
    """Strategic reflection tool for research planning.
    
    Use this to pause and think about your research progress.
    
    Args:
        reflection: Your detailed reflection on research progress
        
    Returns:
        Confirmation message
    """
    logger.info(f"think_tool | Reflection: {reflection[:100]}...")
    return f"Reflection recorded: {reflection}"


# ===========================
# Node Classes (OOP)
# ===========================

class BaseResearchNode(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):
        """Each node must implement this method"""
        pass
    
    async def __call__(self, state: dict, config: RunnableConfig):
        """Entry point for LangGraph"""
        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:
        """Helper method: Extract model configuration from config"""
        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 ResearcherNode(BaseResearchNode):
    """Node responsible for thinking and decision-making"""
    
    def __init__(self):
        super().__init__("Researcher")
        self.tools = [tavily_search_tool, think_tool]
    
    async def execute(
        self, 
        state: dict, 
        config: RunnableConfig
    ) -> Command[Literal["researcher_tools"]]:
        """Execute main researcher logic"""
        configurable = Configuration.from_runnable_config(config)
        researcher_messages = state.get("researcher_messages", [])
        
        # Prepare model
        model_config = self.get_model_config(config, configurable.research_model)
        research_model = (
            configurable_model
            .bind_tools(self.tools)
            .with_retry(stop_after_attempt=configurable.max_structured_output_retries)
            .with_config(model_config)
        )
        
        # Prepare messages
        system_prompt = research_system_prompt.format(
            mcp_prompt="",
            date=time.strftime("%Y-%m-%d")
        )
        messages = [SystemMessage(content=system_prompt)] + researcher_messages
        
        # Call model
        self.logger.info("Calling model...")
        response = await research_model.ainvoke(messages)
        
        if hasattr(response, 'tool_calls') and response.tool_calls:
            self.logger.info(f"Tool calls: {[tc['name'] for tc in response.tool_calls]}")
        
        # Return Command
        return Command(
            goto="researcher_tools",
            update={
                "researcher_messages": [response],
                "tool_call_iterations": state.get("tool_call_iterations", 0) + 1
            }
        )


class ToolExecutorNode(BaseResearchNode):
    """Node responsible for executing tools"""
    
    def __init__(self):
        super().__init__("ToolExecutor")
        self.tools_by_name = {
            "tavily_search_tool": tavily_search_tool,
            "think_tool": think_tool
        }
    
    async def execute(
        self, 
        state: dict, 
        config: RunnableConfig
    ) -> Command[Literal["researcher", "compress_research"]]:
        """Execute tools and decide to continue or stop"""
        configurable = Configuration.from_runnable_config(config)
        researcher_messages = state.get("researcher_messages", [])
        most_recent_message = researcher_messages[-1]
        
        # Check exit conditions
        has_tool_calls = hasattr(most_recent_message, 'tool_calls') and bool(most_recent_message.tool_calls)
        exceeded_iterations = state.get("tool_call_iterations", 0) >= configurable.max_researcher_iterations
        
        if not has_tool_calls or exceeded_iterations:
            self.logger.info(f"Exiting - has_tools: {has_tool_calls}, iterations: {state.get('tool_call_iterations', 0)}")
            return Command(goto="compress_research")
        
        # Execute tools
        tool_outputs = await self._execute_tools(most_recent_message.tool_calls, config)
        
        self.logger.info(f"Executed {len(tool_outputs)} tools")
        
        # Continue research loop
        return Command(
            goto="researcher",
            update={"researcher_messages": tool_outputs}
        )
    
    async def _execute_tools(self, tool_calls: List, config: RunnableConfig) -> List[ToolMessage]:
        """Helper method: Execute list of tool calls"""
        tool_outputs = []
        
        for tool_call in tool_calls:
            tool_name = tool_call["name"]
            tool_args = tool_call["args"]
            
            self.logger.info(f"Executing {tool_name}")
            
            try:
                if tool_name in self.tools_by_name:
                    result = await self.tools_by_name[tool_name].ainvoke(tool_args, config)
                else:
                    result = f"Error: Unknown tool {tool_name}"
                
                tool_outputs.append(ToolMessage(
                    content=str(result),
                    name=tool_name,
                    tool_call_id=tool_call["id"]
                ))
                
            except Exception as e:
                self.logger.error(f"Tool {tool_name} failed: {e}")
                tool_outputs.append(ToolMessage(
                    content=f"Error: {str(e)}",
                    name=tool_name,
                    tool_call_id=tool_call["id"]
                ))
        
        return tool_outputs


class CompressionNode(BaseResearchNode):
    """Node responsible for compressing research findings"""
    
    def __init__(self):
        super().__init__("Compression")
    
    async def execute(self, state: dict, config: RunnableConfig) -> dict:
        """Compress research findings into summary"""
        configurable = Configuration.from_runnable_config(config)
        researcher_messages = state.get("researcher_messages", [])
        
        # Prepare compression model
        model_config = self.get_model_config(config, configurable.summarization_model)
        compression_model = configurable_model.with_config(model_config)
        
        # Prepare research content
        research_content = "\n\n".join([
            f"{'AI' if isinstance(msg, AIMessage) else 'Tool'}: {msg.content[:1000]}..."
            for msg in researcher_messages
            if isinstance(msg, (AIMessage, ToolMessage))
        ])
        
        # Prepare messages
        messages = [
            SystemMessage(content=compress_research_system_prompt),
            HumanMessage(content=f"Research findings:\n\n{research_content}\n\nProvide summary.")
        ]
        
        # Generate summary
        self.logger.info("Generating summary...")
        response = await compression_model.ainvoke(messages)
        
        # Extract raw notes
        raw_notes = "\n".join([
            str(msg.content)
            for msg in filter_messages(researcher_messages, include_types=["tool", "ai"])
        ])
        
        return {
            "compressed_research": str(response.content),
            "raw_notes": [raw_notes]
        }


# ===========================
# Graph Creation
# ===========================

# Create node instances
researcher_node = ResearcherNode()
tool_executor_node = ToolExecutorNode()
compression_node = CompressionNode()


def create_research_graph() -> CompiledStateGraph:
    """Create and compile the research graph"""
    builder = StateGraph(
        ResearcherState,
        output=ResearcherOutputState,
        context_schema=Configuration
    )
    
    # Add nodes
    builder.add_node("researcher", researcher_node)
    builder.add_node("researcher_tools", tool_executor_node)
    builder.add_node("compress_research", compression_node)
    
    # Define edges
    builder.add_edge(START, "researcher")
    builder.add_edge("compress_research", END)
    
    return builder.compile()


# Create compiled graph
research_graph = create_research_graph()


# ===========================
# Main Test Function
# ===========================

async def main():
    """Test the research graph"""
    print("\n" + "="*60)
    print("🚀 Testing OOP-based Research System")
    print("="*60)
    
    # Load API keys
    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!")
        print(f"   NVIDIA_API_KEY: {'✓' if nvidia_api_key else '✗'}")
        print(f"   TAVILY_API_KEY: {'✓' if tavily_api_key else '✗'}")
        return
    
    # Configuration
    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",
            "summarization_model": "meta/llama-3.1-8b-instruct",
            "max_researcher_iterations": 3,
            "max_structured_output_retries": 3
        }
    }
    
    # Initial state
    initial_state = {
        "research_topic": "Latest AI trends in 2024",
        "researcher_messages": [HumanMessage(content="Research: Latest AI trends in 2024")],
        "tool_call_iterations": 0,
        "compressed_research": "",
        "raw_notes": []
    }
    
    try:
        # Run research
        print("\n📡 Starting research workflow...\n")
        result = await research_graph.ainvoke(initial_state, config)
        
        print("\n" + "="*60)
        print("📊 Results")
        print("="*60)
        print(f"✓ Summary length: {len(result.get('compressed_research', ''))} characters")
        print(f"✓ Raw notes length: {len(result.get('raw_notes', [])[0]) if result.get('raw_notes') else 0} characters")
        print(f"\n📝 Summary:\n{result.get('compressed_research', 'N/A')[:800]}...")
        
    except Exception as e:
        print(f"\n❌ Error during execution: {e}")
        logger.error("Execution failed", exc_info=True)


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