Build a Multi-Agent Research Platform with LangGraph and A2A

import os
import logging
from typing import TypedDict, List, Annotated, Union
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
import sqlite3
from langgraph.checkpoint.sqlite import SqliteSaver

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

# --- Environment Configuration ---
# Ensure OPENAI_API_KEY is set in your .env file or environment variables
if not os.environ.get("OPENAI_API_KEY"): # Using os.environ.get for safety
    logger.error("OPENAI_API_KEY environment variable not set. Please set it to proceed.")
    raise ValueError("OPENAI_API_KEY not set")

# --- LangGraph Agent State Definition ---
# This defines the state of our graph. It is typed for clarity and error checking.
class AgentState(TypedDict):
    input: str  # The initial user query
    chat_history: Annotated[List[BaseMessage], "append"]  # Conversation history
    research_results: Annotated[List[str], "append"]  # Accumulated research findings
    next_agent: str # The next agent to be invoked (e.g., 'web_search', 'academic', 'code', 'supervisor')
    current_task: str # The specific task being worked on by an agent

# --- LLM Initialization (for supervisor) ---
def get_llm():
    try:
        return ChatOpenAI(model="gpt-4o", temperature=0.5, streaming=True)
    except Exception as e:
        logger.error(f"Failed to initialize LLM: {e}")
        raise

# --- Main application entry point for testing setup ---
if __name__ == "__main__":
    print("
--- Phase 1: Environment Setup and Core Agent State ---")
    print("Checking environment variables...")
    try:
        llm = get_llm()
        print(f"LLM initialized successfully with model: {llm.model_name}")
    except ValueError as e:
        print(f"Error: {e}")
        print("Please ensure OPENAI_API_KEY is set in your environment or a .env file.")
        exit(1)

    print("
AgentState definition:")
    for key, value in AgentState.__annotations__.items():
        print(f"  {key}: {value}")
    print("
Phase 1 setup complete. Ready to define the graph.")

    # Example of how to use SqliteSaver (will be used in later phases)
    memory = SqliteSaver(conn=sqlite3.connect(":memory:", check_same_thread=False))
    print(f"Initialized in-memory SqliteSaver for future use: {memory}")

import os
import logging
from typing import TypedDict, List, Annotated, Union, Literal
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
import sqlite3
from langgraph.checkpoint.sqlite import SqliteSaver
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel
from langchain_core.output_parsers import JsonOutputParser

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

# --- Environment Configuration ---
if not os.environ.get("OPENAI_API_KEY"): # Using os.environ.get for safety
    logger.error("OPENAI_API_KEY environment variable not set. Please set it to proceed.")
    raise ValueError("OPENAI_API_KEY not set")

# --- LangGraph Agent State Definition ---
class AgentState(TypedDict):
    input: str  # The initial user query
    chat_history: Annotated[List[BaseMessage], "append"]  # Conversation history
    research_results: Annotated[List[str], "append"]  # Accumulated research findings
    next_agent: str # The next agent to be invoked (e.g., 'web_search', 'academic', 'code', 'supervisor')
    current_task: str # The specific task being worked on by an agent

# --- LLM Initialization ---
def get_llm():
    try:
        return ChatOpenAI(model="gpt-4o", temperature=0.5, streaming=True)
    except Exception as e:
        logger.error(f"Failed to initialize LLM: {e}")
        raise

# --- Supervisor Agent Logic ---
class AgentDecision(BaseModel):
    next_agent: str
    current_task: str

class SupervisorAgent:
    def __init__(self, llm: ChatOpenAI):
        self.llm = llm
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a research supervisor. Your goal is to break down complex research queries into sub-tasks and delegate them to specialist agents. "
                       "You have access to 'web_search_agent', 'academic_paper_agent', and 'code_execution_agent'. "
                       "Based on the user's input, decide which agent is best suited to perform the next step. "
                       "If the research is complete or you need to synthesize results, choose 'FINISH'.
"
                       "Respond with a JSON object containing 'next_agent' and 'current_task'.
"
                       "Available agents: web_search_agent, academic_paper_agent, code_execution_agent, FINISH"),
            ("human", "{input}")
        ])
        self.parser = JsonOutputParser(pydantic_object=AgentDecision)
        self.chain = self.prompt | self.llm | self.parser

    def route_agent(self, state: AgentState) -> AgentState:
        logger.info(f"Supervisor: Routing agent for input: {state['input']}")
        try:
            # Use the latest input or the last message in chat_history if input is empty
            current_input = state['input'] if state['input'] else state['chat_history'][-1].content
            response = self.chain.invoke({"input": current_input})
            next_agent = response.get('next_agent', 'FINISH')
            current_task = response.get('current_task', 'Synthesize findings.')
            logger.info(f"Supervisor decided: next_agent={next_agent}, current_task={current_task}")
            return {**state, "next_agent": next_agent, "current_task": current_task}
        except Exception as e:
            logger.error(f"Supervisor routing failed: {e}")
            # Fallback to FINISH in case of error to prevent infinite loops
            return {**state, "next_agent": "FINISH", "current_task": f"Error during routing: {e}"}

    def run_supervisor(self, state: AgentState) -> AgentState:
        logger.info("Supervisor: Running supervisor logic.")
        # In a more complex scenario, this might involve more reasoning or tool use
        # For now, it primarily focuses on routing.
        return self.route_agent(state)


# --- Define the graph --- (placeholder for specialist agents)
def web_search_agent_node(state: AgentState) -> AgentState:
    logger.info("Web Search Agent: Executing placeholder.")
    return {**state, "research_results": state["research_results"] + [f"Web search for: {state['current_task']}"], "next_agent": "supervisor"}

def academic_paper_agent_node(state: AgentState) -> AgentState:
    logger.info("Academic Paper Agent: Executing placeholder.")
    return {**state, "research_results": state["research_results"] + [f"Academic search for: {state['current_task']}"], "next_agent": "supervisor"}

def code_execution_agent_node(state: AgentState) -> AgentState:
    logger.info("Code Execution Agent: Executing placeholder.")
    return {**state, "research_results": state["research_results"] + [f"Code execution for: {state['current_task']}"], "next_agent": "supervisor"}

# --- Build the LangGraph --- 
def create_research_graph(llm: ChatOpenAI, memory: SqliteSaver):
    workflow = StateGraph(AgentState)

    supervisor = SupervisorAgent(llm)

    workflow.add_node("supervisor", supervisor.run_supervisor)
    workflow.add_node("web_search_agent", web_search_agent_node)
    workflow.add_node("academic_paper_agent", academic_paper_agent_node)
    workflow.add_node("code_execution_agent", code_execution_agent_node)

    workflow.set_entry_point("supervisor")

    workflow.add_conditional_edges(
        "supervisor",
        lambda state: state['next_agent'],
        {
            "web_search_agent": "web_search_agent",
            "academic_paper_agent": "academic_paper_agent",
            "code_execution_agent": "code_execution_agent",
            "FINISH": END
        },
    )

    # After specialist agents, return control to supervisor
    workflow.add_edge("web_search_agent", "supervisor")
    workflow.add_edge("academic_paper_agent", "supervisor")
    workflow.add_edge("code_execution_agent", "supervisor")

    app = workflow.compile(checkpointer=memory)
    return app

# --- Main application entry point for testing --- 
if __name__ == "__main__":
    print("
--- Phase 2: Supervisor Agent and Initial Graph Structure ---")
    try:
        llm = get_llm()
        memory = SqliteSaver(conn=sqlite3.connect(":memory:", check_same_thread=False))
        app = create_research_graph(llm, memory)

        # Example invocation
        initial_state = {"input": "Research the latest advancements in large language models for scientific discovery.", "chat_history": [], "research_results": [], "next_agent": "", "current_task": ""}
        print(f"
Invoking graph with input: {initial_state['input']}")
        for s in app.stream(initial_state):
            print(s)

        print("
Phase 2 setup complete. Supervisor can now route tasks.")
    except ValueError as e:
        print(f"Error: {e}")
        print("Please ensure OPENAI_API_KEY is set.")
    except Exception as e:
        logger.exception("An unexpected error occurred during Phase 2 execution.")

import os
import logging
import json
from typing import TypedDict, List, Annotated, Union, Literal
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
import sqlite3
from langgraph.checkpoint.sqlite import SqliteSaver
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel
from langchain_core.output_parsers import JsonOutputParser
from langchain_community.tools.tavily_research import TavilySearchResults
from langchain_community.document_loaders import ArxivLoader
from langchain.tools import tool
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from functools import lru_cache
import subprocess
import sys

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

# --- Environment Configuration ---
required_env_vars = ["OPENAI_API_KEY", "TAVILY_API_KEY"]
for var in required_env_vars:
    if not os.environ.get(var):
        logger.error(f"{var} environment variable not set. Please set it to proceed.")
        raise ValueError(f"{var} not set")

# --- LangGraph Agent State Definition ---
class AgentState(TypedDict):
    input: str  # The initial user query
    chat_history: Annotated[List[BaseMessage], "append"]  # Conversation history
    research_results: Annotated[List[str], "append"]  # Accumulated research findings
    next_agent: str # The next agent to be invoked (e.g., 'web_search', 'academic', 'code', 'supervisor')
    current_task: str # The specific task being worked on by an agent

# --- LLM Initialization ---
def get_llm():
    try:
        return ChatOpenAI(model="gpt-4o", temperature=0.5, streaming=True)
    except Exception as e:
        logger.error(f"Failed to initialize LLM: {e}")
        raise

# --- Episodic Memory Setup ---
@lru_cache(maxsize=1) # Cache the vector store instance
def get_episodic_memory():
    try:
        embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
        # Using a local ChromaDB instance
        vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
        return vectorstore
    except Exception as e:
        logger.error(f"Failed to initialize ChromaDB for episodic memory: {e}")
        raise

def add_to_episodic_memory(query: str, content: str):
    try:
        vectorstore = get_episodic_memory()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        docs = text_splitter.create_documents([content], metadatas=[{"query": query}])
        vectorstore.add_documents(docs)
        logger.info(f"Added {len(docs)} documents to episodic memory for query: {query[:50]}...")
    except Exception as e:
        logger.error(f"Error adding to episodic memory: {e}")

def retrieve_from_episodic_memory(query: str, k: int = 3) -> List[str]:
    try:
        vectorstore = get_episodic_memory()
        retrieved_docs = vectorstore.similarity_search(query, k=k)
        logger.info(f"Retrieved {len(retrieved_docs)} documents from episodic memory for query: {query[:50]}...")
        return [doc.page_content for doc in retrieved_docs]
    except Exception as e:
        logger.error(f"Error retrieving from episodic memory: {e}")
        return []

# --- Supervisor Agent Logic (from Phase 2, updated to use memory) ---
class AgentDecision(BaseModel):
    next_agent: str
    current_task: str

class SupervisorAgent:
    def __init__(self, llm: ChatOpenAI):
        self.llm = llm
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a research supervisor. Your goal is to break down complex research queries into sub-tasks and delegate them to specialist agents. "
                       "You have access to 'web_search_agent', 'academic_paper_agent', and 'code_execution_agent'. "
                       "Based on the user's input and any existing research results, decide which agent is best suited to perform the next step. "
                       "If the research is complete or you need to synthesize results, choose 'FINISH'.
"
                       "Consider the existing research results and episodic memory before delegating. Focus on new, unaddressed aspects.
"
                       "Respond with a JSON object containing 'next_agent' and 'current_task'.
"
                       "Available agents: web_search_agent, academic_paper_agent, code_execution_agent, FINISH
"
                       "Existing Research Results: {research_results}
"
                       "Relevant Episodic Memory: {episodic_memory}"),
            ("human", "{input}")
        ])
        self.parser = JsonOutputParser(pydantic_object=AgentDecision)
        self.chain = self.prompt | self.llm | self.parser

    def route_agent(self, state: AgentState) -> AgentState:
        logger.info(f"Supervisor: Routing agent for input: {state['input']}")
        try:
            current_input = state['input'] if state['input'] else state['chat_history'][-1].content
            # Retrieve relevant info from episodic memory
            relevant_memory = retrieve_from_episodic_memory(current_input + " ".join(state['research_results']))

            response = self.chain.invoke({
                "input": current_input,
                "research_results": "
".join(state['research_results']) if state['research_results'] else "None yet.",
                "episodic_memory": "
".join(relevant_memory) if relevant_memory else "None found."
            })
            next_agent = response.get('next_agent', 'FINISH')
            current_task = response.get('current_task', 'Synthesize findings.')
            logger.info(f"Supervisor decided: next_agent={next_agent}, current_task={current_task}")
            return {**state, "next_agent": next_agent, "current_task": current_task}
        except Exception as e:
            logger.error(f"Supervisor routing failed: {e}")
            return {**state, "next_agent": "FINISH", "current_task": f"Error during routing: {e}"}

    def run_supervisor(self, state: AgentState) -> AgentState:
        logger.info("Supervisor: Running supervisor logic.")
        return self.route_agent(state)

# --- Specialist Agent Implementations ---
# Web Search Agent
class WebSearchAgent:
    def __init__(self):
        self.tool = TavilySearchResults(max_results=5) # Reduced for faster execution

    def run_web_search(self, state: AgentState) -> AgentState:
        logger.info(f"Web Search Agent: Searching for '{state['current_task']}'")
        try:
            search_results = self.tool.invoke({"query": state['current_task']})
            result_str = json.dumps(search_results, indent=2)
            add_to_episodic_memory(state['current_task'], result_str)
            return {**state, "research_results": state["research_results"] + [f"Web Search Results for '{state['current_task']}':
{result_str}"], "next_agent": "supervisor"}
        except Exception as e:
            logger.error(f"Web Search Agent failed: {e}")
            return {**state, "research_results": state["research_results"] + [f"Web Search failed for '{state['current_task']}': {e}"], "next_agent": "supervisor"}

# Academic Paper Agent
class AcademicPaperAgent:
    def run_academic_search(self, state: AgentState) -> AgentState:
        logger.info(f"Academic Paper Agent: Searching ArXiv for '{state['current_task']}'")
        try:
            loader = ArxivLoader(query=state['current_task'], load_max_docs=2, load_all_available_meta=True)
            docs = loader.load()
            if not docs:
                return {**state, "research_results": state["research_results"] + [f"No academic papers found for '{state['current_task']}'."], "next_agent": "supervisor"}

            result_summary = []
            for doc in docs:
                summary = f"Title: {doc.metadata.get('Title', 'N/A')}
Authors: {doc.metadata.get('Authors', 'N/A')}
Published: {doc.metadata.get('Published', 'N/A')}
Summary: {doc.page_content[:500]}...
Link: {doc.metadata.get('entry_id', 'N/A')}"
                result_summary.append(summary)
            
            full_result_str = "
---
".join(result_summary)
            add_to_episodic_memory(state['current_task'], full_result_str)
            return {**state, "research_results": state["research_results"] + [f"Academic Research Results for '{state['current_task']}':
{full_result_str}"], "next_agent": "supervisor"}
        except Exception as e:
            logger.error(f"Academic Paper Agent failed: {e}")
            return {**state, "research_results": state["research_results"] + [f"Academic Search failed for '{state['current_task']}': {e}"], "next_agent": "supervisor"}

# Code Execution Agent
class CodeExecutionAgent:
    @tool("python_interpreter", return_direct=True)
    def python_interpreter(self, code: str) -> str:
        "Executes Python code in a sandboxed environment."
        try:
            # For production, use a secure sandbox like E2B or a Docker container
            # For this example, we'll use a basic subprocess for demonstration.
            # WARNING: Executing arbitrary code locally is a security risk.
            logger.warning("Executing code locally without a proper sandbox. Use with caution.")
            process = subprocess.run(
                [sys.executable, "-c", code],
                capture_output=True,
                text=True,
                check=True,
                timeout=30
            )
            if process.stderr:
                logger.warning(f"Code execution stderr: {process.stderr}")
            return process.stdout.strip()
        except subprocess.CalledProcessError as e:
            error_msg = f"Code execution failed with error: {e.stderr.strip()}"
            logger.error(error_msg)
            raise ValueError(error_msg)
        except subprocess.TimeoutExpired:
            error_msg = "Code execution timed out."
            logger.error(error_msg)
            raise ValueError(error_msg)
        except Exception as e:
            error_msg = f"An unexpected error occurred during code execution: {e}"
            logger.error(error_msg)
            raise ValueError(error_msg)

    def run_code_execution(self, state: AgentState) -> AgentState:
        logger.info(f"Code Execution Agent: Attempting to execute code for task '{state['current_task']}'")
        try:
            # LLM decides the code to execute based on current_task
            code_generation_prompt = ChatPromptTemplate.from_messages([
                ("system", "You are a helpful assistant capable of writing and executing Python code to answer questions or perform data analysis. "
                           "Given a task, generate valid Python code that can be executed. Output ONLY the Python code, no explanations or markdown fences."
                           "Current Task: {current_task}
"
                           "Existing Research Results: {research_results}
"
                           "Relevant Episodic Memory: {episodic_memory}
"
                           "Example: print(2+2)"),
                ("human", "Generate Python code for the following task: {task}")
            ])
            llm_code_chain = code_generation_prompt | get_llm()

            relevant_memory = retrieve_from_episodic_memory(state['current_task'] + " ".join(state['research_results']))
            
            code_to_execute = llm_code_chain.invoke({
                "task": state['current_task'],
                "research_results": "
".join(state['research_results']) if state['research_results'] else "None yet.",
                "episodic_memory": "
".join(relevant_memory) if relevant_memory else "None found."
            }).content
            logger.info(f"Generated code: {code_to_execute[:100]}...")

            execution_output = self.python_interpreter.invoke(code_to_execute)
            add_to_episodic_memory(state['current_task'], f"Code Output: {execution_output}")
            return {**state, "research_results": state["research_results"] + [f"Code Execution Results for '{state['current_task']}':
{execution_output}"], "next_agent": "supervisor"}
        except Exception as e:
            logger.error(f"Code Execution Agent failed: {e}")
            return {**state, "research_results": state["research_results"] + [f"Code Execution failed for '{state['current_task']}': {e}"], "next_agent": "supervisor"}


# --- Build the LangGraph --- 
def create_research_graph(llm: ChatOpenAI, memory: SqliteSaver):
    workflow = StateGraph(AgentState)

    supervisor = SupervisorAgent(llm)
    web_search = WebSearchAgent()
    academic_paper = AcademicPaperAgent()
    code_execution = CodeExecutionAgent()

    workflow.add_node("supervisor", supervisor.run_supervisor)
    workflow.add_node("web_search_agent", web_search.run_web_search)
    workflow.add_node("academic_paper_agent", academic_paper.run_academic_search)
    workflow.add_node("code_execution_agent", code_execution.run_code_execution)

    workflow.set_entry_point("supervisor")

    workflow.add_conditional_edges(
        "supervisor",
        lambda state: state['next_agent'],
        {
            "web_search_agent": "web_search_agent",
            "academic_paper_agent": "academic_paper_agent",
            "code_execution_agent": "code_execution_agent",
            "FINISH": END
        },
    )

    workflow.add_edge("web_search_agent", "supervisor")
    workflow.add_edge("academic_paper_agent", "supervisor")
    workflow.add_edge("code_execution_agent", "supervisor")

    app = workflow.compile(checkpointer=memory)
    return app

# --- Main application entry point for testing --- 
if __name__ == "__main__":
    print("
--- Phase 3: Integrating Specialist Tools and Episodic Memory ---")
    try:
        llm = get_llm()
        memory = SqliteSaver(conn=sqlite3.connect("./langgraph_checkpoints.sqlite", check_same_thread=False)) # Persist checkpoints
        app = create_research_graph(llm, memory)

        # Create chroma_db directory if it doesn't exist
        if not os.path.exists("./chroma_db"):
            os.makedirs("./chroma_db")
            logger.info("Created ./chroma_db directory for episodic memory.")

        # Example invocation
        initial_state = {"input": "Find recent papers on graph neural networks for drug discovery and summarize their key findings. Then, write and execute a Python script to calculate the average number of citations for the top 3 papers.", "chat_history": [], "research_results": [], "next_agent": "", "current_task": ""}
        print(f"
Invoking graph with input: {initial_state['input']}")
        for s in app.stream(initial_state):
            print(s)

        print("
Phase 3 setup complete. Specialist agents are now functional with tools and episodic memory.")
    except ValueError as e:
        print(f"Error: {e}")
        print("Please ensure all required API keys are set.")
    except Exception as e:
        logger.exception("An unexpected error occurred during Phase 3 execution.")

import os
import logging
import json
import time
from typing import TypedDict, List, Annotated, Union, Literal, AsyncIterator, Dict, Any
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
import sqlite3
from langgraph.checkpoint.sqlite import SqliteSaver
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_community.tools.tavily_research import TavilySearchResults
from langchain_community.document_loaders import ArxivLoader
from langchain.tools import tool
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from functools import lru_cache
import subprocess
import sys
import uuid
import asyncio
import uvicorn
from fastapi import FastAPI, Request, Response
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import redis.asyncio as redis
from redis.asyncio import Redis

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

# --- Environment Configuration ---
required_env_vars = ["OPENAI_API_KEY", "TAVILY_API_KEY", "REDIS_URL"]
for var in required_env_vars:
    if not os.environ.get(var):
        logger.error(f"{var} environment variable not set. Please set it to proceed.")
        raise ValueError(f"{var} not set")

# --- LangGraph Agent State Definition (updated with session_id) ---
class AgentState(TypedDict):
    input: str  # The initial user query
    chat_history: Annotated[List[BaseMessage], "append"]  # Conversation history
    research_results: Annotated[List[str], "append"]  # Accumulated research findings
    next_agent: str # The next agent to be invoked (e.g., 'web_search', 'academic', 'code', 'supervisor')
    current_task: str # The specific task being worked on by an agent
    session_id: str # Unique ID for the current research session
    a2a_messages: Annotated[List[Dict[str, Any]], "append"] # A2A messages for inter-agent communication

# --- LLM Initialization ---
def get_llm():
    try:
        return ChatOpenAI(model="gpt-4o", temperature=0.5, streaming=True)
    except Exception as e:
        logger.error(f"Failed to initialize LLM: {e}")
        raise

# --- Episodic Memory Setup (from Phase 3) ---
@lru_cache(maxsize=1)
def get_episodic_memory():
    try:
        embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
        vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
        return vectorstore
    except Exception as e:
        logger.error(f"Failed to initialize ChromaDB for episodic memory: {e}")
        raise

def add_to_episodic_memory(query: str, content: str):
    try:
        vectorstore = get_episodic_memory()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        docs = text_splitter.create_documents([content], metadatas=[{"query": query}])
        vectorstore.add_documents(docs)
        logger.info(f"Added {len(docs)} documents to episodic memory for query: {query[:50]}...")
    except Exception as e:
        logger.error(f"Error adding to episodic memory: {e}")

def retrieve_from_episodic_memory(query: str, k: int = 3) -> List[str]:
    try:
        vectorstore = get_episodic_memory()
        retrieved_docs = vectorstore.similarity_search(query, k=k)
        logger.info(f"Retrieved {len(retrieved_docs)} documents from episodic memory for query: {query[:50]}...")
        return [doc.page_content for doc in retrieved_docs]
    except Exception as e:
        logger.error(f"Error retrieving from episodic memory: {e}")
        return []

# --- A2A Message Queue (Redis) ---
@lru_cache(maxsize=1)
async def get_redis_client() -> Redis:
    try:
        redis_url = os.environ.get("REDIS_URL", "redis://localhost:6379/0")
        r = redis.from_url(redis_url)
        await r.ping() # Test connection
        logger.info(f"Connected to Redis at {redis_url}")
        return r
    except Exception as e:
        logger.error(f"Failed to connect to Redis: {e}")
        raise

async def publish_a2a_message(session_id: str, sender: str, recipient: str, message_type: str, content: Dict[str, Any]):
    try:
        r = await get_redis_client()
        a2a_msg = {
            "session_id": session_id,
            "timestamp": time.time(),
            "sender": sender,
            "recipient": recipient,
            "message_type": message_type,
            "content": content
        }
        await r.publish(f"a2a_channel:{session_id}", json.dumps(a2a_msg))
        logger.info(f"Published A2A message from {sender} to {recipient} ({message_type}) for session {session_id}")
    except Exception as e:
        logger.error(f"Error publishing A2A message: {e}")

# --- Supervisor Agent Logic (updated for A2A) ---
class AgentDecision(BaseModel):
    next_agent: str
    current_task: str

class SupervisorAgent:
    def __init__(self, llm: ChatOpenAI):
        self.llm = llm
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a research supervisor. Your goal is to break down complex research queries into sub-tasks and delegate them to specialist agents. "
                       "You have access to 'web_search_agent', 'academic_paper_agent', and 'code_execution_agent'. "
                       "Based on the user's input and any existing research results, decide which agent is best suited to perform the next step. "
                       "If the research is complete or you need to synthesize results, choose 'FINISH'.
"
                       "Consider the existing research results and episodic memory before delegating. Focus on new, unaddressed aspects.
"
                       "Respond with a JSON object containing 'next_agent' and 'current_task'.
"
                       "Available agents: web_search_agent, academic_paper_agent, code_execution_agent, FINISH
"
                       "Existing Research Results: {research_results}
"
                       "Relevant Episodic Memory: {episodic_memory}"),
            ("human", "{input}")
        ])
        self.parser = JsonOutputParser(pydantic_object=AgentDecision)
        self.chain = self.prompt | self.llm | self.parser

    async def route_agent(self, state: AgentState) -> AgentState:
        logger.info(f"Supervisor: Routing agent for input: {state['input']}")
        try:
            current_input = state['input'] if state['input'] else state['chat_history'][-1].content
            relevant_memory = retrieve_from_episodic_memory(current_input + " ".join(state['research_results']))

            response = self.chain.invoke({
                "input": current_input,
                "research_results": "
".join(state['research_results']) if state['research_results'] else "None yet.",
                "episodic_memory": "
".join(relevant_memory) if relevant_memory else "None found."
            })
            next_agent = response.get('next_agent', 'FINISH')
            current_task = response.get('current_task', 'Synthesize findings.')
            logger.info(f"Supervisor decided: next_agent={next_agent}, current_task={current_task}")

            # Publish A2A message about delegation
            await publish_a2a_message(
                state['session_id'], "supervisor", next_agent,
                "task_delegation", {"task": current_task, "input": current_input}
            )

            return {**state, "next_agent": next_agent, "current_task": current_task}
        except Exception as e:
            logger.error(f"Supervisor routing failed: {e}")
            await publish_a2a_message(
                state['session_id'], "supervisor", "user",
                "error", {"message": f"Supervisor routing failed: {e}"}
            )
            return {**state, "next_agent": "FINISH", "current_task": f"Error during routing: {e}"}

    async def run_supervisor(self, state: AgentState) -> AgentState:
        logger.info("Supervisor: Running supervisor logic.")
        return await self.route_agent(state)

# --- Specialist Agent Implementations (updated for A2A) ---
class WebSearchAgent:
    def __init__(self):
        self.tool = TavilySearchResults(max_results=5)

    async def run_web_search(self, state: AgentState) -> AgentState:
        logger.info(f"Web Search Agent: Searching for '{state['current_task']}'")
        try:
            await publish_a2a_message(
                state['session_id'], "web_search_agent", "supervisor",
                "status_update", {"message": f"Starting web search for: {state['current_task']}"}
            )
            search_results = self.tool.invoke({"query": state['current_task']})
            result_str = json.dumps(search_results, indent=2)
            add_to_episodic_memory(state['current_task'], result_str)
            await publish_a2a_message(
                state['session_id'], "web_search_agent", "supervisor",
                "result", {"task": state['current_task'], "output": result_str}
            )
            return {**state, "research_results": state["research_results"] + [f"Web Search Results for '{state['current_task']}':
{result_str}"], "next_agent": "supervisor"}
        except Exception as e:
            logger.error(f"Web Search Agent failed: {e}")
            await publish_a2a_message(
                state['session_id'], "web_search_agent", "supervisor",
                "error", {"task": state['current_task'], "message": f"Web Search failed: {e}"}
            )
            return {**state, "research_results": state["research_results"] + [f"Web Search failed for '{state['current_task']}': {e}"], "next_agent": "supervisor"}

class AcademicPaperAgent:
    async def run_academic_search(self, state: AgentState) -> AgentState:
        logger.info(f"Academic Paper Agent: Searching ArXiv for '{state['current_task']}'")
        try:
            await publish_a2a_message(
                state['session_id'], "academic_paper_agent", "supervisor",
                "status_update", {"message": f"Starting academic search for: {state['current_task']}"}
            )
            loader = ArxivLoader(query=state['current_task'], load_max_docs=2, load_all_available_meta=True)
            docs = loader.load()
            if not docs:
                await publish_a2a_message(
                    state['session_id'], "academic_paper_agent", "supervisor",
                    "result", {"task": state['current_task'], "output": f"No academic papers found for '{state['current_task']}'."}
                )
                return {**state, "research_results": state["research_results"] + [f"No academic papers found for '{state['current_task']}'."], "next_agent": "supervisor"}

            result_summary = []
            for doc in docs:
                summary = f"Title: {doc.metadata.get('Title', 'N/A')}
Authors: {doc.metadata.get('Authors', 'N/A')}
Published: {doc.metadata.get('Published', 'N/A')}
Summary: {doc.page_content[:500]}...
Link: {doc.metadata.get('entry_id', 'N/A')}"
                result_summary.append(summary)
            
            full_result_str = "
---
".join(result_summary)
            add_to_episodic_memory(state['current_task'], full_result_str)
            await publish_a2a_message(
                state['session_id'], "academic_paper_agent", "supervisor",
                "result", {"task": state['current_task'], "output": full_result_str}
            )
            return {**state, "research_results": state["research_results"] + [f"Academic Research Results for '{state['current_task']}':
{full_result_str}"], "next_agent": "supervisor"}
        except Exception as e:
            logger.error(f"Academic Paper Agent failed: {e}")
            await publish_a2a_message(
                state['session_id'], "academic_paper_agent", "supervisor",
                "error", {"task": state['current_task'], "message": f"Academic Search failed: {e}"}
            )
            return {**state, "research_results": state["research_results"] + [f"Academic Search failed for '{state['current_task']}': {e}"], "next_agent": "supervisor"}

class CodeExecutionAgent:
    @tool("python_interpreter", return_direct=True)
    def python_interpreter(self, code: str) -> str:
        "Executes Python code in a sandboxed environment."
        try:
            logger.warning("Executing code locally without a proper sandbox. Use with caution.")
            process = subprocess.run(
                [sys.executable, "-c", code],
                capture_output=True,
                text=True,
                check=True,
                timeout=30
            )
            if process.stderr:
                logger.warning(f"Code execution stderr: {process.stderr}")
            return process.stdout.strip()
        except subprocess.CalledProcessError as e:
            error_msg = f"Code execution failed with error: {e.stderr.strip()}"
            logger.error(error_msg)
            raise ValueError(error_msg)
        except subprocess.TimeoutExpired:
            error_msg = "Code execution timed out."
            logger.error(error_msg)
            raise ValueError(error_msg)
        except Exception as e:
            error_msg = f"An unexpected error occurred during code execution: {e}"
            logger.error(error_msg)
            raise ValueError(error_msg)

    async def run_code_execution(self, state: AgentState) -> AgentState:
        logger.info(f"Code Execution Agent: Attempting to execute code for task '{state['current_task']}'")
        try:
            await publish_a2a_message(
                state['session_id'], "code_execution_agent", "supervisor",
                "status_update", {"message": f"Starting code execution for: {state['current_task']}"}
            )
            code_generation_prompt = ChatPromptTemplate.from_messages([
                ("system", "You are a helpful assistant capable of writing and executing Python code to answer questions or perform data analysis. "
                           "Given a task, generate valid Python code that can be executed. Output ONLY the Python code, no explanations or markdown fences."
                           "Current Task: {current_task}
"
                           "Existing Research Results: {research_results}
"
                           "Relevant Episodic Memory: {episodic_memory}
"
                           "Example: print(2+2)"),
                ("human", "Generate Python code for the following task: {task}")
            ])
            llm_code_chain = code_generation_prompt | get_llm()

            relevant_memory = retrieve_from_episodic_memory(state['current_task'] + " ".join(state['research_results']))
            
            code_to_execute = llm_code_chain.invoke({
                "task": state['current_task'],
                "research_results": "
".join(state['research_results']) if state['research_results'] else "None yet.",
                "episodic_memory": "
".join(relevant_memory) if relevant_memory else "None found."
            }).content
            logger.info(f"Generated code: {code_to_execute[:100]}...")

            execution_output = self.python_interpreter.invoke(code_to_execute)
            add_to_episodic_memory(state['current_task'], f"Code Output: {execution_output}")
            await publish_a2a_message(
                state['session_id'], "code_execution_agent", "supervisor",
                "result", {"task": state['current_task'], "output": execution_output}
            )
            return {**state, "research_results": state["research_results"] + [f"Code Execution Results for '{state['current_task']}':
{execution_output}"], "next_agent": "supervisor"}
        except Exception as e:
            logger.error(f"Code Execution Agent failed: {e}")
            await publish_a2a_message(
                state['session_id'], "code_execution_agent", "supervisor",
                "error", {"task": state['current_task'], "message": f"Code Execution failed: {e}"}
            )
            return {**state, "research_results": state["research_results"] + [f"Code Execution failed for '{state['current_task']}': {e}"], "next_agent": "supervisor"}


# --- Build the LangGraph --- 
def create_research_graph(llm: ChatOpenAI, memory: SqliteSaver):
    workflow = StateGraph(AgentState)

    supervisor = SupervisorAgent(llm)
    web_search = WebSearchAgent()
    academic_paper = AcademicPaperAgent()
    code_execution = CodeExecutionAgent()

    # Note: LangGraph nodes can be sync or async. For simplicity, we make all nodes async here.
    # LangGraph will handle running async nodes correctly.
    workflow.add_node("supervisor", supervisor.run_supervisor)
    workflow.add_node("web_search_agent", web_search.run_web_search)
    workflow.add_node("academic_paper_agent", academic_paper.run_academic_search)
    workflow.add_node("code_execution_agent", code_execution.run_code_execution)

    workflow.set_entry_point("supervisor")

    workflow.add_conditional_edges(
        "supervisor",
        lambda state: state['next_agent'],
        {
            "web_search_agent": "web_search_agent",
            "academic_paper_agent": "academic_paper_agent",
            "code_execution_agent": "code_execution_agent",
            "FINISH": END
        },
    )

    workflow.add_edge("web_search_agent", "supervisor")
    workflow.add_edge("academic_paper_agent", "supervisor")
    workflow.add_edge("code_execution_agent", "supervisor")

    app = workflow.compile(checkpointer=memory)
    return app

# --- FastAPI Application for Streaming --- 
app_fastapi = FastAPI(title="Multi-Agent Research Platform")

class ResearchRequest(BaseModel):
    query: str
    session_id: str = str(uuid.uuid4()) # Generate a new session_id if not provided

async def stream_a2a_messages(session_id: str) -> AsyncIterator[str]:
    r = await get_redis_client()
    pubsub = r.pubsub()
    await pubsub.subscribe(f"a2a_channel:{session_id}")
    logger.info(f"Subscribed to A2A channel: a2a_channel:{session_id}")

    try:
        while True:
            message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0)
            if message:
                data = json.loads(message['data'].decode('utf-8'))
                yield f"data: {json.dumps(data)}

"
            await asyncio.sleep(0.01) # Yield control
    except asyncio.CancelledError:
        logger.info(f"Streaming for session {session_id} cancelled.")
    except Exception as e:
        logger.error(f"Error in A2A message streaming for session {session_id}: {e}")
        yield f"data: {json.dumps({'event': 'error', 'message': str(e)})}

"
    finally:
        await pubsub.unsubscribe(f"a2a_channel:{session_id}")
        logger.info(f"Unsubscribed from A2A channel: a2a_channel:{session_id}")

@app_fastapi.post("/research/stream")
async def start_research_stream(request: ResearchRequest):
    session_id = request.session_id
    query = request.query
    logger.info(f"Received research request for session {session_id} with query: {query}")

    llm = get_llm()
    memory = SqliteSaver(conn=sqlite3.connect(f"./langgraph_checkpoints_{session_id}.sqlite", check_same_thread=False)) # Per-session checkpoint
    research_graph = create_research_graph(llm, memory)

    async def research_task():
        try:
            initial_state = {"input": query, "chat_history": [HumanMessage(content=query)], "research_results": [], "next_agent": "", "current_task": "", "session_id": session_id}
            final_state = None
            async for s in research_graph.astream(initial_state, config={"configurable": {"thread_id": session_id}}):
                # This loop processes LangGraph state changes but does not directly stream them out.
                # A2A messages are published independently to Redis.
                if "supervisor" in s:
                    current_state = s["supervisor"]
                    if current_state.get("next_agent") == "FINISH":
                        final_state = current_state
                        break
            
            if final_state:
                await publish_a2a_message(
                    session_id, "supervisor", "user",
                    "final_result", {"input": query, "final_results": final_state['research_results']}
                )
            else:
                await publish_a2a_message(
                    session_id, "supervisor", "user",
                    "error", {"message": "Research task completed without a final state.", "input": query}
                )
        except Exception as e:
            logger.error(f"Error during research task for session {session_id}: {e}")
            await publish_a2a_message(
                session_id, "supervisor", "user",
                "error", {"message": f"Research task failed: {e}", "input": query}
            )

    # Run the research task in the background
    asyncio.create_task(research_task())

    # Stream A2A messages to the client
    return StreamingResponse(stream_a2a_messages(session_id), media_type="text/event-stream")

# --- Main application entry point for testing (FastAPI) --- 
if __name__ == "__main__":
    print("
--- Phase 4: A2A Protocol Integration and Streaming API ---")
    print("Starting FastAPI server... access at http://127.0.0.1:8000/docs")
    print("Ensure Redis is running (e.g., `docker run -p 6379:6379 --name my-redis -d redis`).")
    uvicorn.run(app_fastapi, host="0.0.0.0", port=8000)

import os
import logging
import json
import time
from typing import TypedDict, List, Annotated, Union, Literal, AsyncIterator, Dict, Any
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
import sqlite3
from langgraph.checkpoint.sqlite import SqliteSaver
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_community.tools.tavily_research import TavilySearchResults
from langchain_community.document_loaders import ArxivLoader
from langchain.tools import tool
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from functools import lru_cache
import subprocess
import sys
import uuid
import asyncio
import uvicorn
from fastapi import FastAPI, Request, Response
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import redis.asyncio as redis
from redis.asyncio import Redis

# OpenTelemetry imports
from opentelemetry import trace
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.langchain import LangchainInstrumentor
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor

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

# --- Environment Configuration ---
required_env_vars = ["OPENAI_API_KEY", "TAVILY_API_KEY", "REDIS_URL"]
for var in required_env_vars:
    if not os.environ.get(var):
        logger.error(f"{var} environment variable not set. Please set it to proceed.")
        raise ValueError(f"{var} not set")

# LangSmith specific environment variables (optional but recommended)
os.environ["LANGCHAIN_TRACING_V2"] = os.environ.get("LANGCHAIN_TRACING_V2", "true")
os.environ["LANGCHAIN_PROJECT"] = os.environ.get("LANGCHAIN_PROJECT", "Multi-Agent Research Platform")
# LANGCHAIN_API_KEY should be set in .env or environment variables

# --- OpenTelemetry Setup ---
def setup_opentelemetry():
    resource = Resource.create({"service.name": "multi-agent-research-platform"})
    provider = TracerProvider(resource=resource)
    # Configure OTLP HTTP exporter for traces
    exporter = OTLPSpanExporter(endpoint=os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318/v1/traces"))
    span_processor = BatchSpanProcessor(exporter)
    provider.add_span_processor(span_processor)
    trace.set_tracer_provider(provider)
    logger.info("OpenTelemetry configured.")

    # Instrument LangChain/LangGraph
    LangchainInstrumentor().instrument()
    logger.info("LangChain/LangGraph instrumented for OpenTelemetry.")

# --- LangGraph Agent State Definition (updated with session_id) ---
class AgentState(TypedDict):
    input: str
    chat_history: Annotated[List[BaseMessage], "append"]
    research_results: Annotated[List[str], "append"]
    next_agent: str
    current_task: str
    session_id: str
    a2a_messages: Annotated[List[Dict[str, Any]], "append"]

# --- LLM Initialization ---
def get_llm():
    try:
        return ChatOpenAI(model="gpt-4o", temperature=0.5, streaming=True)
    except Exception as e:
        logger.error(f"Failed to initialize LLM: {e}")
        raise

# --- Episodic Memory Setup (from Phase 3) ---
@lru_cache(maxsize=1)
def get_episodic_memory():
    try:
        embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
        vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
        return vectorstore
    except Exception as e:
        logger.error(f"Failed to initialize ChromaDB for episodic memory: {e}")
        raise

def add_to_episodic_memory(query: str, content: str):
    try:
        vectorstore = get_episodic_memory()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        docs = text_splitter.create_documents([content], metadatas=[{"query": query}])
        vectorstore.add_documents(docs)
        logger.info(f"Added {len(docs)} documents to episodic memory for query: {query[:50]}...")
    except Exception as e:
        logger.error(f"Error adding to episodic memory: {e}")

def retrieve_from_episodic_memory(query: str, k: int = 3) -> List[str]:
    try:
        vectorstore = get_episodic_memory()
        retrieved_docs = vectorstore.similarity_search(query, k=k)
        logger.info(f"Retrieved {len(retrieved_docs)} documents from episodic memory for query: {query[:50]}...")
        return [doc.page_content for doc in retrieved_docs]
    except Exception as e:
        logger.error(f"Error retrieving from episodic memory: {e}")
        return []

# --- A2A Message Queue (Redis) ---
@lru_cache(maxsize=1)
async def get_redis_client() -> Redis:
    try:
        redis_url = os.environ.get("REDIS_URL", "redis://localhost:6379/0")
        r = redis.from_url(redis_url)
        await r.ping() # Test connection
        logger.info(f"Connected to Redis at {redis_url}")
        return r
    except Exception as e:
        logger.error(f"Failed to connect to Redis: {e}")
        raise

async def publish_a2a_message(session_id: str, sender: str, recipient: str, message_type: str, content: Dict[str, Any]):
    try:
        r = await get_redis_client()
        a2a_msg = {
            "session_id": session_id,
            "timestamp": time.time(),
            "sender": sender,
            "recipient": recipient,
            "message_type": message_type,
            "content": content
        }
        await r.publish(f"a2a_channel:{session_id}", json.dumps(a2a_msg))
        logger.info(f"Published A2A message from {sender} to {recipient} ({message_type}) for session {session_id}")
    except Exception as e:
        logger.error(f"Error publishing A2A message: {e}")

# --- Supervisor Agent Logic (from Phase 4) ---
class AgentDecision(BaseModel):
    next_agent: str
    current_task: str

class SupervisorAgent:
    def __init__(self, llm: ChatOpenAI):
        self.llm = llm
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a research supervisor. Your goal is to break down complex research queries into sub-tasks and delegate them to specialist agents. "
                       "You have access to 'web_search_agent', 'academic_paper_agent', and 'code_execution_agent'. "
                       "Based on the user's input and any existing research results, decide which agent is best suited to perform the next step. "
                       "If the research is complete or you need to synthesize results, choose 'FINISH'.
"
                       "Consider the existing research results and episodic memory before delegating. Focus on new, unaddressed aspects.
"
                       "Respond with a JSON object containing 'next_agent' and 'current_task'.
"
                       "Available agents: web_search_agent, academic_paper_agent, code_execution_agent, FINISH
"
                       "Existing Research Results: {research_results}
"
                       "Relevant Episodic Memory: {episodic_memory}"),
            ("human", "{input}")
        ])
        self.parser = JsonOutputParser(pydantic_object=AgentDecision)
        self.chain = self.prompt | self.llm | self.parser

    async def route_agent(self, state: AgentState) -> AgentState:
        logger.info(f"Supervisor: Routing agent for input: {state['input']}")
        with trace.get_tracer(__name__).start_as_current_span("supervisor.route_agent") as span:
            span.set_attribute("session_id", state['session_id'])
            span.set_attribute("current_task", state['current_task'])
            try:
                current_input = state['input'] if state['input'] else state['chat_history'][-1].content
                relevant_memory = retrieve_from_episodic_memory(current_input + " ".join(state['research_results']))

                response = self.chain.invoke({
                    "input": current_input,
                    "research_results": "
".join(state['research_results']) if state['research_results'] else "None yet.",
                    "episodic_memory": "
".join(relevant_memory) if relevant_memory else "None found."
                })
                next_agent = response.get('next_agent', 'FINISH')
                current_task = response.get('current_task', 'Synthesize findings.')
                logger.info(f"Supervisor decided: next_agent={next_agent}, current_task={current_task}")
                span.set_attribute("next_agent_decision", next_agent)
                span.set_attribute("delegated_task", current_task)

                await publish_a2a_message(
                    state['session_id'], "supervisor", next_agent,
                    "task_delegation", {"task": current_task, "input": current_input}
                )

                return {**state, "next_agent": next_agent, "current_task": current_task}
            except Exception as e:
                logger.error(f"Supervisor routing failed: {e}")
                span.set_attribute("error", True)
                span.set_attribute("error.message", str(e))
                await publish_a2a_message(
                    state['session_id'], "supervisor", "user",
                    "error", {"message": f"Supervisor routing failed: {e}"}
                )
                return {**state, "next_agent": "FINISH", "current_task": f"Error during routing: {e}"}

    async def run_supervisor(self, state: AgentState) -> AgentState:
        logger.info("Supervisor: Running supervisor logic.")
        return await self.route_agent(state)

# --- Specialist Agent Implementations (updated for A2A and tracing) ---
class WebSearchAgent:
    def __init__(self):
        self.tool = TavilySearchResults(max_results=5)

    async def run_web_search(self, state: AgentState) -> AgentState:
        logger.info(f"Web Search Agent: Searching for '{state['current_task']}'")
        with trace.get_tracer(__name__).start_as_current_span("web_search_agent.run_web_search") as span:
            span.set_attribute("session_id", state['session_id'])
            span.set_attribute("search_query", state['current_task'])
            try:
                await publish_a2a_message(
                    state['session_id'], "web_search_agent", "supervisor",
                    "status_update", {"message": f"Starting web search for: {state['current_task']}"}
                )
                search_results = self.tool.invoke({"query": state['current_task']})
                result_str = json.dumps(search_results, indent=2)
                add_to_episodic_memory(state['current_task'], result_str)
                await publish_a2a_message(
                    state['session_id'], "web_search_agent", "supervisor",
                    "result", {"task": state['current_task'], "output": result_str}
                )
                span.set_attribute("results_summary", result_str[:200])
                return {**state, "research_results": state["research_results"] + [f"Web Search Results for '{state['current_task']}':
{result_str}"], "next_agent": "supervisor"}
            except Exception as e:
                logger.error(f"Web Search Agent failed: {e}")
                span.set_attribute("error", True)
                span.set_attribute("error.message", str(e))
                await publish_a2a_message(
                    state['session_id'], "web_search_agent", "supervisor",
                    "error", {"task": state['current_task'], "message": f"Web Search failed: {e}"}
                )
                return {**state, "research_results": state["research_results"] + [f"Web Search failed for '{state['current_task']}': {e}"], "next_agent": "supervisor"}

class AcademicPaperAgent:
    async def run_academic_search(self, state: AgentState) -> AgentState:
        logger.info(f"Academic Paper Agent: Searching ArXiv for '{state['current_task']}'")
        with trace.get_tracer(__name__).start_as_current_span("academic_paper_agent.run_academic_search") as span:
            span.set_attribute("session_id", state['session_id'])
            span.set_attribute("search_query", state['current_task'])
            try:
                await publish_a2a_message(
                    state['session_id'], "academic_paper_agent", "supervisor",
                    "status_update", {"message": f"Starting academic search for: {state['current_task']}"}
                )
                loader = ArxivLoader(query=state['current_task'], load_max_docs=2, load_all_available_meta=True)
                docs = loader.load()
                if not docs:
                    await publish_a2a_message(
                        state['session_id'], "academic_paper_agent", "supervisor",
                        "result", {"task": state['current_task'], "output": f"No academic papers found for '{state['current_task']}'."}
                    )
                    span.set_attribute("found_papers", 0)
                    return {**state, "research_results": state["research_results"] + [f"No academic papers found for '{state['current_task']}'."], "next_agent": "supervisor"}

                result_summary = []
                for doc in docs:
                    summary = f"Title: {doc.metadata.get('Title', 'N/A')}
Authors: {doc.metadata.get('Authors', 'N/A')}
Published: {doc.metadata.get('Published', 'N/A')}
Summary: {doc.page_content[:500]}...
Link: {doc.metadata.get('entry_id', 'N/A')}"
                    result_summary.append(summary)
                
                full_result_str = "
---
".join(result_summary)
                add_to_episodic_memory(state['current_task'], full_result_str)
                await publish_a2a_message(
                    state['session_id'], "academic_paper_agent", "supervisor",
                    "result", {"task": state['current_task'], "output": full_result_str}
                )
                span.set_attribute("found_papers", len(docs))
                span.set_attribute("results_summary", full_result_str[:200])
                return {**state, "research_results": state["research_results"] + [f"Academic Research Results for '{state['current_task']}':
{full_result_str}"], "next_agent": "supervisor"}
            except Exception as e:
                logger.error(f"Academic Paper Agent failed: {e}")
                span.set_attribute("error", True)
                span.set_attribute("error.message", str(e))
                await publish_a2a_message(
                    state['session_id'], "academic_paper_agent", "supervisor",
                    "error", {"task": state['current_task'], "message": f"Academic Search failed: {e}"}
                )
                return {**state, "research_results": state["research_results"] + [f"Academic Search failed for '{state['current_task']}': {e}"], "next_agent": "supervisor"}

class CodeExecutionAgent:
    @tool("python_interpreter", return_direct=True)
    def python_interpreter(self, code: str) -> str:
        "Executes Python code in a sandboxed environment."
        with trace.get_tracer(__name__).start_as_current_span("code_execution_agent.python_interpreter") as span:
            span.set_attribute("code_snippet", code)
            try:
                logger.warning("Executing code locally without a proper sandbox. Use with caution.")
                process = subprocess.run(
                    [sys.executable, "-c", code],
                    capture_output=True,
                    text=True,
                    check=True,
                    timeout=30
                )
                if process.stderr:
                    logger.warning(f"Code execution stderr: {process.stderr}")
                span.set_attribute("execution_output", process.stdout.strip()[:200])
                return process.stdout.strip()
            except subprocess.CalledProcessError as e:
                error_msg = f"Code execution failed with error: {e.stderr.strip()}"
                logger.error(error_msg)
                span.set_attribute("error", True)
                span.set_attribute("error.message", error_msg)
                raise ValueError(error_msg)
            except subprocess.TimeoutExpired:
                error_msg = "Code execution timed out."
                logger.error(error_msg)
                span.set_attribute("error", True)
                span.set_attribute("error.message", error_msg)
                raise ValueError(error_msg)
            except Exception as e:
                error_msg = f"An unexpected error occurred during code execution: {e}"
                logger.error(error_msg)
                span.set_attribute("error", True)
                span.set_attribute("error.message", error_msg)
                raise ValueError(error_msg)

    async def run_code_execution(self, state: AgentState) -> AgentState:
        logger.info(f"Code Execution Agent: Attempting to execute code for task '{state['current_task']}'")
        with trace.get_tracer(__name__).start_as_current_span("code_execution_agent.run_code_execution") as span:
            span.set_attribute("session_id", state['session_id'])
            span.set_attribute("task_description", state['current_task'])
            try:
                await publish_a2a_message(
                    state['session_id'], "code_execution_agent", "supervisor",
                    "status_update", {"message": f"Starting code execution for: {state['current_task']}"}
                )
                code_generation_prompt = ChatPromptTemplate.from_messages([
                    ("system", "You are a helpful assistant capable of writing and executing Python code to answer questions or perform data analysis. "
                               "Given a task, generate valid Python code that can be executed. Output ONLY the Python code, no explanations or markdown fences."
                               "Current Task: {current_task}
"
                               "Existing Research Results: {research_results}
"
                               "Relevant Episodic Memory: {episodic_memory}
"
                               "Example: print(2+2)"),
                    ("human", "Generate Python code for the following task: {task}")
                ])
                llm_code_chain = code_generation_prompt | get_llm()

                relevant_memory = retrieve_from_episodic_memory(state['current_task'] + " ".join(state['research_results']))
                
                code_to_execute = llm_code_chain.invoke({
                    "task": state['current_task'],
                    "research_results": "
".join(state['research_results']) if state['research_results'] else "None yet.",
                    "episodic_memory": "
".join(relevant_memory) if relevant_memory else "None found."
                }).content
                logger.info(f"Generated code: {code_to_execute[:100]}...")
                span.set_attribute("generated_code_snippet", code_to_execute[:200])

                execution_output = self.python_interpreter.invoke(code_to_execute)
                add_to_episodic_memory(state['current_task'], f"Code Output: {execution_output}")
                await publish_a2a_message(
                    state['session_id'], "code_execution_agent", "supervisor",
                    "result", {"task": state['current_task'], "output": execution_output}
                )
                return {**state, "research_results": state["research_results"] + [f"Code Execution Results for '{state['current_task']}':
{execution_output}"], "next_agent": "supervisor"}
            except Exception as e:
                logger.error(f"Code Execution Agent failed: {e}")
                span.set_attribute("error", True)
                span.set_attribute("error.message", str(e))
                await publish_a2a_message(
                    state['session_id'], "code_execution_agent", "supervisor",
                    "error", {"task": state['current_task'], "message": f"Code Execution failed: {e}"}
                )
                return {**state, "research_results": state["research_results"] + [f"Code Execution failed for '{state['current_task']}': {e}"], "next_agent": "supervisor"}

# --- Build the LangGraph (from Phase 4) ---
def create_research_graph(llm: ChatOpenAI, memory: SqliteSaver):
    workflow = StateGraph(AgentState)

    supervisor = SupervisorAgent(llm)
    web_search = WebSearchAgent()
    academic_paper = AcademicPaperAgent()
    code_execution = CodeExecutionAgent()

    workflow.add_node("supervisor", supervisor.run_supervisor)
    workflow.add_node("web_search_agent", web_search.run_web_search)
    workflow.add_node("academic_paper_agent", academic_paper.run_academic_search)
    workflow.add_node("code_execution_agent", code_execution.run_code_execution)

    workflow.set_entry_point("supervisor")

    workflow.add_conditional_edges(
        "supervisor",
        lambda state: state['next_agent'],
        {
            "web_search_agent": "web_search_agent",
            "academic_paper_agent": "academic_paper_agent",
            "code_execution_agent": "code_execution_agent",
            "FINISH": END
        },
    )

    workflow.add_edge("web_search_agent", "supervisor")
    workflow.add_edge("academic_paper_agent", "supervisor")
    workflow.add_edge("code_execution_agent", "supervisor")

    app = workflow.compile(checkpointer=memory)
    return app

# --- FastAPI Application for Streaming (from Phase 4, updated for OpenTelemetry) ---
app_fastapi = FastAPI(title="Multi-Agent Research Platform")

class ResearchRequest(BaseModel):
    query: str
    session_id: str = str(uuid.uuid4())

async def stream_a2a_messages(session_id: str) -> AsyncIterator[str]:
    r = await get_redis_client()
    pubsub = r.pubsub()
    await pubsub.subscribe(f"a2a_channel:{session_id}")
    logger.info(f"Subscribed to A2A channel: a2a_channel:{session_id}")

    try:
        while True:
            message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0)
            if message:
                data = json.loads(message['data'].decode('utf-8'))
                yield f"data: {json.dumps(data)}

"
            await asyncio.sleep(0.01)
    except asyncio.CancelledError:
        logger.info(f"Streaming for session {session_id} cancelled.")
    except Exception as e:
        logger.error(f"Error in A2A message streaming for session {session_id}: {e}")
        yield f"data: {json.dumps({'event': 'error', 'message': str(e)})}

"
    finally:
        await pubsub.unsubscribe(f"a2a_channel:{session_id}")
        logger.info(f"Unsubscribed from A2A channel: a2a_channel:{session_id}")

@app_fastapi.post("/research/stream")
async def start_research_stream(request: ResearchRequest):
    session_id = request.session_id
    query = request.query
    logger.info(f"Received research request for session {session_id} with query: {query}")

    llm = get_llm()
    memory = SqliteSaver(conn=sqlite3.connect(f"./langgraph_checkpoints_{session_id}.sqlite", check_same_thread=False))
    research_graph = create_research_graph(llm, memory)

    async def research_task():
        try:
            initial_state = {"input": query, "chat_history": [HumanMessage(content=query)], "research_results": [], "next_agent": "", "current_task": "", "session_id": session_id}
            final_state = None
            async for s in research_graph.astream(initial_state, config={"configurable": {"thread_id": session_id}}):
                if "supervisor" in s:
                    current_state = s["supervisor"]
                    if current_state.get("next_agent") == "FINISH":
                        final_state = current_state
                        break
            
            if final_state:
                await publish_a2a_message(
                    session_id, "supervisor", "user",
                    "final_result", {"input": query, "final_results": final_state['research_results']}
                )
            else:
                await publish_a2a_message(
                    session_id, "supervisor", "user",
                    "error", {"message": "Research task completed without a final state.", "input": query}
                )
        except Exception as e:
            logger.error(f"Error during research task for session {session_id}: {e}")
            await publish_a2a_message(
                session_id, "supervisor", "user",
                "error", {"message": f"Research task failed: {e}", "input": query}
            )

    asyncio.create_task(research_task())

    return StreamingResponse(stream_a2a_messages(session_id), media_type="text/event-stream")

# --- Main application entry point for testing (FastAPI) ---
if __name__ == "__main__":
    print("
--- Phase 5: Observability with OpenTelemetry and LangSmith ---")
    setup_opentelemetry() # Initialize OpenTelemetry
    FastAPIInstrumentor.instrument_app(app_fastapi) # Instrument FastAPI

    print("Starting FastAPI server... access at http://127.0.0.1:8000/docs")
    print("Ensure Redis is running (e.g., `docker run -p 6379:6379 --name my-redis -d redis`).")
    print("To view traces, ensure an OpenTelemetry collector is running (e.g., Jaeger) or provide a LangSmith API Key.")
    uvicorn.run(app_fastapi, host="0.0.0.0", port=8000)