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

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

@@ -15,38 +15,6 @@ from . import stream_types
 
 
 type_map = {
-    "types.AudioContent": types.AudioContent,
-    "stream_types.AudioContent": stream_types.AudioContent,
-    "types.ContentLabel": types.ContentLabel,
-    "stream_types.ContentLabel": stream_types.ContentLabel,
-    "types.DefaultContentPrediction": types.DefaultContentPrediction,
-    "stream_types.DefaultContentPrediction": stream_types.DefaultContentPrediction,
-    "types.DynamicKnowledgeGraph": types.DynamicKnowledgeGraph,
-    "stream_types.DynamicKnowledgeGraph": stream_types.DynamicKnowledgeGraph,
-    "types.Edge": types.Edge,
-    "stream_types.Edge": stream_types.Edge,
-    "types.ImageContent": types.ImageContent,
-    "stream_types.ImageContent": stream_types.ImageContent,
-    "types.KnowledgeGraph": types.KnowledgeGraph,
-    "stream_types.KnowledgeGraph": stream_types.KnowledgeGraph,
-    "types.Model3DContent": types.Model3DContent,
-    "stream_types.Model3DContent": stream_types.Model3DContent,
-    "types.MultimediaContent": types.MultimediaContent,
-    "stream_types.MultimediaContent": stream_types.MultimediaContent,
-    "types.Node": types.Node,
-    "stream_types.Node": stream_types.Node,
-    "types.ProceduralContent": types.ProceduralContent,
-    "stream_types.ProceduralContent": stream_types.ProceduralContent,
-    "types.SummarizedClass": types.SummarizedClass,
-    "stream_types.SummarizedClass": stream_types.SummarizedClass,
-    "types.SummarizedCode": types.SummarizedCode,
-    "stream_types.SummarizedCode": stream_types.SummarizedCode,
-    "types.SummarizedContent": types.SummarizedContent,
-    "stream_types.SummarizedContent": stream_types.SummarizedContent,
-    "types.SummarizedFunction": types.SummarizedFunction,
-    "stream_types.SummarizedFunction": stream_types.SummarizedFunction,
-    "types.TextContent": types.TextContent,
-    "stream_types.TextContent": stream_types.TextContent,
-    "types.VideoContent": types.VideoContent,
-    "stream_types.VideoContent": stream_types.VideoContent,
+    "types.ResponseModel": types.ResponseModel,
+    "stream_types.ResponseModel": stream_types.ResponseModel,
 }

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

@@ -48,119 +48,14 @@ def all_succeeded(checks: typing.Dict[CheckName, Check]) -> bool:
 # #########################################################################
 
 # #########################################################################
-# Generated classes (17)
+# Generated classes (1)
 # #########################################################################
 
 
-class AudioContent(BaseModel):
-    type: str
-    subclass: typing.List[str]
-
-
-class ContentLabel(BaseModel):
-    content_type: typing.Union[
-        typing_extensions.Literal["text"],
-        typing_extensions.Literal["audio"],
-        typing_extensions.Literal["image"],
-        typing_extensions.Literal["video"],
-        typing_extensions.Literal["multimedia"],
-        typing_extensions.Literal["3d_model"],
-        typing_extensions.Literal["procedural"],
-    ]
-    type: str
-    subclass: typing.List[str]
-
-
-class DefaultContentPrediction(BaseModel):
-    label: "ContentLabel"
-
-
-class DynamicKnowledgeGraph(BaseModel):
+class ResponseModel(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
-class Edge(BaseModel):
-    # doc string for edge
-    # doc string for source_node_id
-
-    source_node_id: str
-    target_node_id: str
-    relationship_name: str
-
-
-class ImageContent(BaseModel):
-    type: str
-    subclass: typing.List[str]
-
-
-class KnowledgeGraph(BaseModel):
-    nodes: typing.List["Node"]
-    edges: typing.List["Edge"]
-
-
-class Model3DContent(BaseModel):
-    type: str
-    subclass: typing.List[str]
-
-
-class MultimediaContent(BaseModel):
-    type: str
-    subclass: typing.List[str]
-
-
-class Node(BaseModel):
-    model_config = ConfigDict(extra="allow")
-    id: str
-    name: str
-    type: str
-    description: str
-
-
-class ProceduralContent(BaseModel):
-    type: str
-    subclass: typing.List[str]
-
-
-class SummarizedClass(BaseModel):
-    name: str
-    description: str
-    methods: typing.Optional[typing.List["SummarizedFunction"]] = None
-    decorators: typing.Optional[typing.List[str]] = None
-
-
-class SummarizedCode(BaseModel):
-    high_level_summary: str
-    key_features: typing.List[str]
-    imports: typing.List[str]
-    constants: typing.List[str]
-    classes: typing.List["SummarizedClass"]
-    functions: typing.List["SummarizedFunction"]
-    workflow_description: typing.Optional[str] = None
-
-
-class SummarizedContent(BaseModel):
-    summary: str
-    description: str
-
-
-class SummarizedFunction(BaseModel):
-    name: str
-    description: str
-    inputs: typing.Optional[typing.List[str]] = None
-    outputs: typing.Optional[typing.List[str]] = None
-    decorators: typing.Optional[typing.List[str]] = None
-
-
-class TextContent(BaseModel):
-    type: str
-    subclass: typing.List[str]
-
-
-class VideoContent(BaseModel):
-    type: str
-    subclass: typing.List[str]
-
-
 # #########################################################################
 # Generated type aliases (0)
 # #########################################################################

cognee/infrastructure/llm/structured_output_framework/baml/baml_src/acreate_structured_output.baml

@@ -0,0 +1,26 @@
+class ResponseModel {
+    @@dynamic
+}
+
+// OpenAI client with environment model selection
+client<llm> OpenAI {
+    provider openai
+    options {
+        model client_registry.model
+        api_key client_registry.api_key
+    }
+}
+
+function AcreateStructuredOutput(
+    text_input: string,
+    system_prompt: string,
+) -> ResponseModel {
+    client OpenAI
+
+    prompt #"
+        {{ system_prompt }}
+        {{ ctx.output_format }}
+        {{ _.role('user') }}
+        {{ text_input }}
+    "#
+}

cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/__init__.py

@@ -1,2 +1 @@
-from .knowledge_graph.extract_content_graph import extract_content_graph
-from .extract_summary import extract_summary, extract_code_summary
+from .acreate_structured_output import acreate_structured_output

cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/acreate_structured_output.py

@@ -0,0 +1,76 @@
+import asyncio
+from typing import Type
+from cognee.shared.logging_utils import get_logger
+
+from cognee.infrastructure.llm.config import get_llm_config
+from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction.create_dynamic_baml_type import (
+    create_dynamic_baml_type,
+)
+from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.type_builder import (
+    TypeBuilder,
+)
+from cognee.infrastructure.llm.structured_output_framework.baml.baml_client import b
+from pydantic import BaseModel
+
+
+logger = get_logger()
+
+
+async def acreate_structured_output(
+    text_input: str, system_prompt: str, response_model: Type[BaseModel]
+):
+    """
+    Generate a response from a user query.
+
+    This method asynchronously creates structured output by sending a request through BAML
+    using the provided parameters to generate a completion based on the user input and
+    system prompt.
+
+    Parameters:
+    -----------
+
+        - text_input (str): The input text provided by the user for generating a response.
+        - system_prompt (str): The system's prompt to guide the model's response.
+        - response_model (Type[BaseModel]): The expected model type for the response.
+
+    Returns:
+    --------
+
+        - BaseModel: A structured output generated by the model, returned as an instance of
+          BaseModel.
+    """
+    config = get_llm_config()
+
+    # Dynamically create BAML response model
+    tb = TypeBuilder()
+    type_builder = create_dynamic_baml_type(tb, tb.ResponseModel, response_model)
+
+    result = await b.AcreateStructuredOutput(
+        text_input=text_input,
+        system_prompt=system_prompt,
+        baml_options={"client_registry": config.baml_registry, "tb": type_builder},
+    )
+
+    # Transform BAML response to proper pydantic reponse model
+    if response_model is str:
+        # Note: when a response model is set to string in python, result is stored in text property in the BAML response model
+        return str(result.text)
+    return response_model.model_validate(result.dict())
+
+
+if __name__ == "__main__":
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    try:
+        from typing import Optional, Dict, Any, List, Literal
+
+        # Models for representing different entities
+        class TestModel(BaseModel):
+            type: str
+            source: Optional[str] = None
+            target: Optional[str] = None
+            properties: Optional[Dict[str, List[str]]] = None
+
+        loop.run_until_complete(acreate_structured_output("TEST", "THIS IS A TEST", TestModel))
+    finally:
+        loop.run_until_complete(loop.shutdown_asyncgens())

cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/create_dynamic_baml_type.py

@@ -0,0 +1,122 @@
+from typing import Union
+from enum import Enum
+from datetime import datetime
+from typing import get_origin, get_args
+
+from baml_py.baml_py import ClassBuilder
+from cognee.shared.logging_utils import get_logger
+
+logger = get_logger()
+
+
+def create_dynamic_baml_type(tb, baml_model, pydantic_model):
+    if pydantic_model is str:
+        baml_model.add_property("text", tb.string())
+        return tb
+
+    def map_type(field_type, field_info):
+        """
+        Convert a Python / Pydantic type  ->  BAML TypeBuilder representation.
+        """
+
+        origin = get_origin(field_type)  # e.g. list[…]  ->  list
+        args = get_args(field_type)  # e.g. list[int] -> (int,)
+
+        # ------------------------------------------------------------------
+        # 1. Optional / Union  ------------------------------------------------
+        # ------------------------------------------------------------------
+        if origin is Union:
+            non_none_args = [t for t in args if t is not type(None)]
+
+            # Optional[T]  ⇢  exactly (T, NoneType)
+            if len(args) == 2 and len(non_none_args) == 1:
+                return map_type(non_none_args[0], field_info).optional()
+
+            # Plain Union[A, B, …]
+            return tb.union(*(map_type(t, field_info) for t in args))
+
+        # ------------------------------------------------------------------
+        # 2. List / Sequence  -------------------------------------------------
+        # ------------------------------------------------------------------
+        if origin in (list,):
+            (inner_type,) = args  # list has exactly one parameter
+            return map_type(inner_type, field_info).list()
+
+        # ------------------------------------------------------------------
+        # 3. Dict / Map -------------------------------------------------------
+        # ------------------------------------------------------------------
+        def _is_enum_subclass(key_type) -> bool:
+            """Guarded issubclass – returns False when tp is not a class."""
+            return isinstance(key_type, type) and issubclass(key_type, Enum)
+
+        if origin in (dict,):
+            key_type, value_type = args or (str, object)
+
+            if key_type is not str and not _is_enum_subclass(key_type):
+                raise ValueError("BAML maps only allow 'str' or Enum subclasses as keys")
+
+            return tb.map(
+                map_type(key_type, field_info),  # mostly tb.string() or enum
+                map_type(value_type, field_info),
+            )
+
+        # ------------------------------------------------------------------
+        # 4. Enum -------------------------------------------------------------
+        # ------------------------------------------------------------------
+        if _is_enum_subclass(field_type):
+            enum_builder = tb.add_enum(field_type.__name__)
+            for member in field_type:
+                enum_builder.add_value(member.name)
+            return enum_builder.type()
+
+        # ------------------------------------------------------------------
+        # 5. Nested Pydantic model -------------------------------------------
+        # ------------------------------------------------------------------
+        from pydantic import BaseModel  # local import
+
+        if isinstance(field_type, type) and issubclass(field_type, BaseModel):
+            try:
+                # Create nested class if it doesn't exist
+                nested_class = tb.add_class(field_type.__name__)
+                # Find dynamic types of nested class
+                create_dynamic_baml_type(tb, nested_class, field_type)
+            except ValueError:
+                # If nested class already exists get it
+                nested_class = tb._tb.class_(field_type.__name__)
+
+            # Return nested class model
+            if isinstance(nested_class, ClassBuilder):
+                # Different nested_class objects have different syntax for type information
+                # If nested class already exists type information can be found using the field method
+                return nested_class.field()
+            else:
+                # If nested class was created type information can be found using type method
+                return nested_class.type()
+
+        primitive_map = {
+            str: tb.string(),
+            int: tb.int(),
+            float: tb.float(),
+            bool: tb.bool(),
+            datetime: tb.string(),  # BAML has no native datetime
+        }
+        if field_type in primitive_map:
+            return primitive_map[field_type]
+
+        raise ValueError(f"Unsupported type for BAML mapping: {field_type}")
+
+    fields = pydantic_model.model_fields
+
+    # Add fields
+    for field_name, field_info in fields.items():
+        field_type = field_info.annotation
+        baml_type = map_type(field_type, field_info)
+
+        # Add property with type
+        prop = baml_model.add_property(field_name, baml_type)
+
+        # Add description if available
+        if field_info.description:
+            prop.description(field_info.description)
+
+    return tb

cognee/infrastructure/llm/structured_output_framework/baml/baml_src/generators.baml

@@ -6,13 +6,13 @@ generator target {
     output_type "python/pydantic"
 
     // Where the generated code will be saved (relative to baml_src/)
-    output_dir "../baml/"
+    output_dir "../"
 
     // The version of the BAML package you have installed (e.g. same version as your baml-py or @boundaryml/baml).
     // The BAML VSCode extension version should also match this version.
-    version "0.201.0"
+    version "0.206.0"
 
     // Valid values: "sync", "async"
     // This controls what `b.FunctionName()` will be (sync or async).
-    default_client_mode sync
+    default_client_mode async
 }

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py

@@ -70,35 +70,3 @@ class AnthropicAdapter(LLMInterface):
             ],
             response_model=response_model,
         )
-
-    def show_prompt(self, text_input: str, system_prompt: str) -> str:
-        """
-        Format and display the prompt for a user query.
-
-        Parameters:
-        -----------
-
-            - text_input (str): The input text from the user, defaults to a placeholder if
-              empty.
-            - system_prompt (str): The path to the system prompt to be read and formatted.
-
-        Returns:
-        --------
-
-            - str: A formatted string displaying the system prompt and user input.
-        """
-
-        if not text_input:
-            text_input = "No user input provided."
-        if not system_prompt:
-            raise MissingSystemPromptPathError()
-
-        system_prompt = LLMGateway.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