From 236776a69fada8ad70cc4609640581786f16b148 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Mon, 8 Sep 2025 16:36:28 +0200 Subject: [PATCH] refactor: Update baml client --- .../baml/baml_client/__init__.py | 2 +- .../baml/baml_client/async_client.py | 14 ++-- .../baml/baml_client/inlinedbaml.py | 2 +- .../baml/baml_client/parser.py | 8 +- .../baml/baml_client/stream_types.py | 9 ++- .../baml/baml_client/sync_client.py | 14 ++-- .../baml/baml_client/type_builder.py | 80 ++++++++++++++++--- .../baml/baml_client/type_map.py | 6 +- .../baml/baml_client/types.py | 9 ++- .../baml_src/acreate_structured_output.baml | 9 ++- poetry.lock | 20 ++--- pyproject.toml | 2 +- uv.lock | 20 ++--- 13 files changed, 137 insertions(+), 58 deletions(-) diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py index 1f3cf9932..daad5fc69 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py +++ b/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.206.1" +__version__ = "0.206.0" try: from baml_py.safe_import import EnsureBamlPyImport diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/async_client.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/async_client.py index ac35cd1b1..e69544c48 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/async_client.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/async_client.py @@ -85,7 +85,7 @@ class BamlAsyncClient: system_prompt: str, user_prompt: str, baml_options: BamlCallOptions = {}, - ) -> types.DynamicModel: + ) -> types.DynamicOutputModel: # Check if on_tick is provided if "on_tick" in baml_options: # Use streaming internally when on_tick is provided @@ -107,7 +107,8 @@ class BamlAsyncClient: }, ) return typing.cast( - types.DynamicModel, result.cast_to(types, types, stream_types, False, __runtime__) + types.DynamicOutputModel, + result.cast_to(types, types, stream_types, False, __runtime__), ) async def ExtractCategories( @@ -270,7 +271,7 @@ class BamlStreamClient: system_prompt: str, user_prompt: str, baml_options: BamlCallOptions = {}, - ) -> baml_py.BamlStream[stream_types.DynamicModel, types.DynamicModel]: + ) -> baml_py.BamlStream[stream_types.DynamicOutputModel, types.DynamicOutputModel]: ctx, result = self.__options.merge_options(baml_options).create_async_stream( function_name="AcreateStructuredOutput", args={ @@ -279,13 +280,14 @@ class BamlStreamClient: "user_prompt": user_prompt, }, ) - return baml_py.BamlStream[stream_types.DynamicModel, types.DynamicModel]( + return baml_py.BamlStream[stream_types.DynamicOutputModel, types.DynamicOutputModel]( result, lambda x: typing.cast( - stream_types.DynamicModel, x.cast_to(types, types, stream_types, True, __runtime__) + stream_types.DynamicOutputModel, + x.cast_to(types, types, stream_types, True, __runtime__), ), lambda x: typing.cast( - types.DynamicModel, x.cast_to(types, types, stream_types, False, __runtime__) + types.DynamicOutputModel, x.cast_to(types, types, stream_types, False, __runtime__) ), ctx, ) diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/inlinedbaml.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/inlinedbaml.py index 12b61094f..3de9a2753 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/inlinedbaml.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/inlinedbaml.py @@ -11,7 +11,7 @@ # baml-cli is available with the baml package. _file_map = { - "acreate_structured_output.baml": "class DynamicModel {\n test string\n @@dynamic\n}\n\nfunction AcreateStructuredOutput(\n content: string,\n system_prompt: string,\n user_prompt: string,\n) -> DynamicModel {\n client OpenAI\n\n prompt #\"\n {{ system_prompt }}\n {{ ctx.output_format }}\n {{ _.role('user') }}\n {{ user_prompt }}\n {{ content }}\n \"#\n}\n", + "acreate_structured_output.baml": "class DynamicInputModel {\n test string\n @@dynamic\n}\n\nclass DynamicOutputModel {\n test string\n @@dynamic\n}\n\nfunction AcreateStructuredOutput(\n content: string,\n system_prompt: string,\n user_prompt: string,\n) -> DynamicOutputModel {\n client OpenAI\n\n prompt #\"\n {{ system_prompt }}\n {{ ctx.output_format }}\n {{ _.role('user') }}\n {{ user_prompt }}\n {{ content }}\n \"#\n}\n", "extract_categories.baml": '// Content classification data models - matching shared/data_models.py\nclass TextContent {\n type string\n subclass string[]\n}\n\nclass AudioContent {\n type string\n subclass string[]\n}\n\nclass ImageContent {\n type string\n subclass string[]\n}\n\nclass VideoContent {\n type string\n subclass string[]\n}\n\nclass MultimediaContent {\n type string\n subclass string[]\n}\n\nclass Model3DContent {\n type string\n subclass string[]\n}\n\nclass ProceduralContent {\n type string\n subclass string[]\n}\n\nclass ContentLabel {\n content_type "text" | "audio" | "image" | "video" | "multimedia" | "3d_model" | "procedural"\n type string\n subclass string[]\n}\n\nclass DefaultContentPrediction {\n label ContentLabel\n}\n\n// Content classification prompt template\ntemplate_string ClassifyContentPrompt() #"\n You are a classification engine and should classify content. Make sure to use one of the existing classification options and not invent your own.\n\n Classify the content into one of these main categories and their relevant subclasses:\n\n **TEXT CONTENT** (content_type: "text"):\n - type: "TEXTUAL_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"\n - subclass options: ["Articles, essays, and reports", "Books and manuscripts", "News stories and blog posts", "Research papers and academic publications", "Social media posts and comments", "Website content and product descriptions", "Personal narratives and stories", "Spreadsheets and tables", "Forms and surveys", "Databases and CSV files", "Source code in various programming languages", "Shell commands and scripts", "Markup languages (HTML, XML)", "Stylesheets (CSS) and configuration files (YAML, JSON, INI)", "Chat transcripts and messaging history", "Customer service logs and interactions", "Conversational AI training data", "Textbook content and lecture notes", "Exam questions and academic exercises", "E-learning course materials", "Poetry and prose", "Scripts for plays, movies, and television", "Song lyrics", "Manuals and user guides", "Technical specifications and API documentation", "Helpdesk articles and FAQs", "Contracts and agreements", "Laws, regulations, and legal case documents", "Policy documents and compliance materials", "Clinical trial reports", "Patient records and case notes", "Scientific journal articles", "Financial reports and statements", "Business plans and proposals", "Market research and analysis reports", "Ad copies and marketing slogans", "Product catalogs and brochures", "Press releases and promotional content", "Professional and formal correspondence", "Personal emails and letters", "Image and video captions", "Annotations and metadata for various media", "Vocabulary lists and grammar rules", "Language exercises and quizzes", "Other types of text data"]\n\n **AUDIO CONTENT** (content_type: "audio"):\n - type: "AUDIO_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"\n - subclass options: ["Music tracks and albums", "Podcasts and radio broadcasts", "Audiobooks and audio guides", "Recorded interviews and speeches", "Sound effects and ambient sounds", "Other types of audio recordings"]\n\n **IMAGE CONTENT** (content_type: "image"):\n - type: "IMAGE_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"\n - subclass options: ["Photographs and digital images", "Illustrations, diagrams, and charts", "Infographics and visual data representations", "Artwork and paintings", "Screenshots and graphical user interfaces", "Other types of images"]\n\n **VIDEO CONTENT** (content_type: "video"):\n - type: "VIDEO_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"\n - subclass options: ["Movies and short films", "Documentaries and educational videos", "Video tutorials and how-to guides", "Animated features and cartoons", "Live event recordings and sports broadcasts", "Other types of video content"]\n\n **MULTIMEDIA CONTENT** (content_type: "multimedia"):\n - type: "MULTIMEDIA_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"\n - subclass options: ["Interactive web content and games", "Virtual reality (VR) and augmented reality (AR) experiences", "Mixed media presentations and slide decks", "E-learning modules with integrated multimedia", "Digital exhibitions and virtual tours", "Other types of multimedia content"]\n\n **3D MODEL CONTENT** (content_type: "3d_model"):\n - type: "3D_MODEL_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"\n - subclass options: ["Architectural renderings and building plans", "Product design models and prototypes", "3D animations and character models", "Scientific simulations and visualizations", "Virtual objects for AR/VR applications", "Other types of 3D models"]\n\n **PROCEDURAL CONTENT** (content_type: "procedural"):\n - type: "PROCEDURAL_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"\n - subclass options: ["Tutorials and step-by-step guides", "Workflow and process descriptions", "Simulation and training exercises", "Recipes and crafting instructions", "Other types of procedural content"]\n\n Select the most appropriate content_type, type, and relevant subclasses.\n"#\n\n// OpenAI client defined once for all BAML files\n\n// Classification function\nfunction ExtractCategories(content: string) -> DefaultContentPrediction {\n client OpenAI\n\n prompt #"\n {{ ClassifyContentPrompt() }}\n\n {{ ctx.output_format(prefix="Answer in this schema:\\n") }}\n\n {{ _.role(\'user\') }}\n {{ content }}\n "#\n}\n\n// Test case for classification\ntest ExtractCategoriesExample {\n functions [ExtractCategories]\n args {\n content #"\n Natural language processing (NLP) is an interdisciplinary subfield of computer science and information retrieval.\n It deals with the interaction between computers and human language, in particular how to program computers to process and analyze large amounts of natural language data.\n "#\n }\n}\n', "extract_content_graph.baml": 'class Node {\n id string\n name string\n type string\n description string\n @@dynamic\n}\n\n/// doc string for edge\nclass Edge {\n /// doc string for source_node_id\n source_node_id string\n target_node_id string\n relationship_name string\n}\n\nclass KnowledgeGraph {\n nodes (Node @stream.done)[]\n edges Edge[]\n}\n\n// Summarization classes\nclass SummarizedContent {\n summary string\n description string\n}\n\nclass SummarizedFunction {\n name string\n description string\n inputs string[]?\n outputs string[]?\n decorators string[]?\n}\n\nclass SummarizedClass {\n name string\n description string\n methods SummarizedFunction[]?\n decorators string[]?\n}\n\nclass SummarizedCode {\n high_level_summary string\n key_features string[]\n imports string[]\n constants string[]\n classes SummarizedClass[]\n functions SummarizedFunction[]\n workflow_description string?\n}\n\nclass DynamicKnowledgeGraph {\n @@dynamic\n}\n\n\n// Simple template for basic extraction (fast, good quality)\ntemplate_string ExtractContentGraphPrompt() #"\n You are an advanced algorithm that extracts structured data into a knowledge graph.\n\n - **Nodes**: Entities/concepts (like Wikipedia articles).\n - **Edges**: Relationships (like Wikipedia links). Use snake_case (e.g., `acted_in`).\n\n **Rules:**\n\n 1. **Node Labeling & IDs**\n - Use basic types only (e.g., "Person", "Date", "Organization").\n - Avoid overly specific or generic terms (e.g., no "Mathematician" or "Entity").\n - Node IDs must be human-readable names from the text (no numbers).\n\n 2. **Dates & Numbers**\n - Label dates as **"Date"** in "YYYY-MM-DD" format (use available parts if incomplete).\n - Properties are key-value pairs; do not use escaped quotes.\n\n 3. **Coreference Resolution**\n - Use a single, complete identifier for each entity (e.g., always "John Doe" not "Joe" or "he").\n\n 4. **Relationship Labels**:\n - Use descriptive, lowercase, snake_case names for edges.\n - *Example*: born_in, married_to, invented_by.\n - Avoid vague or generic labels like isA, relatesTo, has.\n - Avoid duplicated relationships like produces, produced by.\n\n 5. **Strict Compliance**\n - Follow these rules exactly. Non-compliance results in termination.\n"#\n\n// Summarization prompt template\ntemplate_string SummarizeContentPrompt() #"\n You are a top-tier summarization engine. Your task is to summarize text and make it versatile.\n Be brief and concise, but keep the important information and the subject.\n Use synonym words where possible in order to change the wording but keep the meaning.\n"#\n\n// Code summarization prompt template\ntemplate_string SummarizeCodePrompt() #"\n You are an expert code analyst. Analyze the provided source code and extract key information:\n\n 1. Provide a high-level summary of what the code does\n 2. List key features and functionality\n 3. Identify imports and dependencies\n 4. List constants and global variables\n 5. Summarize classes with their methods\n 6. Summarize standalone functions\n 7. Describe the overall workflow if applicable\n\n Be precise and technical while remaining clear and concise.\n"#\n\n// Detailed template for complex extraction (slower, higher quality)\ntemplate_string DetailedExtractContentGraphPrompt() #"\n You are a top-tier algorithm designed for extracting information in structured formats to build a knowledge graph.\n **Nodes** represent entities and concepts. They\'re akin to Wikipedia nodes.\n **Edges** represent relationships between concepts. They\'re akin to Wikipedia links.\n\n The aim is to achieve simplicity and clarity in the knowledge graph.\n\n # 1. Labeling Nodes\n **Consistency**: Ensure you use basic or elementary types for node labels.\n - For example, when you identify an entity representing a person, always label it as **"Person"**.\n - Avoid using more specific terms like "Mathematician" or "Scientist", keep those as "profession" property.\n - Don\'t use too generic terms like "Entity".\n **Node IDs**: Never utilize integers as node IDs.\n - Node IDs should be names or human-readable identifiers found in the text.\n\n # 2. Handling Numerical Data and Dates\n - For example, when you identify an entity representing a date, make sure it has type **"Date"**.\n - Extract the date in the format "YYYY-MM-DD"\n - If not possible to extract the whole date, extract month or year, or both if available.\n - **Property Format**: Properties must be in a key-value format.\n - **Quotation Marks**: Never use escaped single or double quotes within property values.\n - **Naming Convention**: Use snake_case for relationship names, e.g., `acted_in`.\n\n # 3. Coreference Resolution\n - **Maintain Entity Consistency**: When extracting entities, it\'s vital to ensure consistency.\n If an entity, such as "John Doe", is mentioned multiple times in the text but is referred to by different names or pronouns (e.g., "Joe", "he"),\n always use the most complete identifier for that entity throughout the knowledge graph. In this example, use "John Doe" as the Person\'s ID.\n Remember, the knowledge graph should be coherent and easily understandable, so maintaining consistency in entity references is crucial.\n\n # 4. Strict Compliance\n Adhere to the rules strictly. Non-compliance will result in termination.\n"#\n\n// Guided template with step-by-step instructions\ntemplate_string GuidedExtractContentGraphPrompt() #"\n You are an advanced algorithm designed to extract structured information to build a clean, consistent, and human-readable knowledge graph.\n\n **Objective**:\n - Nodes represent entities and concepts, similar to Wikipedia articles.\n - Edges represent typed relationships between nodes, similar to Wikipedia hyperlinks.\n - The graph must be clear, minimal, consistent, and semantically precise.\n\n **Node Guidelines**:\n\n 1. **Label Consistency**:\n - Use consistent, basic types for all node labels.\n - Do not switch between granular or vague labels for the same kind of entity.\n - Pick one label for each category and apply it uniformly.\n - Each entity type should be in a singular form and in a case of multiple words separated by whitespaces\n\n 2. **Node Identifiers**:\n - Node IDs must be human-readable and derived directly from the text.\n - Prefer full names and canonical terms.\n - Never use integers or autogenerated IDs.\n - *Example*: Use "Marie Curie", "Theory of Evolution", "Google".\n\n 3. **Coreference Resolution**:\n - Maintain one consistent node ID for each real-world entity.\n - Resolve aliases, acronyms, and pronouns to the most complete form.\n - *Example*: Always use "John Doe" even if later referred to as "Doe" or "he".\n\n **Edge Guidelines**:\n\n 4. **Relationship Labels**:\n - Use descriptive, lowercase, snake_case names for edges.\n - *Example*: born_in, married_to, invented_by.\n - Avoid vague or generic labels like isA, relatesTo, has.\n\n 5. **Relationship Direction**:\n - Edges must be directional and logically consistent.\n - *Example*:\n - "Marie Curie" —[born_in]→ "Warsaw"\n - "Radioactivity" —[discovered_by]→ "Marie Curie"\n\n **Compliance**:\n Strict adherence to these guidelines is required. Any deviation will result in immediate termination of the task.\n"#\n\n// Strict template with zero-tolerance rules\ntemplate_string StrictExtractContentGraphPrompt() #"\n You are a top-tier algorithm for **extracting structured information** from unstructured text to build a **knowledge graph**.\n\n Your primary goal is to extract:\n - **Nodes**: Representing **entities** and **concepts** (like Wikipedia nodes).\n - **Edges**: Representing **relationships** between those concepts (like Wikipedia links).\n\n The resulting knowledge graph must be **simple, consistent, and human-readable**.\n\n ## 1. Node Labeling and Identification\n\n ### Node Types\n Use **basic atomic types** for node labels. Always prefer general types over specific roles or professions:\n - "Person" for any human.\n - "Organization" for companies, institutions, etc.\n - "Location" for geographic or place entities.\n - "Date" for any temporal expression.\n - "Event" for historical or scheduled occurrences.\n - "Work" for books, films, artworks, or research papers.\n - "Concept" for abstract notions or ideas.\n\n ### Node IDs\n - Always assign **human-readable and unambiguous identifiers**.\n - Never use numeric or autogenerated IDs.\n - Prioritize **most complete form** of entity names for consistency.\n\n ## 2. Relationship Handling\n - Use **snake_case** for all relationship (edge) types.\n - Keep relationship types semantically clear and consistent.\n - Avoid vague relation names like "related_to" unless no better alternative exists.\n\n ## 3. Strict Compliance\n Follow all rules exactly. Any deviation may lead to rejection or incorrect graph construction.\n"#\n\n// OpenAI client with environment model selection\nclient OpenAI {\n provider openai\n options {\n model client_registry.model\n api_key client_registry.api_key\n }\n}\n\n\n\n// Function that returns raw structured output (for custom objects - to be handled in Python)\nfunction ExtractContentGraphGeneric(\n content: string,\n mode: "simple" | "base" | "guided" | "strict" | "custom"?,\n custom_prompt_content: string?\n) -> KnowledgeGraph {\n client OpenAI\n\n prompt #"\n {% if mode == "base" %}\n {{ DetailedExtractContentGraphPrompt() }}\n {% elif mode == "guided" %}\n {{ GuidedExtractContentGraphPrompt() }}\n {% elif mode == "strict" %}\n {{ StrictExtractContentGraphPrompt() }}\n {% elif mode == "custom" and custom_prompt_content %}\n {{ custom_prompt_content }}\n {% else %}\n {{ ExtractContentGraphPrompt() }}\n {% endif %}\n\n {{ ctx.output_format(prefix="Answer in this schema:\\n") }}\n\n Before answering, briefly describe what you\'ll extract from the text, then provide the structured output.\n\n Example format:\n I\'ll extract the main entities and their relationships from this text...\n\n { ... }\n\n {{ _.role(\'user\') }}\n {{ content }}\n "#\n}\n\n// Backward-compatible function specifically for KnowledgeGraph\nfunction ExtractDynamicContentGraph(\n content: string,\n mode: "simple" | "base" | "guided" | "strict" | "custom"?,\n custom_prompt_content: string?\n) -> DynamicKnowledgeGraph {\n client OpenAI\n\n prompt #"\n {% if mode == "base" %}\n {{ DetailedExtractContentGraphPrompt() }}\n {% elif mode == "guided" %}\n {{ GuidedExtractContentGraphPrompt() }}\n {% elif mode == "strict" %}\n {{ StrictExtractContentGraphPrompt() }}\n {% elif mode == "custom" and custom_prompt_content %}\n {{ custom_prompt_content }}\n {% else %}\n {{ ExtractContentGraphPrompt() }}\n {% endif %}\n\n {{ ctx.output_format(prefix="Answer in this schema:\\n") }}\n\n Before answering, briefly describe what you\'ll extract from the text, then provide the structured output.\n\n Example format:\n I\'ll extract the main entities and their relationships from this text...\n\n { ... }\n\n {{ _.role(\'user\') }}\n {{ content }}\n "#\n}\n\n\n// Summarization functions\nfunction SummarizeContent(content: string) -> SummarizedContent {\n client OpenAI\n\n prompt #"\n {{ SummarizeContentPrompt() }}\n\n {{ ctx.output_format(prefix="Answer in this schema:\\n") }}\n\n {{ _.role(\'user\') }}\n {{ content }}\n "#\n}\n\nfunction SummarizeCode(content: string) -> SummarizedCode {\n client OpenAI\n\n prompt #"\n {{ SummarizeCodePrompt() }}\n\n {{ ctx.output_format(prefix="Answer in this schema:\\n") }}\n\n {{ _.role(\'user\') }}\n {{ content }}\n "#\n}\n\ntest ExtractStrictExample {\n functions [ExtractContentGraphGeneric]\n args {\n content #"\n The Python programming language was created by Guido van Rossum in 1991.\n "#\n mode "strict"\n }\n}\n', "generators.baml": '// This helps use auto generate libraries you can use in the language of\n// your choice. You can have multiple generators if you use multiple languages.\n// Just ensure that the output_dir is different for each generator.\ngenerator target {\n // Valid values: "python/pydantic", "typescript", "ruby/sorbet", "rest/openapi"\n output_type "python/pydantic"\n\n // Where the generated code will be saved (relative to baml_src/)\n output_dir "../"\n\n // The version of the BAML package you have installed (e.g. same version as your baml-py or @boundaryml/baml).\n // The BAML VSCode extension version should also match this version.\n version "0.206.0"\n\n // Valid values: "sync", "async"\n // This controls what `b.FunctionName()` will be (sync or async).\n default_client_mode async\n}\n', diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/parser.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/parser.py index a5535994a..7482217fd 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/parser.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/parser.py @@ -27,11 +27,11 @@ class LlmResponseParser: self, llm_response: str, baml_options: BamlCallOptions = {}, - ) -> types.DynamicModel: + ) -> types.DynamicOutputModel: result = self.__options.merge_options(baml_options).parse_response( function_name="AcreateStructuredOutput", llm_response=llm_response, mode="request" ) - return typing.cast(types.DynamicModel, result) + return typing.cast(types.DynamicOutputModel, result) def ExtractCategories( self, @@ -94,11 +94,11 @@ class LlmStreamParser: self, llm_response: str, baml_options: BamlCallOptions = {}, - ) -> stream_types.DynamicModel: + ) -> stream_types.DynamicOutputModel: result = self.__options.merge_options(baml_options).parse_response( function_name="AcreateStructuredOutput", llm_response=llm_response, mode="stream" ) - return typing.cast(stream_types.DynamicModel, result) + return typing.cast(stream_types.DynamicOutputModel, result) def ExtractCategories( self, diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/stream_types.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/stream_types.py index 8c6f12f09..7c60ebab1 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/stream_types.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/stream_types.py @@ -27,7 +27,7 @@ class StreamState(BaseModel, typing.Generic[StreamStateValueT]): # ######################################################################### -# Generated classes (18) +# Generated classes (19) # ######################################################################### @@ -46,11 +46,16 @@ class DefaultContentPrediction(BaseModel): label: typing.Optional["ContentLabel"] = None +class DynamicInputModel(BaseModel): + model_config = ConfigDict(extra="allow") + test: typing.Optional[str] = None + + class DynamicKnowledgeGraph(BaseModel): model_config = ConfigDict(extra="allow") -class DynamicModel(BaseModel): +class DynamicOutputModel(BaseModel): model_config = ConfigDict(extra="allow") test: typing.Optional[str] = None diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/sync_client.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/sync_client.py index 3af7f5741..148cf2ea8 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/sync_client.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/sync_client.py @@ -98,7 +98,7 @@ class BamlSyncClient: system_prompt: str, user_prompt: str, baml_options: BamlCallOptions = {}, - ) -> types.DynamicModel: + ) -> types.DynamicOutputModel: # Check if on_tick is provided if "on_tick" in baml_options: stream = self.stream.AcreateStructuredOutput( @@ -119,7 +119,8 @@ class BamlSyncClient: }, ) return typing.cast( - types.DynamicModel, result.cast_to(types, types, stream_types, False, __runtime__) + types.DynamicOutputModel, + result.cast_to(types, types, stream_types, False, __runtime__), ) def ExtractCategories( @@ -277,7 +278,7 @@ class BamlStreamClient: system_prompt: str, user_prompt: str, baml_options: BamlCallOptions = {}, - ) -> baml_py.BamlSyncStream[stream_types.DynamicModel, types.DynamicModel]: + ) -> baml_py.BamlSyncStream[stream_types.DynamicOutputModel, types.DynamicOutputModel]: ctx, result = self.__options.merge_options(baml_options).create_sync_stream( function_name="AcreateStructuredOutput", args={ @@ -286,13 +287,14 @@ class BamlStreamClient: "user_prompt": user_prompt, }, ) - return baml_py.BamlSyncStream[stream_types.DynamicModel, types.DynamicModel]( + return baml_py.BamlSyncStream[stream_types.DynamicOutputModel, types.DynamicOutputModel]( result, lambda x: typing.cast( - stream_types.DynamicModel, x.cast_to(types, types, stream_types, True, __runtime__) + stream_types.DynamicOutputModel, + x.cast_to(types, types, stream_types, True, __runtime__), ), lambda x: typing.cast( - types.DynamicModel, x.cast_to(types, types, stream_types, False, __runtime__) + types.DynamicOutputModel, x.cast_to(types, types, stream_types, False, __runtime__) ), ctx, ) diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_builder.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_builder.py index c4a6a10ca..d9c704a8f 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_builder.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_builder.py @@ -27,8 +27,9 @@ class TypeBuilder(type_builder.TypeBuilder): "AudioContent", "ContentLabel", "DefaultContentPrediction", + "DynamicInputModel", "DynamicKnowledgeGraph", - "DynamicModel", + "DynamicOutputModel", "Edge", "ImageContent", "KnowledgeGraph", @@ -53,7 +54,7 @@ class TypeBuilder(type_builder.TypeBuilder): # ######################################################################### # ######################################################################### - # Generated classes 18 + # Generated classes 19 # ######################################################################### @property @@ -68,13 +69,17 @@ class TypeBuilder(type_builder.TypeBuilder): def DefaultContentPrediction(self) -> "DefaultContentPredictionViewer": return DefaultContentPredictionViewer(self) + @property + def DynamicInputModel(self) -> "DynamicInputModelBuilder": + return DynamicInputModelBuilder(self) + @property def DynamicKnowledgeGraph(self) -> "DynamicKnowledgeGraphBuilder": return DynamicKnowledgeGraphBuilder(self) @property - def DynamicModel(self) -> "DynamicModelBuilder": - return DynamicModelBuilder(self) + def DynamicOutputModel(self) -> "DynamicOutputModelBuilder": + return DynamicOutputModelBuilder(self) @property def Edge(self) -> "EdgeViewer": @@ -135,7 +140,7 @@ class TypeBuilder(type_builder.TypeBuilder): # ######################################################################### -# Generated classes 18 +# Generated classes 19 # ######################################################################### @@ -274,6 +279,59 @@ class DefaultContentPredictionProperties: return type_builder.ClassPropertyViewer(self.__bldr.property("label")) +class DynamicInputModelAst: + def __init__(self, tb: type_builder.TypeBuilder): + _tb = tb._tb # type: ignore (we know how to use this private attribute) + self._bldr = _tb.class_("DynamicInputModel") + self._properties: typing.Set[str] = set( + [ + "test", + ] + ) + self._props = DynamicInputModelProperties(self._bldr, self._properties) + + def type(self) -> baml_py.FieldType: + return self._bldr.field() + + @property + def props(self) -> "DynamicInputModelProperties": + return self._props + + +class DynamicInputModelBuilder(DynamicInputModelAst): + def __init__(self, tb: type_builder.TypeBuilder): + super().__init__(tb) + + def add_property(self, name: str, type: baml_py.FieldType) -> baml_py.ClassPropertyBuilder: + if name in self._properties: + raise ValueError(f"Property {name} already exists.") + return self._bldr.property(name).type(type) + + def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]: + return self._bldr.list_properties() + + def remove_property(self, name: str) -> None: + self._bldr.remove_property(name) + + def reset(self) -> None: + self._bldr.reset() + + +class DynamicInputModelProperties: + def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): + self.__bldr = bldr + self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 + + def __getattr__(self, name: str) -> baml_py.ClassPropertyBuilder: + if name not in self.__properties: + raise AttributeError(f"Property {name} not found.") + return self.__bldr.property(name) + + @property + def test(self) -> baml_py.ClassPropertyBuilder: + return self.__bldr.property("test") + + class DynamicKnowledgeGraphAst: def __init__(self, tb: type_builder.TypeBuilder): _tb = tb._tb # type: ignore (we know how to use this private attribute) @@ -319,26 +377,26 @@ class DynamicKnowledgeGraphProperties: return self.__bldr.property(name) -class DynamicModelAst: +class DynamicOutputModelAst: def __init__(self, tb: type_builder.TypeBuilder): _tb = tb._tb # type: ignore (we know how to use this private attribute) - self._bldr = _tb.class_("DynamicModel") + self._bldr = _tb.class_("DynamicOutputModel") self._properties: typing.Set[str] = set( [ "test", ] ) - self._props = DynamicModelProperties(self._bldr, self._properties) + self._props = DynamicOutputModelProperties(self._bldr, self._properties) def type(self) -> baml_py.FieldType: return self._bldr.field() @property - def props(self) -> "DynamicModelProperties": + def props(self) -> "DynamicOutputModelProperties": return self._props -class DynamicModelBuilder(DynamicModelAst): +class DynamicOutputModelBuilder(DynamicOutputModelAst): def __init__(self, tb: type_builder.TypeBuilder): super().__init__(tb) @@ -357,7 +415,7 @@ class DynamicModelBuilder(DynamicModelAst): self._bldr.reset() -class DynamicModelProperties: +class DynamicOutputModelProperties: def __init__(self, bldr: baml_py.ClassBuilder, properties: typing.Set[str]): self.__bldr = bldr self.__properties = properties # type: ignore (we know how to use this private attribute) # noqa: F821 diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_map.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_map.py index d56e660a7..2e0427595 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_map.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_map.py @@ -21,10 +21,12 @@ type_map = { "stream_types.ContentLabel": stream_types.ContentLabel, "types.DefaultContentPrediction": types.DefaultContentPrediction, "stream_types.DefaultContentPrediction": stream_types.DefaultContentPrediction, + "types.DynamicInputModel": types.DynamicInputModel, + "stream_types.DynamicInputModel": stream_types.DynamicInputModel, "types.DynamicKnowledgeGraph": types.DynamicKnowledgeGraph, "stream_types.DynamicKnowledgeGraph": stream_types.DynamicKnowledgeGraph, - "types.DynamicModel": types.DynamicModel, - "stream_types.DynamicModel": stream_types.DynamicModel, + "types.DynamicOutputModel": types.DynamicOutputModel, + "stream_types.DynamicOutputModel": stream_types.DynamicOutputModel, "types.Edge": types.Edge, "stream_types.Edge": stream_types.Edge, "types.ImageContent": types.ImageContent, diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/types.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/types.py index f2c51ed89..fd1b10eef 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/types.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_client/types.py @@ -48,7 +48,7 @@ def all_succeeded(checks: typing.Dict[CheckName, Check]) -> bool: # ######################################################################### # ######################################################################### -# Generated classes (18) +# Generated classes (19) # ######################################################################### @@ -75,11 +75,16 @@ class DefaultContentPrediction(BaseModel): label: "ContentLabel" +class DynamicInputModel(BaseModel): + model_config = ConfigDict(extra="allow") + test: str + + class DynamicKnowledgeGraph(BaseModel): model_config = ConfigDict(extra="allow") -class DynamicModel(BaseModel): +class DynamicOutputModel(BaseModel): model_config = ConfigDict(extra="allow") test: str diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/acreate_structured_output.baml b/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/acreate_structured_output.baml index b826d9be1..36f434e18 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/acreate_structured_output.baml +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/acreate_structured_output.baml @@ -1,4 +1,9 @@ -class DynamicModel { +class DynamicInputModel { + test string + @@dynamic +} + +class DynamicOutputModel { test string @@dynamic } @@ -7,7 +12,7 @@ function AcreateStructuredOutput( content: string, system_prompt: string, user_prompt: string, -) -> DynamicModel { +) -> DynamicOutputModel { client OpenAI prompt #" diff --git a/poetry.lock b/poetry.lock index 64c1bb050..4065fd51c 100644 --- a/poetry.lock +++ b/poetry.lock @@ -587,20 +587,20 @@ extras = ["regex"] [[package]] name = "baml-py" -version = "0.201.0" +version = "0.206.0" description = "BAML python bindings (pyproject.toml)" optional = false python-versions = "*" groups = ["main"] files = [ - {file = "baml_py-0.201.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:83228d2af2b0e845bbbb4e14f7cbd3376cec385aee01210ac522ab6076e07bec"}, - {file = "baml_py-0.201.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:2a9d016139e3ae5b5ce98c7b05b5fbd53d5d38f04dc810ec4d70fb17dd6c10e4"}, - {file = "baml_py-0.201.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b5058505b1a3c5f04fc1679aec4d730fa9bef2cbd96209b3ed50152f60b96baf"}, - {file = "baml_py-0.201.0-cp38-abi3-manylinux_2_24_aarch64.whl", hash = "sha256:36289d548581ba4accd5eaaab3246872542dd32dc6717e537654fa0cad884071"}, - {file = "baml_py-0.201.0-cp38-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:5ab70e7bd6481d71edca8a33313347b29faccec78b9960138aa437522813ac9a"}, - {file = "baml_py-0.201.0-cp38-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:7efc5c693a7142c230a4f3d6700415127fee0b9f5fdbb36db63e04e27ac4c0f1"}, - {file = "baml_py-0.201.0-cp38-abi3-win_amd64.whl", hash = "sha256:56499857b7a27ae61a661c8ce0dddd0fb567a45c0b826157e44048a14cf586f9"}, - {file = "baml_py-0.201.0-cp38-abi3-win_arm64.whl", hash = "sha256:1e52dc1151db84a302b746590fe2bc484bdd794f83fa5da7216d9394c559f33a"}, + {file = "baml_py-0.206.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:d9e95f0b481a18ae6936d720b8fc609baec4ea1eabbdde48f1536ffc94ebf39f"}, + {file = "baml_py-0.206.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:0f698127da030b728c7aa2641c3164a3ab19779594019234361dd48f0784f195"}, + {file = "baml_py-0.206.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:095a4039851fd6a2f2b516446aa9c04d633a7f4ddc72aa70c6805af007a62ce0"}, + {file = "baml_py-0.206.0-cp38-abi3-manylinux_2_24_aarch64.whl", hash = "sha256:4df110271c4861850fc58f30c49d7c3ca68cde28dd4d6f613a64c80d281e3c02"}, + {file = "baml_py-0.206.0-cp38-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:ecfae45310a47e669a5b00ba7f51bece300da1fe0df39ad687fd590a7bc8ff44"}, + {file = "baml_py-0.206.0-cp38-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:a104969f7fc04932f9ab6976f7fdd173097d6e8fc64b3faa5ad4b3ad5f870c65"}, + {file = "baml_py-0.206.0-cp38-abi3-win_amd64.whl", hash = "sha256:efbff3f4b49639f25a52433a7eacda039d3118e7965124ced0705017f9847916"}, + {file = "baml_py-0.206.0-cp38-abi3-win_arm64.whl", hash = "sha256:5560becf0c8af22b415d3eb6000629727ad1eb6cde4a6bdd9584c112c42f966d"}, ] [[package]] @@ -11728,4 +11728,4 @@ posthog = ["posthog"] [metadata] lock-version = "2.1" python-versions = ">=3.10,<=3.13" -content-hash = "576318d370b89d128a7c3e755fe3c898fef4e359acdd3f05f952ae497751fb04" +content-hash = "03cd2dee06a9e89967b9287499956bcc18c6eb1849daddeb9504ff71d2453152" diff --git a/pyproject.toml b/pyproject.toml index ece238338..b120b327c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,7 +56,7 @@ dependencies = [ "dlt[sqlalchemy]>=1.9.0,<2", "sentry-sdk[fastapi]>=2.9.0,<3", "structlog>=25.2.0,<26", - "baml-py (>=0.201.0,<0.202.0)", + "baml-py (==0.206.0)", "pympler>=1.1,<2.0.0", "onnxruntime>=1.0.0,<2.0.0", "pylance>=0.22.0,<1.0.0", diff --git a/uv.lock b/uv.lock index fb8ecd9bd..eeaab0b66 100644 --- a/uv.lock +++ b/uv.lock @@ -422,17 +422,17 @@ wheels = [ [[package]] name = "baml-py" -version = "0.201.0" +version = "0.206.0" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/54/54/2b0edb3d22e95ce56f36610391c11108a4ef26ba2837736a32001687ae34/baml_py-0.201.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:83228d2af2b0e845bbbb4e14f7cbd3376cec385aee01210ac522ab6076e07bec", size = 17387971, upload-time = "2025-07-03T19:29:05.844Z" }, - { url = "https://files.pythonhosted.org/packages/c9/08/1d48c28c63eadea2c04360cbb7f64968599e99cd6b8fc0ec0bd4424d3cf1/baml_py-0.201.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:2a9d016139e3ae5b5ce98c7b05b5fbd53d5d38f04dc810ec4d70fb17dd6c10e4", size = 16191010, upload-time = "2025-07-03T19:29:09.323Z" }, - { url = "https://files.pythonhosted.org/packages/73/1a/20b2d46501e3dd0648af339825106a6ac5eeb5d22d7e6a10cf16b9aa1cb8/baml_py-0.201.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b5058505b1a3c5f04fc1679aec4d730fa9bef2cbd96209b3ed50152f60b96baf", size = 19950249, upload-time = "2025-07-03T19:29:11.974Z" }, - { url = "https://files.pythonhosted.org/packages/38/24/bc871059e905159ae1913c2e3032dd6ef2f5c3d0983999d2c2f1eebb65a4/baml_py-0.201.0-cp38-abi3-manylinux_2_24_aarch64.whl", hash = "sha256:36289d548581ba4accd5eaaab3246872542dd32dc6717e537654fa0cad884071", size = 19231310, upload-time = "2025-07-03T19:29:14.857Z" }, - { url = "https://files.pythonhosted.org/packages/0e/11/4268a0b82b02c7202fe5aa0d7175712158d998c491cac723b2bac3d5d495/baml_py-0.201.0-cp38-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:5ab70e7bd6481d71edca8a33313347b29faccec78b9960138aa437522813ac9a", size = 19490012, upload-time = "2025-07-03T19:29:18.512Z" }, - { url = "https://files.pythonhosted.org/packages/31/21/c9f9aea1adba2a5978ffab11ba0948a9f3f81ec6ed3056067713260e93a1/baml_py-0.201.0-cp38-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:7efc5c693a7142c230a4f3d6700415127fee0b9f5fdbb36db63e04e27ac4c0f1", size = 20090620, upload-time = "2025-07-03T19:29:21.072Z" }, - { url = "https://files.pythonhosted.org/packages/99/cf/92123d8d753f1d1473e080c4c182139bfe3b9a6418e891cf1d96b6c33848/baml_py-0.201.0-cp38-abi3-win_amd64.whl", hash = "sha256:56499857b7a27ae61a661c8ce0dddd0fb567a45c0b826157e44048a14cf586f9", size = 17253005, upload-time = "2025-07-03T19:29:23.722Z" }, - { url = "https://files.pythonhosted.org/packages/59/88/5056aa1bc9480f758cd6e210d63bd1f9ad90b44c87f4121285906526495e/baml_py-0.201.0-cp38-abi3-win_arm64.whl", hash = "sha256:1e52dc1151db84a302b746590fe2bc484bdd794f83fa5da7216d9394c559f33a", size = 15612701, upload-time = "2025-07-03T19:29:26.712Z" }, + { url = "https://files.pythonhosted.org/packages/16/e6/b14c0892281fe40560420674be85b50a6849370819185919f4a951ff77f0/baml_py-0.206.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:d9e95f0b481a18ae6936d720b8fc609baec4ea1eabbdde48f1536ffc94ebf39f", size = 19198123, upload-time = "2025-08-28T04:40:13.4Z" }, + { url = "https://files.pythonhosted.org/packages/72/b4/7eff7673fe47803cd369d6732e94244c7cb343ca425671227603b6cb91f0/baml_py-0.206.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:0f698127da030b728c7aa2641c3164a3ab19779594019234361dd48f0784f195", size = 17937150, upload-time = "2025-08-28T04:40:15.932Z" }, + { url = "https://files.pythonhosted.org/packages/ff/f5/fdaf5839693f60e48f7736f29b255f57b9b611edff83703873232855ac36/baml_py-0.206.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:095a4039851fd6a2f2b516446aa9c04d633a7f4ddc72aa70c6805af007a62ce0", size = 21779193, upload-time = "2025-08-28T04:40:18.141Z" }, + { url = "https://files.pythonhosted.org/packages/6e/82/53747e5cb911f33a0f44d6e23669b050dc55a4581e28831f2bd5ff376831/baml_py-0.206.0-cp38-abi3-manylinux_2_24_aarch64.whl", hash = "sha256:4df110271c4861850fc58f30c49d7c3ca68cde28dd4d6f613a64c80d281e3c02", size = 21019066, upload-time = "2025-08-28T04:40:20.568Z" }, + { url = "https://files.pythonhosted.org/packages/26/9f/a1d456af11551dd277c3c7445bfb7eec7f6d5a9f399123d72a3348ccac18/baml_py-0.206.0-cp38-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:ecfae45310a47e669a5b00ba7f51bece300da1fe0df39ad687fd590a7bc8ff44", size = 21257819, upload-time = "2025-08-28T04:40:22.708Z" }, + { url = "https://files.pythonhosted.org/packages/ae/97/f8f4f2cc298f44c01c2fc2a99f216805538c92bfc4ec81e13bd38b9e05df/baml_py-0.206.0-cp38-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:a104969f7fc04932f9ab6976f7fdd173097d6e8fc64b3faa5ad4b3ad5f870c65", size = 21938777, upload-time = "2025-08-28T04:40:25.174Z" }, + { url = "https://files.pythonhosted.org/packages/eb/6a/e9edf5c9e244fac7f33c515da4507e0f9561037670d36b92bb671cda2437/baml_py-0.206.0-cp38-abi3-win_amd64.whl", hash = "sha256:efbff3f4b49639f25a52433a7eacda039d3118e7965124ced0705017f9847916", size = 19015250, upload-time = "2025-08-28T04:40:27.212Z" }, + { url = "https://files.pythonhosted.org/packages/9a/62/7bc8227374790761e6482d77aa9b702422ba3efaccd014f57488091b1fb6/baml_py-0.206.0-cp38-abi3-win_arm64.whl", hash = "sha256:5560becf0c8af22b415d3eb6000629727ad1eb6cde4a6bdd9584c112c42f966d", size = 17289275, upload-time = "2025-08-28T04:40:29.755Z" }, ] [[package]] @@ -972,7 +972,7 @@ requires-dist = [ { name = "anthropic", marker = "extra == 'anthropic'", specifier = ">=0.26.1,<0.27" }, { name = "asyncpg", marker = "extra == 'postgres'", specifier = ">=0.30.0,<1.0.0" }, { name = "asyncpg", marker = "extra == 'postgres-binary'", specifier = ">=0.30.0,<1.0.0" }, - { name = "baml-py", specifier = ">=0.201.0,<0.202.0" }, + { name = "baml-py", specifier = "==0.206.0" }, { name = "chromadb", marker = "extra == 'chromadb'", specifier = ">=0.3.0,<0.7" }, { name = "coverage", marker = "extra == 'dev'", specifier = ">=7.3.2,<8" }, { name = "debugpy", marker = "extra == 'debug'", specifier = ">=1.8.9,<2.0.0" },