Structured LLM Generation via BAML
Relevant source files
The following files were used as context for generating this wiki page:
Structured LLM generation in the am_diag project is facilitated by BAML (Boundary ML), a framework for building type-safe LLM functions. It ensures that the outputs from Large Language Models (LLMs) adhere to specific Pydantic schemas, which is critical for medical knowledge graph construction and diagnostic reasoning. BAML is used to constrain the LLM to a runtime-injectable medical schema, allowing the system to extract entities and relations, generate community reports, and translate natural language into Cypher queries.
The system utilizes a generated BAML client that provides both synchronous and asynchronous interfaces. This architecture allows the project to maintain strict type safety while interacting with various model providers such as OpenAI, Anthropic, and generic OpenAI-compatible endpoints.
Sources: README.md, am_diag/llm/baml_client/inlinedbaml.py:270-286, am_diag/graph_construction/extract/llm.py:84-88
System Architecture
The BAML integration consists of a source layer (.baml files), a generation step, and a typed Python client. The Python client acts as the primary interface for feature modules like the LLMExtractor and CommunitySummarizer.
The diagram above illustrates the flow from BAML source definition to the final consumption of typed data by the project's core modules.
Sources: CONTRIBUTING.md:73-77, am_diag/graph_construction/extract/llm.py:30-33, am_diag/graph_construction/community/summarize.py:29-35
Core BAML Functions
The project defines several specialized BAML functions for different stages of the GraphRAG pipeline. These functions are categorized by their role in knowledge extraction, graph management, and retrieval.
Knowledge Extraction
The ExtractKnowledgeGraph function is the centerpiece of LLM-based ingestion. It uses dynamic enums (EntityType and RelationType) populated at runtime via a TypeBuilder to ensure the LLM only extracts concepts defined in the system's medical schema.
| Function | Input | Output | Description |
|---|---|---|---|
ExtractKnowledgeGraph | text, source_type | ExtractedKnowledgeGraph | Extracts entities and relations constrained by dynamic schema. |
AreEntityMentionsSame | entityA, typeA, entityB, typeB | SamenessVerdict | Determines if two medical mentions refer to the same real-world entity. |
IdentifyClusterDuplicates | item, itemKind, candidateSet | ClusterDeduplicationVerdict | Identifies exact duplicates within a semantic cluster. |
Sources: am_diag/llm/baml_client/inlinedbaml.py:125-168, am_diag/llm/baml_client/inlinedbaml.py:186-210, am_diag/llm/baml_client/inlinedbaml.py:221-263
Graph Analysis and Retrieval
BAML is also used for synthesizing information from the graph and generating queries.
- Community Reporting:
GenerateCommunityReportrolls up findings from specific graph clusters into a structured report. - Cypher Generation:
GenerateCypherQuerytranslates natural language questions into Neo4j Cypher queries.
Sources: am_diag/llm/baml_client/inlinedbaml.py:76-92, am_diag/llm/baml_client/inlinedbaml.py:291-305
Data Models and Schemas
BAML generates Pydantic-based classes in am_diag/llm/baml_client/types.py that represent the structured output of the LLM.
Key Data Structures
This class diagram reflects the structured output format for knowledge extraction. Note that EntityType and RelationType are defined as dynamic enums that are empty in the static code but populated via the TypeBuilder.
Sources: am_diag/llm/baml_client/types.py:53-83, am_diag/llm/baml_client/inlinedbaml.py:127-128
Implementation Details
Asynchronous Client usage
Feature modules typically use the BamlAsyncClient (aliased as b in the client package) to perform non-blocking LLM calls. The client supports with_options for injecting the TypeBuilder (tb).
# am_diag/graph_construction/extract/llm.py:30-33
async def _call_llm(text: str, source: str, schema: GraphSchema) -> object:
from am_diag.llm.baml_client import b
tb = build_schema_typebuilder(schema)
return await b.ExtractKnowledgeGraph(text, source, baml_options={"tb": tb})
Sources: am_diag/graph_construction/extract/llm.py:30-33, am_diag/llm/baml_client/async_client.py:461
Provider Configuration
The system uses a round-robin and fallback strategy for LLM clients, ensuring high availability. It defines a generic client OpenAIGeneric which is often used as the primary provider for extraction and summarization tasks.
| Client Name | Provider | Retry Policy | Description |
|---|---|---|---|
CustomFast | round-robin | N/A | Alternates between GPT-5-Mini and Claude Haiku. |
OpenaiFallback | fallback | N/A | Tries GPT-5-Mini, then GPT-5. |
OpenAIGeneric | openai-generic | GenerousRetry | Primary client for extraction and retrieval. |
Sources: am_diag/llm/baml_client/inlinedbaml.py:48-74, am_diag/llm/baml_client/inlinedbaml.py:101-113
Conclusion
Structured generation via BAML is the foundational mechanism for maintaining medical domain integrity in the am_diag system. By combining dynamic type building with rigid Pydantic output schemas, the project successfully bridges the gap between unstructured medical text and a structured knowledge graph.