Iterative Research Agent

import os
from typing import Annotated, Literal, TypedDict
from dotenv import load_dotenv
import noveum_trace
from noveum_trace import NoveumTraceCallbackHandler
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph

load_dotenv()

# Initialize Noveum Trace
noveum_trace.init(
    api_key=os.getenv("NOVEUM_API_KEY"),
    project="customer-support-bot",
    environment="development"
)

# Define the research state
class ResearchState(TypedDict):
    messages: Annotated[list, "The messages in the conversation"]
    research_topic: str
    research_notes: Annotated[list, "Research notes gathered"]
    evaluation_score: float
    max_iterations: int
    current_iteration: int
    research_complete: bool

# Define research tools
@tool
def search_web(query: str) -> str:
    """Search the web for information about a query."""
    # Simulate web search with realistic results
    search_results = {
        "artificial intelligence": "AI is a branch of computer science focused on creating intelligent machines...",
        "machine learning": "Machine learning is a subset of AI that enables computers to learn without explicit programming...",
        "deep learning": "Deep learning uses neural networks with multiple layers to process data...",
        "natural language processing": "NLP is a field of AI that focuses on the interaction between computers and human language..."
    }
    
    # Return relevant results based on query
    for key, value in search_results.items():
        if key in query.lower():
            return f"Search results for '{query}': {value}"
    
    return f"Search results for '{query}': General information about the topic."

@tool
def analyze_information(info: str) -> str:
    """Analyze and summarize information."""
    return f"Analysis: {info} contains valuable insights and detailed information about the topic."

def research_node(state: ResearchState):
    """Node that performs research using tools."""
    print(f"🔍 Research iteration {state['current_iteration']}: {state['research_topic']}")
    
    # Search for information
    search_query = f"research about {state['research_topic']}"
    search_results = search_web(search_query)
    
    # Analyze the results
    analysis = analyze_information(search_results)
    
    # Add to research notes
    state["research_notes"].append({
        "iteration": state["current_iteration"],
        "query": search_query,
        "results": search_results,
        "analysis": analysis
    })
    
    # Add research message
    state["messages"].append(AIMessage(content=f"Research iteration {state['current_iteration']} completed: {analysis}"))
    
    return state

def evaluate_node(state: ResearchState):
    """Node that evaluates the quality of research gathered."""
    print(f"📊 Evaluating research quality...")
    
    # Simple evaluation based on research notes
    total_notes = len(state["research_notes"])
    quality_score = min(0.9, 0.3 + (total_notes * 0.1))
    
    state["evaluation_score"] = quality_score
    
    # Add evaluation message
    evaluation_msg = f"Research evaluation: {quality_score:.2f} quality score based on {total_notes} research iterations"
    state["messages"].append(AIMessage(content=evaluation_msg))
    
    print(f"📈 Quality score: {quality_score:.2f}")
    
    return state

def should_continue(state: ResearchState) -> Literal["research", "synthesize", "end"]:
    """Decide whether to continue researching, synthesize, or end."""
    print(f"🤔 Deciding next action...")
    
    # Check if we've reached max iterations
    if state["current_iteration"] >= state["max_iterations"]:
        print("⏰ Max iterations reached, synthesizing...")
        return "synthesize"
    
    # Check if quality is sufficient
    if state["evaluation_score"] >= 0.8:
        print("✅ Quality sufficient, synthesizing...")
        return "synthesize"
    
    # Continue researching
    print("🔄 Quality insufficient, continuing research...")
    state["current_iteration"] += 1
    return "research"

def synthesize_node(state: ResearchState):
    """Node that synthesizes all research into a final report."""
    print("📝 Synthesizing final research report...")
    
    # Create comprehensive report
    report = f"""
    # Research Report: {state['research_topic']}
    
    ## Summary
    Based on {state['current_iteration']} research iterations, here's what I found:
    
    """
    
    # Add findings from each iteration
    for i, note in enumerate(state["research_notes"], 1):
        report += f"### Iteration {i}\n{note['analysis']}\n\n"
    
    report += f"""
    ## Final Evaluation
    Quality Score: {state['evaluation_score']:.2f}
    Total Iterations: {state['current_iteration']}
    
    ## Conclusion
    This research provides comprehensive coverage of {state['research_topic']} with detailed analysis and insights.
    """
    
    # Add final message
    state["messages"].append(AIMessage(content=report))
    state["research_complete"] = True
    
    print("✅ Research synthesis completed!")
    
    return state

def create_iterative_research_agent():
    """Create an iterative research agent with tracing."""
    # Create LLM (without callbacks - will be passed at graph level)
    llm = ChatOpenAI(
        model="gpt-4",
        temperature=0.7
    )
    
    # Create the graph
    graph = StateGraph(ResearchState)
    
    # Add nodes
    graph.add_node("research", research_node)
    graph.add_node("evaluate", evaluate_node)
    graph.add_node("synthesize", synthesize_node)
    
    # Add edges
    graph.add_edge("research", "evaluate")
    graph.add_conditional_edges(
        "evaluate",
        should_continue,
        {
            "research": "research",
            "synthesize": "synthesize",
            "end": END
        }
    )
    graph.add_edge("synthesize", END)
    
    # Set entry point
    graph.set_entry_point("research")
    
    return graph.compile()

def run_iterative_research():
    """Run the iterative research agent with tracing."""
    print("=== Iterative Research Agent Tracing ===")
    
    # Create callback handler
    callback_handler = NoveumTraceCallbackHandler()
    
    # Create the agent
    agent = create_iterative_research_agent()
    
    # Run with callbacks via config (recommended approach)
    result = agent.invoke(
        {
            "messages": [HumanMessage(content="Research artificial intelligence and its applications")],
            "research_topic": "artificial intelligence and its applications",
            "research_notes": [],
            "evaluation_score": 0.0,
            "max_iterations": 3,
            "current_iteration": 1,
            "research_complete": False
        },
        config={
            "callbacks": [callback_handler],
            "tags": ["iterative_research"]
        }
    )
    
    print(f"\n🎉 Research completed!")
    print(f"📊 Final quality score: {result['evaluation_score']:.2f}")
    print(f"🔄 Total iterations: {result['current_iteration']}")
    print(f"📝 Research notes: {len(result['research_notes'])}")
    
    return result

if __name__ == "__main__":
    run_iterative_research()
    noveum_trace.flush()

pip install noveum-trace langchain-openai langgraph python-dotenv

export NOVEUM_API_KEY="your-noveum-api-key"
export OPENAI_API_KEY="your-openai-api-key"

def create_adaptive_research_agent():
    """Create an agent that adapts its research strategy."""
    llm = ChatOpenAI()
    
    def adaptive_research_node(state: ResearchState):
        """Adapt research strategy based on previous results."""
        # Analyze previous research to determine next steps
        if state["current_iteration"] > 1:
            # Look for gaps in previous research
            previous_queries = [note["query"] for note in state["research_notes"]]
            # Adapt search strategy based on gaps
            pass
        
        # Continue with research
        return research_node(state)
    
    # Rest of the implementation...

@tool
def search_academic(query: str) -> str:
    """Search academic databases."""
    return f"Academic search results for: {query}"

@tool
def search_news(query: str) -> str:
    """Search news sources."""
    return f"News search results for: {query}"

def multi_source_research_node(state: ResearchState):
    """Research using multiple sources."""
    # Search different sources
    academic_results = search_academic(state["research_topic"])
    news_results = search_news(state["research_topic"])
    web_results = search_web(state["research_topic"])
    
    # Combine results
    combined_analysis = f"""
    Academic: {academic_results}
    News: {news_results}
    Web: {web_results}
    """
    
    # Add to research notes
    state["research_notes"].append({
        "iteration": state["current_iteration"],
        "sources": ["academic", "news", "web"],
        "results": combined_analysis
    })
    
    return state