BAML x cognee: Type-Safe LLMs in Production

We’ve added BAML (from BoundaryML) as a first-class structured-output framework in cognee! BAML gives you a type-safe, schema-driven way to define LLM functions, validate responses, and version prompts—without fragile JSON parsing or excess boilerplate.

BAML lives alongside our existing Instructor support, and you can switch between these two frameworks with a single config flag, with no pipeline rewrites.

Teams using cognee to extract knowledge graphs, classify content, or summarize code now get the more reliable, validated outputs and easier prompt lifecycle management that this integration brings.

What Is BAML?

BAML is a domain-specific language (DSL) for building applications with LLMs. It ships a typed client and a DSL that enforce contracts at runtime. Instead of relying on free-form prompts and post-hoc JSON parsing, BAML lets you:

• Declare schemas and function contracts for LLM calls.
• Enforce types at runtime (and benefit from IDE autocompletion).
• Package and version prompts like code so teams can review changes via PRs.
• Swap providers/models via a client registry without touching call sites.

As a result, you get predictable outputs, simpler error handling, and a smoother path to production for LLM features.

Why Bring BAML Into cognee?

Many of our users have been asking for stronger guarantees around structured outputs and less prompt drift as projects grow.

While Instructor + LiteLLM remain supported and work well for many workloads, they are starting to feel outdated. BAML adds a complementary, opinionated model that reduces the boilerplate needed to manage prompts and centralizes schema + prompt management.

With BAML in cognee, you get:

• Reliability: BAML uses a json-repair algorithm to parse and validate outputs.
• Maintainability: Prompts and extraction logic live in a declarative DSL, easy to diff and review.
• Speed: Fewer glue layers (retries, validation, streaming) to hand-roll.
• Flexibility: Choose BAML or Instructor per environment via config, and keep shipping.

This integration is all about giving you the choice of output style that best suits your team and workload, while not burdening you with having to change how you design pipelines in cognee.

Use BAML when you need stricter guarantees (compliance, audits, multi-team); use Instructor for rapid prototyping where you prefer looser constraints.

What Ships in This Release

cognee's BAML integration is elegantly simple: it dynamically transforms your Pydantic models into BAML types and provides a single unified function for structured LLM outputs.

Here's how it works under the hood:

1. Dynamic Type Mapping: Your Pydantic models are automatically converted to BAML class definitions at runtime
2. Single Function Interface: One BAML function (AcreateStructuredOutput) handles all structured output requests
3. Automatic Validation: Responses are validated against your Pydantic schema before being returned

The BAML Function

function AcreateStructuredOutput(
    text_input: string,
    system_prompt: string,
) -> ResponseModel {
    client OpenAI

    prompt #"
        {{ system_prompt }}
        {{ ctx.output_format }}
        {{ _.role('user') }}
        {{ text_input }}
    "#
}

The ResponseModel is dynamically generated from your Pydantic model, ensuring type safety without manual schema duplication.

Dynamic Type Generation

The magic happens in the type mapping system. When you pass a Pydantic model to cognee, it:

• Maps primitive types: str, int, float, bool, datetime
• Handles collections: List[T], Dict[str, T]
• Supports optionals: Optional[T] and Union types
• Processes enums: Python Enum classes become BAML enums
• Recursively handles nested models: Complex Pydantic models with nested structures

# Your Pydantic model
class Person(BaseModel):
    name: str
    age: Optional[int]
    skills: List[str]

# Automatically becomes this BAML class:
class Person {
    name string
    age int?
    skills string[]
}

Implementation Notes (Clean, Swappable)

Framework Selection (No Recursion, Clear Fallback)

We expose a single entry point and route to the chosen framework at runtime:

from cognee.infrastructure.llm import LLMGateway
from pydantic import BaseModel
from typing import List

class KnowledgeGraph(BaseModel):
    entities: List[str]
    relationships: List[str]

# This automatically uses BAML if configured, Instructor otherwise
result = await LLMGateway.acreate_structured_output(
    text_input="Extract entities from this text...",
    system_prompt="You are an expert at knowledge extraction.",
    response_model=KnowledgeGraph
)

Why this pattern?

• No accidental recursion (the router never calls itself).
• Testable boundaries (you can unit-test each backend independently).
• Easy A/B (flip a flag per env to compare outputs/latency/cost).

Generated Client Behavior (What You Can Rely On)

BAML’s generated Python clients provide:

• Sync / async variants for pipelines and services.
• Automatic retries & backoff for transient failures.
• Streaming support for large responses.
• Runtime type validation for every response.

Tip: expose retry/backoff knobs and max token settings via config so ops can tune them without code changes.

How This Benefits Developers

• Fewer parsing bugs: You work with typed objects, not ad-hoc JSON.
• Faster iteration: Swap prompts/models via config; treat prompts & schemas like code (review, version, roll back).
• Lower boilerplate: Retries, validation, streaming are built-in.
• Safer refactors: IDE-friendly types surface breaking changes in tests—not in prod.
• Zero schema duplication: Your Pydantic models automatically become BAML types.

Quickstart Guide

0) Prereqs

• cognee updated to the version that includes BAML support.
• BAML installed & initialized (see BoundaryML docs).
• Provider API key set (OpenAI/Bedrock/Anthropic/Azure, etc.).
• Network egress to your LLM provider.

1) Pick Your Framework & Model

# .env or environment variables
STRUCTURED_OUTPUT_FRAMEWORK=BAML
BAML_LLM_PROVIDER=openai        # or bedrock, anthropic, azure, etc.
BAML_LLM_MODEL=gpt-4o           # match your provider
BAML_LLM_API_KEY=your_api_key

EMBEDDING_API_KEY=sk-...        # required for cognee

Cognee uses OpenAI as the default embedding model provider. Refer to our .env.template for other options.

2) Use Clearly Defined Pydantic Models

from pydantic import BaseModel, Field
from cognee.infrastructure.llm import LLMGateway
from typing import List, Optional
from your_prompts import your_system_prompt

class Entity(BaseModel):
    name: str
    type: str = Field(description="The category or type of this entity")
    properties: Optional[dict] = Field(default=None, description="Additional properties")

class KnowledgeGraph(BaseModel):
    entities: List[Entity] = Field(description="List of extracted entities")
    summary: str = Field(description="Brief summary of the content")

# Extract structured data - BAML handles the rest
graph = await LLMGateway.acreate_structured_output(
    text_input="Alice purchased 'The Graph Handbook' from CityBooks on March 3rd, 2025 in Seattle.She later reviewed it and recommended it to Bob.",
    system_prompt=your_system_prompt,
    response_model=KnowledgeGraph
)

print(f"Found {len(graph.entities)} entities")
for entity in graph.entities:
    print(f"- {entity.name} ({entity.type})")

Note that types need to be defined clearly. Pydantic models that don’t have specific types won’t work as there is no mapping in BAML to handle it.

End-to-End Example: From Text to Validated Knowledge Graph

0) Environment & Config

import os

# Framework toggle
os.environ["STRUCTURED_OUTPUT_FRAMEWORK"] = "BAML"

# BAML client configuration (example: OpenAI)
os.environ["BAML_LLM_PROVIDER"] = "openai"
os.environ["BAML_LLM_MODEL"]    = "gpt-4o"          # use your approved model
os.environ["BAML_LLM_API_KEY"]  = "sk-..."          # or rely on your secrets manager
os.environ["EMBEDDING_API_KEY"] = "sk-..."          # required for cognee

# Optional: temperature etc. if you’ve surfaced them in config
os.environ["BAML_LLM_TEMPERATURE"] = "0.1"

1) Imports & Model Types

from pydantic import BaseModel, Field
from typing import List, Optional

# Example KG schema used by your BAML function output
class Node(BaseModel):
    id: str
    type: str
    properties: dict = Field(default_factory=dict)

class Edge(BaseModel):
    source: str
    target: str
    type: str
    properties: dict = Field(default_factory=dict)

class KnowledgeGraph(BaseModel):
    nodes: List[Node] = Field(default_factory=list)
    edges: List[Edge] = Field(default_factory=list)

2) Wire Up cognee (BAML routing happens under the hood)

from cognee.infrastructure.llm import LLMGateway

sample_text = """
Alice purchased 'The Graph Handbook' from CityBooks on March 3rd, 2025 in Seattle.
She later reviewed it and recommended it to Bob.
"""

kg = await acreate_structured_output(
    text_input=sample_text,
    system_prompt=system_prompt,
    response_model=KnowledgeGraph,
)

You should see a validated KnowledgeGraph with nodes like Person:Alice, Book:The Graph Handbook, Store:CityBooks, City:Seattle, Person:Bob, and edges such as PURCHASED, REVIEWED, RECOMMENDED_TO, LOCATED_IN (with properties like date).

3) Inspect & Sanity-Check

print(f"Nodes: {len(kg.nodes)} | Edges: {len(kg.edges)}")
for n in kg.nodes:
    print(f"- {n.id} [{n.type}]")
for e in kg.edges:
    print(f"- {e.source} -[{e.type}]-> {e.target}  props={e.properties}")

Migrating From Instructor (Minimal Changes):

• Keep your Pydantic models. They become the response schema for BAML as-is.
• Flip the framework flag (STRUCTURED_OUTPUT_FRAMEWORK=BAML).
• Optional: Move prompt logic into BAML DSL so you can version & review it independently.

Structured Outputs for a Scalable Future

We built this integration to remove daily dev paper-cuts: brittle prompts, silent schema drift, surprise JSON failures. With BAML in cognee, you get production hygiene: fewer parsing errors, clearer versioning, safer changes, faster iteration.

BAML’s DSL lets you treat prompts like code—typed, diffable, testable—while cognee turns those outputs into knowledge graphs, classification, and code summaries you can trust. Keep using Instructor where it fits; switch to BAML where type-safety and governance matter most.

Ready to try it for yourself?

• See our implementation in the cognee repo
• Learn more about BAML in BoundaryML's documentation

Thanks & Credits

Huge thanks to the BAML team for their round-the-clock support and quick feedback during integration. We deeply appreciate how engaged and hard working they are and we’re excited about what this update unlocks for cognee users worldwide.