From 2413b7272b172121e679fd2e6dc47739006f1c10 Mon Sep 17 00:00:00 2001 From: vasilije Date: Wed, 27 Aug 2025 17:34:55 +0200 Subject: [PATCH] clean up poetry --- cognee/infrastructure/llm/config.py | 1 + .../baml/baml_src/extraction/extract_summary.py | 8 ++++++-- .../extraction/knowledge_graph/extract_content_graph.py | 4 +++- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/cognee/infrastructure/llm/config.py b/cognee/infrastructure/llm/config.py index 1ff8f6433..b2a058ca6 100644 --- a/cognee/infrastructure/llm/config.py +++ b/cognee/infrastructure/llm/config.py @@ -3,6 +3,7 @@ from typing import Optional, ClassVar from functools import lru_cache from pydantic_settings import BaseSettings, SettingsConfigDict from pydantic import model_validator + try: from baml_py import ClientRegistry except ImportError: diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/extract_summary.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/extract_summary.py index 8caabd937..89889d294 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/extract_summary.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/extract_summary.py @@ -38,7 +38,9 @@ async def extract_summary(content: str, response_model: Type[BaseModel]): config = get_llm_config() if config.baml_registry is None: - raise ImportError("BAML is not available. Please install with 'pip install cognee[baml]' to use BAML extraction features.") + raise ImportError( + "BAML is not available. Please install with 'pip install cognee[baml]' to use BAML extraction features." + ) # Use BAML's SummarizeContent function summary_result = await b.SummarizeContent( @@ -81,7 +83,9 @@ async def extract_code_summary(content: str): config = get_llm_config() if config.baml_registry is None: - raise ImportError("BAML is not available. Please install with 'pip install cognee[baml]' to use BAML extraction features.") + raise ImportError( + "BAML is not available. Please install with 'pip install cognee[baml]' to use BAML extraction features." + ) result = await b.SummarizeCode( content, baml_options={"client_registry": config.baml_registry} diff --git a/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/knowledge_graph/extract_content_graph.py b/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/knowledge_graph/extract_content_graph.py index aa4fce637..f87d87d1b 100644 --- a/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/knowledge_graph/extract_content_graph.py +++ b/cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/knowledge_graph/extract_content_graph.py @@ -17,7 +17,9 @@ async def extract_content_graph( get_logger(level="INFO") if config.baml_registry is None: - raise ImportError("BAML is not available. Please install with 'pip install cognee[baml]' to use BAML extraction features.") + raise ImportError( + "BAML is not available. Please install with 'pip install cognee[baml]' to use BAML extraction features." + ) # if response_model: # # tb = TypeBuilder()