Text-to-Cypher Retrieval

Text-to-Cypher Retrieval is a specialized retrieval strategy within the am_diag project that translates natural language clinical questions into executable Cypher queries. This component enables the system to perform complex, multi-hop reasoning over the Neo4j knowledge graph that might be difficult for standard vector or hybrid search methods to capture. It leverages Large Language Models (LLMs) via the BAML framework to generate queries constrained by the project's runtime-injectable medical schema.

Sources: README.md:179-181, am_diag/retrieval/retrievers/text2cypher.py:11-23

Architecture and Core Logic

The system utilizes a Text2CypherRetriever class to orchestrate the conversion of text to graph queries. This process involves schema serialization, LLM-based query generation, and a protective "EXPLAIN" gate to ensure query validity before execution against the graph store.

Retrieval Flow

The execution flow follows a structured loop to ensure reliability:

1. Schema Preparation: The GraphSchema is serialized into a textual description that explains the :RELATES_TO generic-edge model to the LLM.
2. BAML Generation: The GenerateCypherQuery function is called with the user's question, the schema description, and optional few-shot examples.
3. Validation (EXPLAIN Gate): The generated Cypher is prepended with EXPLAIN and sent to Neo4j. This validates the syntax and read-only status without executing the full query.
4. Bounded Retry: If validation fails, the retriever attempts the process again, up to a configured maximum number of attempts.
5. Execution: Validated queries are executed, and results are wrapped in SearchResult objects.

Sources: am_diag/retrieval/retrievers/text2cypher.py:25-50, am_diag/llm/baml_client/inlinedbaml.py:270-298

Visualizing the Flow

The following diagram illustrates the retry logic and validation gate within the retriever.

The diagram shows the iterative cycle between LLM generation and database-side validation. Sources: am_diag/retrieval/retrievers/text2cypher.py:53-90

LLM Integration (BAML)

The project uses BAML (Boundary Machine Learning) to define type-safe LLM functions. The GenerateCypherQuery function is the core interface for Text-to-Cypher.

Data Models for Retrieval

Sources: am_diag/llm/baml_client/inlinedbaml.py:263-268

Prompting and Constraints

The system enforces strict rules on the LLM to maintain security and accuracy:

• Read-only Restriction: Queries must only use MATCH, OPTIONAL MATCH, WITH, WHERE, and RETURN. Mutations (CREATE, MERGE, DELETE, SET) are forbidden.
• Identifier Usage: The system uses n.uuid for node identity and n.name for human-readable labels.
• Result Limits: A mandatory LIMIT clause (maximum 50) is required to prevent runaway queries.
• Schema Adherence: The LLM must use the specific node labels and relationship types defined in the injected schema.

Sources: am_diag/llm/baml_client/inlinedbaml.py:284-290

Schema Injection

The GraphSchema is a runtime-injectable value that drives the retriever. It translates the internal medical knowledge graph structure into a format the LLM can understand.

The diagram represents the generic-edge schema pattern described to the LLM. Sources: am_diag/common/data_models/schema.py:92-108, README.md:57-65

The as_text() method in GraphSchema produces a summary of the graph's nodes and relationships, specifically highlighting that all relationships use the :RELATES_TO label with a type property corresponding to specific medical relations like TREATED_BY or HAS_SYMPTOM.

Sources: am_diag/common/data_models/schema.py:92-108

Configuration and Components

The behavior of the Text-to-Cypher module is governed by the RetrievalConfig and the GraphStore.

Sources: am_diag/retrieval/retrievers/text2cypher.py:25-45, am_diag/retrieval/config.py (inferred via imports)

Retrieval Sequence

Implementation Details

The implementation emphasizes safety and lazily cached schema descriptions to optimize performance. The retriever identifies nodes using the uuid property and returns results with a descending score based on their row position to fit the standard SearchResult interface used across the engine.

Code Example: Retriever Logic

# am_diag/retrieval/retrievers/text2cypher.py:53-81
async def retrieve(self, query: str, filters: dict | None = None) -> list[SearchResult]:
    if not self._config.use_text_to_cypher:
        return []

    schema_desc = self._get_schema_description()
    cypher_result = None

    for attempt in range(self._config.text2cypher_max_attempts):
        result = await b.GenerateCypherQuery(query, schema_desc, "")
        cypher = result.cypher.strip()
        try:
            await self._graph.execute_query(f"EXPLAIN {cypher}")
            cypher_result = cypher
            break
        except Exception as exc:
            logger.warning(f"EXPLAIN rejected (attempt {attempt + 1}): {exc}")

Sources: am_diag/retrieval/retrievers/text2cypher.py:53-81

Conclusion

Text-to-Cypher Retrieval serves as a powerful bridge between natural language clinical inquiries and structured graph data. By combining BAML-driven LLM generation with a mandatory validation gate and bounded retries, it provides a robust mechanism for extracting precise medical insights from the Neo4j knowledge graph while maintaining system stability and read-only security.