LangGraph + Neo4j Integration

from langgraph.graph import END, START, StateGraph
from typing_extensions import Annotated, TypedDict

class MessagesState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]
    llm_calls: int

# Build your custom workflow
agent_builder = StateGraph(MessagesState)

# Add nodes - each is a function that transforms state
agent_builder.add_node("llm_call", llm_call)
agent_builder.add_node("tool_node", tool_node)

# Define the flow with edges
agent_builder.add_edge(START, "llm_call")
agent_builder.add_conditional_edges("llm_call", should_continue, ["tool_node", END])
agent_builder.add_edge("tool_node", "llm_call")

agent = agent_builder.compile()

from langchain_mcp_adapters.client import MultiServerMCPClient

cypher_mcp_config = {
    "neo4j-database": {
        "transport": "stdio",
        "command": "uvx",
        "args": ["mcp-neo4j-cypher"],
        "env": {
            "NEO4J_URI": os.environ["NEO4J_URI"],
            "NEO4J_USERNAME": os.environ["NEO4J_USERNAME"],
            "NEO4J_PASSWORD": os.environ["NEO4J_PASSWORD"],
            "NEO4J_DATABASE": os.environ["NEO4J_DATABASE"]
        }
    }
}

client = MultiServerMCPClient(cypher_mcp_config)
mcp_tools = await client.get_tools()

from langchain.tools import tool
from langchain_neo4j import Neo4jGraph

neo4j_graph = Neo4jGraph()

@tool
async def get_investments(company: str) -> str:
    """Returns the investments by a company by name."""
    results = neo4j_graph.query("""
        MATCH (o:Organization)-[:HAS_INVESTOR]->(i)
        WHERE o.name = $company
        RETURN i.id as id, i.name as name, head(labels(i)) as type
    """, {"company": company})
    return json.dumps(results, indent=2)

from langgraph.checkpoint.neo4j import AsyncNeo4jSaver

async with await AsyncNeo4jSaver.from_conn_string(
    uri=NEO4J_URI,
    user=NEO4J_USERNAME,
    password=NEO4J_PASSWORD,
    database=NEO4J_DATABASE
) as checkpointer:
    await checkpointer.setup()

    agent = agent_builder.compile(checkpointer=checkpointer)

from neo4j_graphrag.retrievers import HybridCypherRetriever

# `node` and `score` are in scope: continue into the graph from each vector hit.
RETRIEVAL_QUERY = """
WITH node AS chunk, score
MATCH (article:Article)-[:HAS_CHUNK]->(chunk)
OPTIONAL MATCH (article)-[:MENTIONS]->(org:Organization)
RETURN chunk.text AS text, article.title AS title,
       collect(DISTINCT org.name)[..5] AS companies, score
"""

graph_retriever = HybridCypherRetriever(
    driver=driver,
    vector_index_name="news",
    fulltext_index_name="news_fulltext",
    retrieval_query=RETRIEVAL_QUERY,
    embedder=OpenAIEmbeddings(),
)

def graph_search(state: RetrievalState):
    results = graph_retriever.search(query_text=state["question"], top_k=8)
    return {"documents": [item.content for item in results.items]}

builder.add_conditional_edges(
    "route_question", pick_retriever,
    ["vector_search", "hybrid_search", "graph_search"],
)

# Every retriever feeds the same grading node, which cannot be skipped.
for node in ["vector_search", "hybrid_search", "graph_search"]:
    builder.add_edge(node, "grade_context")

builder.add_conditional_edges("grade_context", decide_next, ["graph_search", "generate"])

client = MultiServerMCPClient({
    "aura-agent": {
        "transport": "streamable_http",
        "url": os.environ["AURA_AGENT_MCP_URL"],
        "headers": {"Authorization": f"Bearer {token}"},
    }
})
aura_tools = await client.get_tools()

# The hosted agent is one tool among your local ones.
all_tools = aura_tools + [portfolio_weight]
tools_by_name = {t.name: t for t in all_tools}
model_with_tools = model.bind_tools(all_tools)

async def tool_node(state: MessagesState):
    """Run the requested tools. One of them is an entire agent hosted in Aura."""
    results = []
    for call in state["messages"][-1].tool_calls:
        observation = await tools_by_name[call["name"]].ainvoke(call["args"])
        results.append(ToolMessage(content=str(observation), tool_call_id=call["id"]))
    return {"messages": results}