Query Agent Architecture¶

This document describes the three query agent implementations in the Graffold and their relationships.

Overview¶

The system provides three query agents with different execution strategies:

Agent	Type	Complexity	Use Case
DynamicQueryAgent	Base pipeline	Medium	Comprehensive multi-strategy search
Neo4jQueryAgent	Enhanced pipeline	High	Neo4j-optimized with RAG and APOC
TwoPhaseAgent	Deterministic + LLM	Medium	Guaranteed tool execution + LLM synthesis
LangGraphAgent	Agentic ReAct	Medium	(Shelved) LLM-driven tool selection

Agent Hierarchy¶

┌─────────────────────────────────────────────────────────────────────┐
│                        DynamicQueryAgent                            │
│  (Base class - src/agents/dynamic_query_agent.py)                   │
│  - Direct database connection via DatabaseInterface                 │
│  - QueryProcessor, EntityDiscovery, SearchStrategies                │
│  - CrossEncoder re-ranking                                          │
│  - REFRAG context compression                                       │
│  - 6 query strategies executed in deterministic pipeline            │
└───────────────────────────┬─────────────────────────────────────────┘
                            │
                            │ extends (inheritance)
                            ▼
┌─────────────────────────────────────────────────────────────────────┐
│                        Neo4jQueryAgent                              │
│  (Subclass - src/agents/neo4j_query_agent.py)                       │
│  - Inherits all DynamicQueryAgent capabilities                      │
│  - Adds graph-optimized components (work on both Neo4j & Memgraph): │
│    • VectorRetriever (chunk_embeddings index)                       │
│    • HybridRetriever (node_embeddings + fulltext)                   │
│    • GraphRAG (LLM + chunk retriever)                               │
│    • APOCKNNAlgorithms (KNN graph expansion, Neo4j only)            │
│  - Uses neo4j_graphrag library for optimized retrieval              │
└─────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────┐
│                        LangGraphAgent  [SHELVED]                    │
│  (Standalone - src/agents/graph_agent.py)                           │
│  - NO inheritance - completely independent architecture             │
│  - Creates database via DatabaseFactory (Neo4j or Memgraph)         │
│  - ReAct loop via LangGraph's create_react_agent                    │
│  - 8 tools with direct Cypher queries                               │
│  - LLM decides which tools to call and when                         │
│  - MemorySaver for conversation persistence                         │
│  - PROBLEM: LLM often skips tool calls, hallucinating answers       │
└─────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────┐
│                        TwoPhaseAgent  [ACTIVE - replaces Agentic]   │
│  (Standalone - src/agents/two_phase_agent.py)                       │
│  - NO inheritance - completely independent architecture             │
│  - Creates database via DatabaseFactory                             │
│  - Phase 1: Deterministic execution of ALL 8+ tools (no LLM)       │
│  - Phase 2: Single LLM call to synthesize results into report       │
│  - Reuses GraphTools from src/tools/graph_tools.py                  │
│  - Guaranteed tool execution — 0 hallucinated data                  │
└─────────────────────────────────────────────────────────────────────┘

1. DynamicQueryAgent¶

File: src/agents/dynamic_query_agent.py
Lines: 1,209
Type: Base class with procedural pipeline

Architecture¶

class DynamicQueryAgent:
    def __init__(self, service: str, database: str, connection_details: dict | None):
        self.db = DatabaseFactory.create_database(database=database)  # Neo4j or Memgraph
        self.llm = get_llm(service)
        self.embedder = get_embedder("huggingface", "all-mpnet-base-v2")

        # Pipeline components
        self.query_processor = QueryProcessor(service)
        self.entity_discovery = EntityDiscovery(self.embedder)
        self.search_strategies = SearchStrategies(self.db, self.embedder)
        self.result_processor = ResultProcessor(self.cross_encoder)
        self.refrag_processor = REFRAGProcessor()

Execution Flow¶

Preprocessing: Query rewriting and expansion with LLM
Entity Discovery: Extract and validate relevant entities
Multi-Strategy Search: Execute 6 complementary query strategies:
Direct entity queries
Schema-guided expansion
Relationship traversal
Topic-based queries
LLM-generated Cypher queries
Literature search
Dynamic Expansion: Heuristic rules for additional searches
Scoring & Ranking: CrossEncoder re-ranking (top 500)
Answer Generation: LLM synthesizes comprehensive answer

Key Features¶

Comprehensive coverage: Multiple search strategies ensure no relevant data is missed
Intelligent expansion: Dynamic rules trigger additional searches based on results
Context awareness: Query rewriting uses conversation history
Scoring pipeline: Combines relevance scores from multiple sources

When to Use¶

Complex queries requiring multi-strategy coverage
Need for deterministic, reproducible results
Situations where you want all relevant information gathered systematically

2. Neo4jQueryAgent¶

File: src/agents/neo4j_query_agent.py
Lines: 411
Type: Subclass extending DynamicQueryAgent

Architecture¶

class Neo4jQueryAgent(DynamicQueryAgent):
    def __init__(self, service: str, database: str, connection_details: dict | None):
        # Inherits parent init but doesn't call super().__init__()
        # Reimplements with graph-optimized stack

        self.db = DatabaseFactory.create_database(database=database)  # Neo4j or Memgraph
        self.embedder = get_embedder("huggingface", "all-distilroberta-v1")

        # Neo4j-specific retrievers
        self.chunk_retriever = VectorRetriever(
            driver=self.db.driver,
            index_name='chunk_embeddings',
            embedder=self.embedder
        )

        self.hybrid_retriever = HybridRetriever(
            driver=self.db.driver,
            vector_index_name='node_embeddings',
            fulltext_index_name='node_fulltext',
            embedder=self.embedder
        )

        # GraphRAG for LLM-powered retrieval
        self.rag = GraphRAG(
            retriever=self.chunk_retriever,
            llm=self.llm
        )

        # APOC for KNN expansion
        self.knn_expander = APOCKNNAlgorithms(
            driver=self.db.driver,
            database=self.database
        )

Additional Capabilities¶

Beyond DynamicQueryAgent:

VectorRetriever: Optimized vector similarity search via neo4j_graphrag
HybridRetriever: Combined vector + fulltext search
GraphRAG: LLM-powered retrieval augmented generation
APOC KNN: Graph algorithm-based neighbor expansion (Neo4j only)
Different embedder: Uses all-distilroberta-v1 (matching ingestion)

Note: Despite the class name, Neo4jQueryAgent works with both Neo4j and Memgraph backends. APOC-based features (KNN expansion) are only available with Neo4j.

Execution Flow¶

Inherits DynamicQueryAgent's pipeline but with enhanced retrievers:

Uses neo4j_graphrag library for optimized index queries
APOC algorithms for efficient graph traversal (Neo4j only)
GraphRAG for context-aware chunk retrieval

When to Use¶

Production deployment with Neo4j or Memgraph database
Need for optimized graph-specific operations
Queries benefit from APOC graph algorithms (Neo4j)
Want GraphRAG-style retrieval

3. LangGraphAgent¶

File: src/agents/graph_agent.py
Lines: 706
Type: Standalone ReAct agent (no inheritance)

Architecture¶

class LangGraphAgent:
    def __init__(self, service: str, database: str, connection_details: dict | None, ...):
        self.llm = get_chat_model(service)

        # Create database via factory
        self.db = DatabaseFactory.create_database(database=database)

        # Create embedder for vector tools
        self.embedder = get_embedder("huggingface", "all-mpnet-base-v2")

        # Create 8 tools
        self.tools = create_graph_tools(
            db=self.db,
            retriever=graph_retriever,
            embedder=self.embedder
        )

        # Build ReAct agent with LangGraph
        self.graph = create_react_agent(
            model=self.llm,
            tools=self.tools,
            prompt=SYSTEM_PROMPT,
            checkpointer=MemorySaver()
        )

Tools (8 total)¶

Graph Traversal Tools¶

count_nodes - Count entities by label
lookup_entity - Find entity by name (returns JSON)
get_entity_neighbors - Explore connections from an entity
get_disease_proteins - Disease → proteins relationships
get_protein_diseases - Protein → diseases relationships

Semantic Search Tools¶

search_knowledge_graph - Hybrid text + keyword search
vector_search_entities - Semantic similarity search for entities
vector_search_literature - Semantic search over PubMed chunks

Execution Flow¶

ReAct Loop (LLM-driven):

Thought: LLM analyzes the query and decides strategy
Action: LLM selects appropriate tool(s) to call
Observation: Tool results returned as JSON
Repeat: LLM decides if more tools needed or answer ready
Answer: Final synthesis with citations

Key Features¶

Autonomous reasoning: LLM decides which tools to use
Flexible strategy: Not locked into fixed query pipeline
Conversation memory: MemorySaver persists thread context
Structured output: Pydantic models for entities, relationships, table_data
Direct Cypher: Tools use db.index.vector.queryNodes() for reliability

Tool Selection Strategy¶

LLM chooses tools based on query type:

Query Type	Recommended Tool
"How many proteins..."	`count_nodes`
"Tell me about BRCA1"	`lookup_entity` → `get_entity_neighbors`
"Proteins linked to Alzheimer's"	`get_disease_proteins`
"What diseases involve TP53"	`get_protein_diseases`
"Biomarkers similar to..."	`vector_search_entities`
"Research about X"	`vector_search_literature`

When to Use¶

Queries that benefit from flexible reasoning
Need for conversation context and memory
Want transparent tool selection (can inspect agent decisions)
Prototyping or research scenarios
Complex multi-step reasoning tasks

Database Access Patterns¶

DynamicQueryAgent & Neo4jQueryAgent¶

# Database access via DatabaseInterface (works with Neo4j or Memgraph)
self.db = DatabaseFactory.create_database(database=database)

# Execute Cypher directly (Cypher is supported by both Neo4j and Memgraph)
with self.db.driver.session(database=self.database) as session:
    result = session.run(cypher_query, parameters)

# Use neo4j_graphrag retrievers (Neo4jQueryAgent only)
results = self.hybrid_retriever.search(query_text=question, top_k=50)

LangGraphAgent¶

# Database via factory (auto-selects Neo4j, Memgraph, or FalkorDB based on DATABASE_TYPE)
self.db = DatabaseFactory.create_database(database=database)

# Tools execute Cypher
results = self.db._execute_cypher(query, parameters)

# Vector search via direct Cypher
query_embedding = self.embedder.embed_query(query)
results = self.db._execute_cypher(
    "CALL db.index.vector.queryNodes('node_embeddings', $limit, $embedding)",
    {"limit": limit, "embedding": query_embedding}
)

Embedder Models¶

Agent	Embedder	Dimension	Use Case
DynamicQueryAgent	all-mpnet-base-v2	768	General purpose
Neo4jQueryAgent	all-distilroberta-v1	768	Matches ingestion model
LangGraphAgent	all-mpnet-base-v2	768	General purpose

Note: Neo4jQueryAgent uses all-distilroberta-v1 to match the model used during data ingestion for optimal retrieval performance.

UI Integration¶

In query_service.py, agents are selected via agent_type:

def _build_neo4j_agent(self, agent_type: str, service: str, database: str, ...):
    normalized = agent_type.strip().lower()

    if normalized.startswith("dynamic"):
        return DynamicQueryAgent(service=service, database=database, ...)

    if normalized.startswith("hybrid"):
        return Neo4jQueryAgent(service=service, database=database, ...)

    if normalized.startswith("agentic"):
        return TwoPhaseAgent(service=service, database=database, ...)

    # Default
    return Neo4jQueryAgent(service=service, database=database, ...)

Frontend Options¶

Recommend these 3 options in UI:

Standard → Neo4jQueryAgent
"Production-optimized with GraphRAG and APOC (works with Neo4j and Memgraph)"
Dynamic → DynamicQueryAgent
"Comprehensive multi-strategy pipeline"
Agentic → TwoPhaseAgent
"Deterministic tool execution + LLM synthesis (guaranteed data)"

Performance Characteristics¶

Agent	Speed	Coverage	Transparency	Determinism
DynamicQueryAgent	Medium	High	Medium	High
Neo4jQueryAgent	Fast	High	Medium	High
TwoPhaseAgent	Medium	High	High	High
LangGraphAgent	Variable*	Medium	High**	Low***

* Speed depends on LLM reasoning steps
** Can inspect tool calls in conversation history
*** LLM may choose different tools for same query

Common Interface¶

All agents implement the same sync_query() interface:

response = agent.sync_query(
    question="What proteins are associated with Alzheimer's disease?",
    include_graph_analytics=True
)

# Returns consistent format:
{
    "answer": str,
    "results": list[dict],
    "detailed_results": list[dict],
    "table_data": dict,
    "sources": list[str],
    "graph_analytics": dict,
    "confidence": float,
    "relationships": list[dict],
}

Migration Path¶

From Neo4jQueryAgent to LangGraphAgent¶

# Before
from src.agents.neo4j_query_agent import Neo4jQueryAgent
agent = Neo4jQueryAgent(service="bedrock", database="psx")

# After
from src.agents.graph_agent import LangGraphAgent
agent = LangGraphAgent(service="bedrock", database="psx")

# API stays the same
response = agent.sync_query("your question")

From DynamicQueryAgent to the graph databaseQueryAgent¶

# Before
from src.agents.dynamic_query_agent import DynamicQueryAgent
agent = DynamicQueryAgent(service="ollama", database="olink")

# After - inherits all Dynamic capabilities + Neo4j optimizations
from src.agents.neo4j_query_agent import Neo4jQueryAgent
agent = Neo4jQueryAgent(service="ollama", database="olink")

Testing¶

Unit Tests¶

DynamicQueryAgent: archive_tests/test_dynamic_query_agent.py
Neo4jQueryAgent: archive_tests/test_neo4j_query_agent.py
LangGraphAgent: Create tests/test_langgraph_agent.py

Comparison Evaluation¶

# src/evaluation/evaluate_agent_comparison.py
from src.agents.neo4j_query_agent import Neo4jQueryAgent
from src.agents.graph_agent import LangGraphAgent

baseline = Neo4jQueryAgent(service="bedrock", database="neo4j")
new_agent = LangGraphAgent(service="bedrock", database="neo4j")

# Compare performance on test queries

Future Enhancements¶

DynamicQueryAgent¶

[ ] Async pipeline execution for parallel strategies
[ ] Configurable strategy weights
[ ] Query plan caching

Neo4jQueryAgent¶

[ ] Optimize APOC algorithm selection
[ ] Adaptive retriever selection based on query
[ ] Batch query optimization

LangGraphAgent¶

[ ] Add more specialized tools (pathway search, GO term lookup)
[ ] Implement tool use history for learning
[ ] Add streaming support for real-time responses
[ ] Tool result caching within conversation

Best Practices¶

Development: Use LangGraphAgent for flexibility and transparency
Production: Use Neo4jQueryAgent for speed and reliability (works with both Neo4j and Memgraph)
Research: Use DynamicQueryAgent for comprehensive coverage
Testing: Compare all three agents on representative queries

Query Agent Architecture¶

Overview¶

Agent Hierarchy¶

1. DynamicQueryAgent¶

Architecture¶

Execution Flow¶

Key Features¶

When to Use¶

2. Neo4jQueryAgent¶

Architecture¶

Additional Capabilities¶

Execution Flow¶

When to Use¶

3. LangGraphAgent¶

Architecture¶

Tools (8 total)¶

Graph Traversal Tools¶

Semantic Search Tools¶

Execution Flow¶

Key Features¶

Tool Selection Strategy¶

When to Use¶

Database Access Patterns¶

DynamicQueryAgent & Neo4jQueryAgent¶

LangGraphAgent¶

Embedder Models¶

UI Integration¶

Frontend Options¶

Performance Characteristics¶

Common Interface¶

Migration Path¶

From Neo4jQueryAgent to LangGraphAgent¶

From DynamicQueryAgent to the graph databaseQueryAgent¶

Testing¶

Unit Tests¶

Comparison Evaluation¶

Future Enhancements¶

DynamicQueryAgent¶

Neo4jQueryAgent¶

LangGraphAgent¶

Best Practices¶

References¶