From 9ca050a70d36c36b5047810fb5344b241b50e09c Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Tue, 5 Aug 2025 21:23:40 +0200
Subject: [PATCH] feat: Add baml specific llm config options

---
 .env.template                       |  6 ++++++
 .github/workflows/basic_tests.yml   | 13 ++++++-------
 cognee/infrastructure/llm/config.py | 18 +++++++++++++-----
 3 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/.env.template b/.env.template
index a90171806..0beba7d1c 100644
--- a/.env.template
+++ b/.env.template
@@ -31,6 +31,12 @@ EMBEDDING_MAX_TOKENS=8191
 # If embedding key is not provided same key set for LLM_API_KEY will be used
 #EMBEDDING_API_KEY="your_api_key"
 
+# If using BAML structured output these env variables will be used
+BAML_LLM_PROVIDER=openai
+BAML_LLM_MODEL="gpt-4o-mini"
+BAML_LLM_ENDPOINT=""
+BAML_LLM_API_KEY="your_api_key"
+
 ################################################################################
 # 🗄️ Relational database settings
 ################################################################################
diff --git a/.github/workflows/basic_tests.yml b/.github/workflows/basic_tests.yml
index 33b442065..7aa87cc29 100644
--- a/.github/workflows/basic_tests.yml
+++ b/.github/workflows/basic_tests.yml
@@ -150,18 +150,17 @@ jobs:
     name: Run Simple Examples BAML
     runs-on: ubuntu-22.04
     env:
-      LLM_PROVIDER: openai-generic
-      LLM_MODEL: ${{ secrets.LLM_MODEL }}
-      LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
-      LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
-      LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
+      STRUCTURED_OUTPUT_FRAMEWORK: "BAML"
+      BAML_LLM_PROVIDER: openai-generic
+      BAML_LLM_MODEL: ${{ secrets.LLM_MODEL }}
+      BAML_LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
+      BAML_LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
 
-      EMBEDDING_PROVIDER: openai-generic
+      EMBEDDING_PROVIDER: openai
       EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
       EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
       EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
       EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
-      STRUCTURED_OUTPUT_FRAMEWORK: "BAML"
     steps:
       - name: Check out repository
         uses: actions/checkout@v4
diff --git a/cognee/infrastructure/llm/config.py b/cognee/infrastructure/llm/config.py
index 6c2f8c0c2..54153ed38 100644
--- a/cognee/infrastructure/llm/config.py
+++ b/cognee/infrastructure/llm/config.py
@@ -42,6 +42,13 @@ class LLMConfig(BaseSettings):
     llm_temperature: float = 0.0
     llm_streaming: bool = False
     llm_max_tokens: int = 16384
+
+    baml_llm_provider: str = "openai"
+    baml_llm_model: str = "gpt-4o-mini"
+    baml_llm_endpoint: str = ""
+    baml_llm_api_key: Optional[str] = None
+    baml_llm_temperature: float = 0.0
+
     transcription_model: str = "whisper-1"
     graph_prompt_path: str = "generate_graph_prompt.txt"
     llm_rate_limit_enabled: bool = False
@@ -62,12 +69,13 @@ class LLMConfig(BaseSettings):
     def model_post_init(self, __context) -> None:
         """Initialize the BAML registry after the model is created."""
         self.baml_registry.add_llm_client(
-            name=self.llm_provider,
-            provider=self.llm_provider,
+            name=self.baml_llm_provider,
+            provider=self.baml_llm_provider,
             options={
-                "model": self.llm_model,
-                "temperature": self.llm_temperature,
-                "api_key": self.llm_api_key,
+                "model": self.baml_llm_model,
+                "temperature": self.baml_llm_temperature,
+                "api_key": self.baml_llm_api_key,
+                "endpoint": self.baml_llm_endpoint,
             },
         )
         # Sets the primary client