cognee 0.3.4.dev3__py3-none-any.whl → 0.3.5__py3-none-any.whl

cognee/infrastructure/llm/LLMGateway.py

@@ -9,24 +9,30 @@ class LLMGateway:
     Class used as a namespace for LLM related functions, should not be instantiated, all methods are static.
     """
 
-    @staticmethod
-    def render_prompt(filename: str, context: dict, base_directory: str = None):
-        from cognee.infrastructure.llm.prompts import render_prompt
-
-        return render_prompt(filename=filename, context=context, base_directory=base_directory)
-
     @staticmethod
     def acreate_structured_output(
         text_input: str, system_prompt: str, response_model: Type[BaseModel]
     ) -> Coroutine:
-        from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.get_llm_client import (
-            get_llm_client,
-        )
+        llm_config = get_llm_config()
+        if llm_config.structured_output_framework.upper() == "BAML":
+            from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction import (
+                acreate_structured_output,
+            )
 
-        llm_client = get_llm_client()
-        return llm_client.acreate_structured_output(
-            text_input=text_input, system_prompt=system_prompt, response_model=response_model
-        )
+            return acreate_structured_output(
+                text_input=text_input,
+                system_prompt=system_prompt,
+                response_model=response_model,
+            )
+        else:
+            from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.get_llm_client import (
+                get_llm_client,
+            )
+
+            llm_client = get_llm_client()
+            return llm_client.acreate_structured_output(
+                text_input=text_input, system_prompt=system_prompt, response_model=response_model
+            )
 
     @staticmethod
     def create_structured_output(
@@ -58,107 +64,3 @@ class LLMGateway:
 
         llm_client = get_llm_client()
         return llm_client.transcribe_image(input=input)
-
-    @staticmethod
-    def show_prompt(text_input: str, system_prompt: str) -> str:
-        from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.get_llm_client import (
-            get_llm_client,
-        )
-
-        llm_client = get_llm_client()
-        return llm_client.show_prompt(text_input=text_input, system_prompt=system_prompt)
-
-    @staticmethod
-    def read_query_prompt(prompt_file_name: str, base_directory: str = None):
-        from cognee.infrastructure.llm.prompts import (
-            read_query_prompt,
-        )
-
-        return read_query_prompt(prompt_file_name=prompt_file_name, base_directory=base_directory)
-
-    @staticmethod
-    def extract_content_graph(
-        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":
-            from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction import (
-                extract_content_graph,
-            )
-
-            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, custom_prompt=custom_prompt
-            )
-
-    @staticmethod
-    def extract_categories(content: str, response_model: Type[BaseModel]) -> Coroutine:
-        # TODO: Add BAML version of category and extraction and update function
-        from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.extraction import (
-            extract_categories,
-        )
-
-        return extract_categories(content=content, response_model=response_model)
-
-    @staticmethod
-    def extract_code_summary(content: str) -> Coroutine:
-        llm_config = get_llm_config()
-        if llm_config.structured_output_framework.upper() == "BAML":
-            from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction import (
-                extract_code_summary,
-            )
-
-            return extract_code_summary(content=content)
-        else:
-            from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.extraction import (
-                extract_code_summary,
-            )
-
-            return extract_code_summary(content=content)
-
-    @staticmethod
-    def extract_summary(content: str, response_model: Type[BaseModel]) -> Coroutine:
-        llm_config = get_llm_config()
-        if llm_config.structured_output_framework.upper() == "BAML":
-            from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction import (
-                extract_summary,
-            )
-
-            return extract_summary(content=content, response_model=response_model)
-        else:
-            from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.extraction import (
-                extract_summary,
-            )
-
-            return extract_summary(content=content, response_model=response_model)
-
-    @staticmethod
-    def extract_event_graph(content: str, response_model: Type[BaseModel]) -> Coroutine:
-        # TODO: Add BAML version of category and extraction and update function (consulted with Igor)
-        from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.extraction import (
-            extract_event_graph,
-        )
-
-        return extract_event_graph(content=content, response_model=response_model)
-
-    @staticmethod
-    def extract_event_entities(content: str, response_model: Type[BaseModel]) -> Coroutine:
-        # TODO: Add BAML version of category and extraction and update function (consulted with Igor)
-        from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.extraction import (
-            extract_event_entities,
-        )
-
-        return extract_event_entities(content=content, response_model=response_model)

cognee/infrastructure/llm/config.py

@@ -1,9 +1,13 @@
 import os
-from typing import Optional, ClassVar
+from typing import Optional, ClassVar, Any
 from functools import lru_cache
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from pydantic import model_validator
-from baml_py import ClientRegistry
+
+try:
+    from baml_py import ClientRegistry
+except ImportError:
+    ClientRegistry = None
 
 
 class LLMConfig(BaseSettings):
@@ -35,7 +39,7 @@ class LLMConfig(BaseSettings):
 
     structured_output_framework: str = "instructor"
     llm_provider: str = "openai"
-    llm_model: str = "openai/gpt-4o-mini"
+    llm_model: str = "openai/gpt-5-mini"
     llm_endpoint: str = ""
     llm_api_key: Optional[str] = None
     llm_api_version: Optional[str] = None
@@ -44,7 +48,7 @@ class LLMConfig(BaseSettings):
     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
@@ -65,25 +69,36 @@ class LLMConfig(BaseSettings):
     fallback_endpoint: str = ""
     fallback_model: str = ""
 
-    baml_registry: ClassVar[ClientRegistry] = ClientRegistry()
+    baml_registry: Optional[Any] = None
 
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
     def model_post_init(self, __context) -> None:
         """Initialize the BAML registry after the model is created."""
-        self.baml_registry.add_llm_client(
-            name=self.baml_llm_provider,
-            provider=self.baml_llm_provider,
-            options={
+        # Check if BAML is selected as structured output framework but not available
+        if self.structured_output_framework.lower() == "baml" and ClientRegistry is None:
+            raise ImportError(
+                "BAML is selected as structured output framework but not available. "
+                "Please install with 'pip install cognee\"[baml]\"' to use BAML extraction features."
+            )
+        elif self.structured_output_framework.lower() == "baml" and ClientRegistry is not None:
+            self.baml_registry = ClientRegistry()
+
+            raw_options = {
                 "model": self.baml_llm_model,
                 "temperature": self.baml_llm_temperature,
                 "api_key": self.baml_llm_api_key,
                 "base_url": self.baml_llm_endpoint,
                 "api_version": self.baml_llm_api_version,
-            },
-        )
-        # Sets the primary client
-        self.baml_registry.set_primary(self.baml_llm_provider)
+            }
+
+            # Note: keep the item only when the value is not None or an empty string (they would override baml default values)
+            options = {k: v for k, v in raw_options.items() if v not in (None, "")}
+            self.baml_registry.add_llm_client(
+                name=self.baml_llm_provider, provider=self.baml_llm_provider, options=options
+            )
+            # Sets the primary client
+            self.baml_registry.set_primary(self.baml_llm_provider)
 
     @model_validator(mode="after")
     def ensure_env_vars_for_ollama(self) -> "LLMConfig":

cognee/infrastructure/llm/{structured_output_framework/litellm_instructor/extraction → extraction}/extract_categories.py

@@ -1,11 +1,12 @@
 from typing import Type
 from pydantic import BaseModel
 
+from cognee.infrastructure.llm.prompts import read_query_prompt
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
 
 
 async def extract_categories(content: str, response_model: Type[BaseModel]):
-    system_prompt = LLMGateway.read_query_prompt("classify_content.txt")
+    system_prompt = read_query_prompt("classify_content.txt")
 
     llm_output = await LLMGateway.acreate_structured_output(content, system_prompt, response_model)
 

cognee/infrastructure/llm/{structured_output_framework/litellm_instructor/extraction → extraction}/extract_event_entities.py

@@ -1,6 +1,7 @@
 import os
-from typing import List, Type
+from typing import Type
 from pydantic import BaseModel
+from cognee.infrastructure.llm.prompts.render_prompt import render_prompt
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.infrastructure.llm.config import (
     get_llm_config,
@@ -35,7 +36,7 @@ async def extract_event_entities(content: str, response_model: Type[BaseModel]):
     else:
         base_directory = None
 
-    system_prompt = LLMGateway.render_prompt(prompt_path, {}, base_directory=base_directory)
+    system_prompt = render_prompt(prompt_path, {}, base_directory=base_directory)
 
     content_graph = await LLMGateway.acreate_structured_output(
         content, system_prompt, response_model

cognee/infrastructure/llm/{structured_output_framework/litellm_instructor/extraction → extraction}/extract_summary.py

@@ -2,7 +2,8 @@ from cognee.shared.logging_utils import get_logger
 import os
 from typing import Type
 
-from instructor.exceptions import InstructorRetryException
+from instructor.core import InstructorRetryException
+from cognee.infrastructure.llm.prompts import read_query_prompt
 from pydantic import BaseModel
 
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
@@ -25,7 +26,7 @@ def get_mock_summarized_code():
 
 
 async def extract_summary(content: str, response_model: Type[BaseModel]):
-    system_prompt = LLMGateway.read_query_prompt("summarize_content.txt")
+    system_prompt = read_query_prompt("summarize_content.txt")
 
     llm_output = await LLMGateway.acreate_structured_output(content, system_prompt, response_model)
 

cognee/infrastructure/llm/{structured_output_framework/litellm_instructor/extraction → extraction}/knowledge_graph/extract_content_graph.py

@@ -2,6 +2,7 @@ import os
 from typing import Type, Optional
 from pydantic import BaseModel
 
+from cognee.infrastructure.llm.prompts import render_prompt
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.infrastructure.llm.config import (
     get_llm_config,
@@ -26,7 +27,7 @@ async def extract_content_graph(
         else:
             base_directory = None
 
-        system_prompt = LLMGateway.render_prompt(prompt_path, {}, base_directory=base_directory)
+        system_prompt = render_prompt(prompt_path, {}, base_directory=base_directory)
 
     content_graph = await LLMGateway.acreate_structured_output(
         content, system_prompt, response_model

cognee/infrastructure/llm/{structured_output_framework/litellm_instructor/extraction → extraction}/knowledge_graph/extract_event_graph.py

@@ -1,8 +1,9 @@
 import os
 from pydantic import BaseModel
 from typing import Type
-from cognee.infrastructure.llm.LLMGateway import LLMGateway
 
+from cognee.infrastructure.llm.prompts import render_prompt
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.infrastructure.llm.config import (
     get_llm_config,
 )
@@ -37,7 +38,7 @@ async def extract_event_graph(content: str, response_model: Type[BaseModel]):
     else:
         base_directory = None
 
-    system_prompt = LLMGateway.render_prompt(prompt_path, {}, base_directory=base_directory)
+    system_prompt = render_prompt(prompt_path, {}, base_directory=base_directory)
 
     content_graph = await LLMGateway.acreate_structured_output(
         content, system_prompt, response_model

cognee/infrastructure/llm/prompts/read_query_prompt.py

@@ -26,6 +26,7 @@ def read_query_prompt(prompt_file_name: str, base_directory: str = None):
         read due to an error.
     """
     logger = get_logger(level=ERROR)
+
     try:
         if base_directory is None:
             base_directory = get_absolute_path("./infrastructure/llm/prompts")
@@ -35,8 +36,8 @@ def read_query_prompt(prompt_file_name: str, base_directory: str = None):
         with open(file_path, "r", encoding="utf-8") as file:
             return file.read()
     except FileNotFoundError:
-        logger.error(f"Error: Prompt file not found. Attempted to read: %s {file_path}")
+        logger.error(f"Error: Prompt file not found. Attempted to read: {file_path}")
         return None
     except Exception as e:
-        logger.error(f"An error occurred: %s {e}")
+        logger.error(f"An error occurred: {e}")
         return None

cognee/infrastructure/llm/prompts/show_prompt.py

@@ -0,0 +1,35 @@
+from cognee.infrastructure.llm.exceptions import MissingSystemPromptPathError
+from cognee.infrastructure.llm.prompts import read_query_prompt
+
+
+def show_prompt(text_input: str, system_prompt: str) -> str:
+    """
+    Format and display the prompt for a user query.
+
+    This method formats the prompt using the provided user input and system prompt,
+    returning a string representation. Raises MissingSystemPromptPathError if the system prompt is not
+    provided.
+
+    Parameters:
+    -----------
+
+        - text_input (str): The input text provided by the user.
+        - system_prompt (str): The system's prompt to guide the model's response.
+
+    Returns:
+    --------
+
+        - str: A formatted string representing the user input and system prompt.
+    """
+    if not text_input:
+        text_input = "No user input provided."
+    if not system_prompt:
+        raise MissingSystemPromptPathError()
+    system_prompt = read_query_prompt(system_prompt)
+
+    formatted_prompt = (
+        f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n"""
+        if system_prompt
+        else None
+    )
+    return formatted_prompt

cognee/infrastructure/llm/prompts/test.txt

@@ -0,0 +1 @@
+Respond with: test

cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py

@@ -10,7 +10,7 @@
 # BAML files and re-generate this code using: baml-cli generate
 # baml-cli is available with the baml package.
 
-__version__ = "0.201.0"
+__version__ = "0.206.0"
 
 try:
     from baml_py.safe_import import EnsureBamlPyImport
@@ -39,7 +39,7 @@ with EnsureBamlPyImport(__version__) as e:
     from . import config
     from .config import reset_baml_env_vars
 
-    from .sync_client import b
+    from .async_client import b
 
 
 # FOR LEGACY COMPATIBILITY, expose "partial_types" as an alias for "stream_types"