cognee 0.2.3.dev1__py3-none-any.whl → 0.2.4__py3-none-any.whl

cognee/cli/exceptions.py

@@ -0,0 +1,23 @@
+from typing import Optional
+
+
+class CliCommandException(Exception):
+    """Exception raised by CLI commands with additional context"""
+
+    def __init__(
+        self,
+        message: str,
+        error_code: int = -1,
+        docs_url: Optional[str] = None,
+        raiseable_exception: Optional[Exception] = None,
+    ) -> None:
+        super().__init__(message)
+        self.error_code = error_code
+        self.docs_url = docs_url
+        self.raiseable_exception = raiseable_exception
+
+
+class CliCommandInnerException(Exception):
+    """Inner exception for wrapping other exceptions in CLI context"""
+
+    pass

cognee/cli/minimal_cli.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+"""
+Minimal CLI entry point for cognee that avoids early initialization
+"""
+
+import sys
+import os
+from typing import Any, Sequence
+
+# CRITICAL: Prevent verbose logging initialization for CLI-only usage
+# This must be set before any cognee imports to be effective
+os.environ["COGNEE_MINIMAL_LOGGING"] = "true"
+os.environ["COGNEE_CLI_MODE"] = "true"
+
+
+def get_version() -> str:
+    """Get cognee version without importing the main package"""
+    try:
+        # Try to get version from pyproject.toml first (for development)
+        from pathlib import Path
+
+        pyproject_path = Path(__file__).parent.parent.parent / "pyproject.toml"
+        if pyproject_path.exists():
+            with open(pyproject_path, encoding="utf-8") as f:
+                for line in f:
+                    if line.startswith("version"):
+                        version = line.split("=")[1].strip("'\"\n ")
+                        return f"{version}-local"
+
+        # Fallback to installed package version
+        import importlib.metadata
+
+        return importlib.metadata.version("cognee")
+    except Exception:
+        return "unknown"
+
+
+def get_command_info() -> dict:
+    """Get command information without importing cognee"""
+    return {
+        "add": "Add data to Cognee for knowledge graph processing",
+        "search": "Search and query the knowledge graph for insights, information, and connections",
+        "cognify": "Transform ingested data into a structured knowledge graph",
+        "delete": "Delete data from cognee knowledge base",
+        "config": "Manage cognee configuration settings",
+    }
+
+
+def print_help() -> None:
+    """Print help message with dynamic command descriptions"""
+    commands = get_command_info()
+    command_list = "\n".join(f"    {cmd:<12} {desc}" for cmd, desc in commands.items())
+
+    print(f"""
+usage: cognee [-h] [--version] [--debug] {{{"|".join(commands.keys())}}} ...
+
+Cognee CLI - Manage your knowledge graphs and cognitive processing pipelines.
+
+options:
+  -h, --help            show this help message and exit
+  --version             show program's version number and exit
+  --debug               Enable debug mode to show full stack traces on exceptions
+
+Available commands:
+  {{{",".join(commands.keys())}}}
+{command_list}
+
+For more information on each command, use: cognee <command> --help
+""")
+
+
+def main() -> int:
+    """Minimal CLI main function"""
+    # Handle help and version without any imports - purely static
+    if len(sys.argv) == 1 or (len(sys.argv) == 2 and sys.argv[1] in ["-h", "--help"]):
+        print_help()
+        return 0
+
+    if len(sys.argv) == 2 and sys.argv[1] == "--version":
+        print(f"cognee {get_version()}")
+        return 0
+
+    # For actual commands, import the full CLI with minimal logging
+    try:
+        from cognee.cli._cognee import main as full_main
+
+        return full_main()
+    except Exception as e:
+        if "--debug" in sys.argv:
+            raise
+        print(f"Error: {e}")
+        print("Use --debug for full stack trace")
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

cognee/cli/reference.py

@@ -0,0 +1,26 @@
+from abc import abstractmethod
+from typing import Protocol, Optional
+import argparse
+
+
+class SupportsCliCommand(Protocol):
+    """Protocol for defining one cognee cli command"""
+
+    command_string: str
+    """name of the command"""
+    help_string: str
+    """the help string for argparse"""
+    description: Optional[str]
+    """the more detailed description for argparse, may include markdown for the docs"""
+    docs_url: Optional[str]
+    """the default docs url to be printed in case of an exception"""
+
+    @abstractmethod
+    def configure_parser(self, parser: argparse.ArgumentParser) -> None:
+        """Configures the parser for the given argument"""
+        ...
+
+    @abstractmethod
+    def execute(self, args: argparse.Namespace) -> None:
+        """Executes the command with the given arguments"""
+        ...

cognee/cli/suppress_logging.py

@@ -0,0 +1,12 @@
+"""
+Module to suppress verbose logging before any cognee imports.
+This must be imported before any other cognee modules.
+"""
+
+import os
+
+# Set CLI mode to suppress verbose logging
+os.environ["COGNEE_CLI_MODE"] = "true"
+
+# Also set log level to ERROR for extra safety
+os.environ["LOG_LEVEL"] = "ERROR"

cognee/eval_framework/corpus_builder/corpus_builder_executor.py

@@ -5,7 +5,7 @@ from typing import Optional, Tuple, List, Dict, Union, Any, Callable, Awaitable
 from cognee.eval_framework.benchmark_adapters.benchmark_adapters import BenchmarkAdapter
 from cognee.modules.chunking.TextChunker import TextChunker
 from cognee.modules.pipelines.tasks.task import Task
-from cognee.modules.pipelines import cognee_pipeline
+from cognee.modules.pipelines import run_pipeline
 
 logger = get_logger(level=ERROR)
 
@@ -61,7 +61,7 @@ class CorpusBuilderExecutor:
         await cognee.add(self.raw_corpus)
 
         tasks = await self.task_getter(chunk_size=chunk_size, chunker=chunker)
-        pipeline_run = cognee_pipeline(tasks=tasks)
+        pipeline_run = run_pipeline(tasks=tasks)
 
         async for run_info in pipeline_run:
             print(run_info)

cognee/eval_framework/eval_config.py

@@ -25,7 +25,7 @@ class EvalConfig(BaseSettings):
         "EM",
         "f1",
     ]  # Use only 'correctness' for DirectLLM
-    deepeval_model: str = "gpt-4o-mini"
+    deepeval_model: str = "gpt-5-mini"
 
     # Metrics params
     calculate_metrics: bool = True

cognee/infrastructure/databases/graph/get_graph_engine.py

@@ -21,10 +21,6 @@ async def get_graph_engine() -> GraphDBInterface:
     if hasattr(graph_client, "initialize"):
         await graph_client.initialize()
 
-    # Handle loading of graph for NetworkX
-    if config["graph_database_provider"].lower() == "networkx" and graph_client.graph is None:
-        await graph_client.load_graph_from_file()
-
     return graph_client
 
 
@@ -181,8 +177,7 @@ def create_graph_engine(
             graph_id=graph_identifier,
         )
 
-    from .networkx.adapter import NetworkXAdapter
-
-    graph_client = NetworkXAdapter(filename=graph_file_path)
-
-    return graph_client
+    raise EnvironmentError(
+        f"Unsupported graph database provider: {graph_database_provider}. "
+        f"Supported providers are: {', '.join(list(supported_databases.keys()) + ['neo4j', 'falkordb', 'kuzu', 'kuzu-remote', 'memgraph', 'neptune', 'neptune_analytics'])}"
+    )

cognee/infrastructure/databases/graph/kuzu/adapter.py

@@ -138,8 +138,9 @@ class KuzuAdapter(GraphDBInterface):
 
             s3_file_storage = S3FileStorage("")
 
-            async with self.KUZU_ASYNC_LOCK:
-                self.connection.execute("CHECKPOINT;")
+            if self.connection:
+                async with self.KUZU_ASYNC_LOCK:
+                    self.connection.execute("CHECKPOINT;")
 
             s3_file_storage.s3.put(self.temp_graph_file, self.db_path, recursive=True)
 
@@ -1631,3 +1632,64 @@ class KuzuAdapter(GraphDBInterface):
         """
         result = await self.query(query)
         return [record[0] for record in result] if result else []
+
+    async def get_last_user_interaction_ids(self, limit: int) -> List[str]:
+        """
+        Retrieve the IDs of the most recent CogneeUserInteraction nodes.
+        Parameters:
+        -----------
+        - limit (int): The maximum number of interaction IDs to return.
+        Returns:
+        --------
+        - List[str]: A list of interaction IDs, sorted by created_at descending.
+        """
+
+        query = """
+        MATCH (n)
+        WHERE n.type = 'CogneeUserInteraction'
+        RETURN n.id as id
+        ORDER BY n.created_at DESC
+        LIMIT $limit
+        """
+        rows = await self.query(query, {"limit": limit})
+
+        id_list = [row[0] for row in rows]
+        return id_list
+
+    async def apply_feedback_weight(
+        self,
+        node_ids: List[str],
+        weight: float,
+    ) -> None:
+        """
+        Increment `feedback_weight` inside r.properties JSON for edges where
+        relationship_name = 'used_graph_element_to_answer'.
+
+        """
+        # Step 1: fetch matching edges
+        query = """
+            MATCH (n:Node)-[r:EDGE]->()
+            WHERE n.id IN $node_ids AND r.relationship_name = 'used_graph_element_to_answer'
+            RETURN r.properties, n.id
+            """
+        results = await self.query(query, {"node_ids": node_ids})
+
+        # Step 2: update JSON client-side
+        updates = []
+        for props_json, source_id in results:
+            try:
+                props = json.loads(props_json) if props_json else {}
+            except json.JSONDecodeError:
+                props = {}
+
+            props["feedback_weight"] = props.get("feedback_weight", 0) + weight
+            updates.append((source_id, json.dumps(props)))
+
+        # Step 3: write back
+        for node_id, new_props in updates:
+            update_query = """
+                MATCH (n:Node)-[r:EDGE]->()
+                WHERE n.id = $node_id AND r.relationship_name = 'used_graph_element_to_answer'
+                SET r.properties = $props
+                """
+            await self.query(update_query, {"node_id": node_id, "props": new_props})

cognee/infrastructure/databases/graph/neo4j_driver/adapter.py

@@ -1322,3 +1322,52 @@ class Neo4jAdapter(GraphDBInterface):
         """
         result = await self.query(query)
         return [record["n"] for record in result] if result else []
+
+    async def get_last_user_interaction_ids(self, limit: int) -> List[str]:
+        """
+        Retrieve the IDs of the most recent CogneeUserInteraction nodes.
+        Parameters:
+        -----------
+        - limit (int): The maximum number of interaction IDs to return.
+        Returns:
+        --------
+        - List[str]: A list of interaction IDs, sorted by created_at descending.
+        """
+
+        query = """
+        MATCH (n)
+        WHERE n.type = 'CogneeUserInteraction'
+        RETURN n.id as id
+        ORDER BY n.created_at DESC
+        LIMIT $limit
+        """
+        rows = await self.query(query, {"limit": limit})
+
+        id_list = [row["id"] for row in rows if "id" in row]
+        return id_list
+
+    async def apply_feedback_weight(
+        self,
+        node_ids: List[str],
+        weight: float,
+    ) -> None:
+        """
+        Increment `feedback_weight` on relationships `:used_graph_element_to_answer`
+        outgoing from nodes whose `id` is in `node_ids`.
+
+        Args:
+            node_ids: List of node IDs to match.
+            weight: Amount to add to `r.feedback_weight` (can be negative).
+
+        Side effects:
+            Updates relationship property `feedback_weight`, defaulting missing values to 0.
+        """
+        query = """
+        MATCH (n)-[r]->()
+        WHERE n.id IN $node_ids AND r.relationship_name = 'used_graph_element_to_answer'
+        SET r.feedback_weight = coalesce(r.feedback_weight, 0) + $weight
+        """
+        await self.query(
+            query,
+            params={"weight": float(weight), "node_ids": list(node_ids)},
+        )

cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py

@@ -41,11 +41,11 @@ class FastembedEmbeddingEngine(EmbeddingEngine):
         self,
         model: Optional[str] = "openai/text-embedding-3-large",
         dimensions: Optional[int] = 3072,
-        max_tokens: int = 512,
+        max_completion_tokens: int = 512,
     ):
         self.model = model
         self.dimensions = dimensions
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
         self.tokenizer = self.get_tokenizer()
         # self.retry_count = 0
         self.embedding_model = TextEmbedding(model_name=model)
@@ -112,7 +112,9 @@ class FastembedEmbeddingEngine(EmbeddingEngine):
         """
         logger.debug("Loading tokenizer for FastembedEmbeddingEngine...")
 
-        tokenizer = TikTokenTokenizer(model="gpt-4o", max_tokens=self.max_tokens)
+        tokenizer = TikTokenTokenizer(
+            model="gpt-4o", max_completion_tokens=self.max_completion_tokens
+        )
 
         logger.debug("Tokenizer loaded for for FastembedEmbeddingEngine")
         return tokenizer

cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py

@@ -57,7 +57,7 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
         api_key: str = None,
         endpoint: str = None,
         api_version: str = None,
-        max_tokens: int = 512,
+        max_completion_tokens: int = 512,
     ):
         self.api_key = api_key
         self.endpoint = endpoint
@@ -65,7 +65,7 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
         self.provider = provider
         self.model = model
         self.dimensions = dimensions
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
         self.tokenizer = self.get_tokenizer()
         self.retry_count = 0
 
@@ -179,20 +179,29 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
         model = self.model.split("/")[-1]
 
         if "openai" in self.provider.lower():
-            tokenizer = TikTokenTokenizer(model=model, max_tokens=self.max_tokens)
+            tokenizer = TikTokenTokenizer(
+                model=model, max_completion_tokens=self.max_completion_tokens
+            )
         elif "gemini" in self.provider.lower():
-            tokenizer = GeminiTokenizer(model=model, max_tokens=self.max_tokens)
+            tokenizer = GeminiTokenizer(
+                model=model, max_completion_tokens=self.max_completion_tokens
+            )
         elif "mistral" in self.provider.lower():
-            tokenizer = MistralTokenizer(model=model, max_tokens=self.max_tokens)
+            tokenizer = MistralTokenizer(
+                model=model, max_completion_tokens=self.max_completion_tokens
+            )
         else:
             try:
                 tokenizer = HuggingFaceTokenizer(
-                    model=self.model.replace("hosted_vllm/", ""), max_tokens=self.max_tokens
+                    model=self.model.replace("hosted_vllm/", ""),
+                    max_completion_tokens=self.max_completion_tokens,
                 )
             except Exception as e:
                 logger.warning(f"Could not get tokenizer from HuggingFace due to: {e}")
                 logger.info("Switching to TikToken default tokenizer.")
-                tokenizer = TikTokenTokenizer(model=None, max_tokens=self.max_tokens)
+                tokenizer = TikTokenTokenizer(
+                    model=None, max_completion_tokens=self.max_completion_tokens
+                )
 
         logger.debug(f"Tokenizer loaded for model: {self.model}")
         return tokenizer

cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

@@ -30,7 +30,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
     Instance variables:
     - model
     - dimensions
-    - max_tokens
+    - max_completion_tokens
     - endpoint
     - mock
     - huggingface_tokenizer_name
@@ -39,7 +39,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
 
     model: str
     dimensions: int
-    max_tokens: int
+    max_completion_tokens: int
     endpoint: str
     mock: bool
     huggingface_tokenizer_name: str
@@ -50,13 +50,13 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         self,
         model: Optional[str] = "avr/sfr-embedding-mistral:latest",
         dimensions: Optional[int] = 1024,
-        max_tokens: int = 512,
+        max_completion_tokens: int = 512,
         endpoint: Optional[str] = "http://localhost:11434/api/embeddings",
         huggingface_tokenizer: str = "Salesforce/SFR-Embedding-Mistral",
     ):
         self.model = model
         self.dimensions = dimensions
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
         self.endpoint = endpoint
         self.huggingface_tokenizer_name = huggingface_tokenizer
         self.tokenizer = self.get_tokenizer()
@@ -132,7 +132,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         """
         logger.debug("Loading HuggingfaceTokenizer for OllamaEmbeddingEngine...")
         tokenizer = HuggingFaceTokenizer(
-            model=self.huggingface_tokenizer_name, max_tokens=self.max_tokens
+            model=self.huggingface_tokenizer_name, max_completion_tokens=self.max_completion_tokens
         )
         logger.debug("Tokenizer loaded for OllamaEmbeddingEngine")
         return tokenizer

cognee/infrastructure/databases/vector/embeddings/config.py

@@ -18,7 +18,7 @@ class EmbeddingConfig(BaseSettings):
     embedding_endpoint: Optional[str] = None
     embedding_api_key: Optional[str] = None
     embedding_api_version: Optional[str] = None
-    embedding_max_tokens: Optional[int] = 8191
+    embedding_max_completion_tokens: Optional[int] = 8191
     huggingface_tokenizer: Optional[str] = None
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
@@ -38,7 +38,7 @@ class EmbeddingConfig(BaseSettings):
             "embedding_endpoint": self.embedding_endpoint,
             "embedding_api_key": self.embedding_api_key,
             "embedding_api_version": self.embedding_api_version,
-            "embedding_max_tokens": self.embedding_max_tokens,
+            "embedding_max_completion_tokens": self.embedding_max_completion_tokens,
             "huggingface_tokenizer": self.huggingface_tokenizer,
         }
 

cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py

@@ -27,7 +27,7 @@ def get_embedding_engine() -> EmbeddingEngine:
         config.embedding_provider,
         config.embedding_model,
         config.embedding_dimensions,
-        config.embedding_max_tokens,
+        config.embedding_max_completion_tokens,
         config.embedding_endpoint,
         config.embedding_api_key,
         config.embedding_api_version,
@@ -41,7 +41,7 @@ def create_embedding_engine(
     embedding_provider,
     embedding_model,
     embedding_dimensions,
-    embedding_max_tokens,
+    embedding_max_completion_tokens,
     embedding_endpoint,
     embedding_api_key,
     embedding_api_version,
@@ -58,7 +58,7 @@ def create_embedding_engine(
           'ollama', or another supported provider.
         - embedding_model: The model to be used for the embedding engine.
         - embedding_dimensions: The number of dimensions for the embeddings.
-        - embedding_max_tokens: The maximum number of tokens for the embeddings.
+        - embedding_max_completion_tokens: The maximum number of tokens for the embeddings.
         - embedding_endpoint: The endpoint for the embedding service, relevant for certain
           providers.
         - embedding_api_key: API key to authenticate with the embedding service, if
@@ -81,7 +81,7 @@ def create_embedding_engine(
         return FastembedEmbeddingEngine(
             model=embedding_model,
             dimensions=embedding_dimensions,
-            max_tokens=embedding_max_tokens,
+            max_completion_tokens=embedding_max_completion_tokens,
         )
 
     if embedding_provider == "ollama":
@@ -90,7 +90,7 @@ def create_embedding_engine(
         return OllamaEmbeddingEngine(
             model=embedding_model,
             dimensions=embedding_dimensions,
-            max_tokens=embedding_max_tokens,
+            max_completion_tokens=embedding_max_completion_tokens,
             endpoint=embedding_endpoint,
             huggingface_tokenizer=huggingface_tokenizer,
         )
@@ -104,5 +104,5 @@ def create_embedding_engine(
         api_version=embedding_api_version,
         model=embedding_model,
         dimensions=embedding_dimensions,
-        max_tokens=embedding_max_tokens,
+        max_completion_tokens=embedding_max_completion_tokens,
     )

cognee/infrastructure/files/utils/get_data_file_path.py

@@ -5,19 +5,24 @@ from urllib.parse import urlparse
 def get_data_file_path(file_path: str):
     # Check if this is a file URI BEFORE normalizing (which corrupts URIs)
     if file_path.startswith("file://"):
+        # Remove first occurrence of file:// prefix
+        pure_file_path = file_path.replace("file://", "", 1)
         # Normalize the file URI for Windows - replace backslashes with forward slashes
-        normalized_file_uri = os.path.normpath(file_path)
+        normalized_file_uri = os.path.normpath(pure_file_path)
 
-        parsed_url = urlparse(normalized_file_uri)
-
-        # Convert URI path to file system path
+        # Convert path to proper file system path
         if os.name == "nt":  # Windows
             # Handle Windows drive letters correctly
-            fs_path = parsed_url.path
-            if fs_path.startswith("/") and len(fs_path) > 1 and fs_path[2] == ":":
-                fs_path = fs_path[1:]  # Remove leading slash for Windows drive paths
-        else:  # Unix-like systems
-            fs_path = parsed_url.path
+            fs_path = normalized_file_uri
+            if (
+                (fs_path.startswith("/") or fs_path.startswith("\\"))
+                and len(fs_path) > 1
+                and fs_path[2] == ":"
+            ):
+                fs_path = fs_path[1:]
+        else:
+            # Unix - like systems
+            fs_path = normalized_file_uri
 
         # Now split the actual filesystem path
         actual_fs_path = os.path.normpath(fs_path)

cognee/infrastructure/files/utils/get_file_metadata.py

@@ -1,6 +1,7 @@
 import io
 import os.path
 from typing import BinaryIO, TypedDict
+from pathlib import Path
 
 from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.files.utils.get_file_content_hash import get_file_content_hash
@@ -55,7 +56,7 @@ async def get_file_metadata(file: BinaryIO) -> FileMetadata:
     file_type = guess_file_type(file)
 
     file_path = getattr(file, "name", None) or getattr(file, "full_name", None)
-    file_name = str(file_path).split("/")[-1].split(".")[0] if file_path else None
+    file_name = Path(file_path).stem if file_path else None
 
     # Get file size
     pos = file.tell()  # remember current pointer

cognee/infrastructure/llm/LLMGateway.py

@@ -1,6 +1,5 @@
-from typing import Type
+from typing import Type, Optional, Coroutine
 from pydantic import BaseModel
-from typing import Coroutine
 from cognee.infrastructure.llm import get_llm_config
 
 
@@ -79,7 +78,10 @@ class LLMGateway:
 
     @staticmethod
     def extract_content_graph(
-        content: str, response_model: Type[BaseModel], mode: str = "simple"
+        content: str,
+        response_model: Type[BaseModel],
+        mode: str = "simple",
+        custom_prompt: Optional[str] = None,
     ) -> Coroutine:
         llm_config = get_llm_config()
         if llm_config.structured_output_framework.upper() == "BAML":
@@ -87,13 +89,20 @@ class LLMGateway:
                 extract_content_graph,
             )
 
-            return extract_content_graph(content=content, response_model=response_model, mode=mode)
+            return extract_content_graph(
+                content=content,
+                response_model=response_model,
+                mode=mode,
+                custom_prompt=custom_prompt,
+            )
         else:
             from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.extraction import (
                 extract_content_graph,
             )
 
-            return extract_content_graph(content=content, response_model=response_model)
+            return extract_content_graph(
+                content=content, response_model=response_model, custom_prompt=custom_prompt
+            )
 
     @staticmethod
     def extract_categories(content: str, response_model: Type[BaseModel]) -> Coroutine:

cognee/infrastructure/llm/config.py

@@ -18,7 +18,7 @@ class LLMConfig(BaseSettings):
     - llm_api_version
     - llm_temperature
     - llm_streaming
-    - llm_max_tokens
+    - llm_max_completion_tokens
     - transcription_model
     - graph_prompt_path
     - llm_rate_limit_enabled
@@ -35,16 +35,16 @@ class LLMConfig(BaseSettings):
 
     structured_output_framework: str = "instructor"
     llm_provider: str = "openai"
-    llm_model: str = "gpt-4o-mini"
+    llm_model: str = "gpt-5-mini"
     llm_endpoint: str = ""
     llm_api_key: Optional[str] = None
     llm_api_version: Optional[str] = None
     llm_temperature: float = 0.0
     llm_streaming: bool = False
-    llm_max_tokens: int = 16384
+    llm_max_completion_tokens: int = 16384
 
     baml_llm_provider: str = "openai"
-    baml_llm_model: str = "gpt-4o-mini"
+    baml_llm_model: str = "gpt-5-mini"
     baml_llm_endpoint: str = ""
     baml_llm_api_key: Optional[str] = None
     baml_llm_temperature: float = 0.0
@@ -171,7 +171,7 @@ class LLMConfig(BaseSettings):
             "api_version": self.llm_api_version,
             "temperature": self.llm_temperature,
             "streaming": self.llm_streaming,
-            "max_tokens": self.llm_max_tokens,
+            "max_completion_tokens": self.llm_max_completion_tokens,
             "transcription_model": self.transcription_model,
             "graph_prompt_path": self.graph_prompt_path,
             "rate_limit_enabled": self.llm_rate_limit_enabled,