gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fixes

Browse source
This commit is contained in:
vasilije 2025-07-06 13:08:59 +02:00
parent a7ce271e74
commit 058a73c073
3 changed files with 9 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

0
cognee/infrastructure/llm/structured_output_framework/baml_src/__init__.py Normal file
View file

0
cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/knowledge_graph/__init__.py Normal file
View file

9
cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/knowledge_graph/extract_content_graph.py
View file

@ -5,14 +5,23 @@ from cognee.infrastructure.llm.structured_output_framework.baml.async_client imp
from cognee.infrastructure.llm.structured_output_framework.baml.type_builder import TypeBuilder from cognee.infrastructure.llm.structured_output_framework.baml.type_builder import TypeBuilder
from cognee.infrastructure.llm.structured_output_framework.baml_src.config import get_llm_config from cognee.infrastructure.llm.structured_output_framework.baml_src.config import get_llm_config
from cognee.shared.logging_utils import get_logger, setup_logging
async def extract_content_graph(content: str, response_model: Type[BaseModel]): async def extract_content_graph(content: str, response_model: Type[BaseModel]):
# tb = TypeBuilder() # tb = TypeBuilder()
config = get_llm_config() config = get_llm_config()
setup_logging()
get_logger(level="INFO")
# country = tb.union \ # country = tb.union \
# ([tb.literal_string("USA"), tb.literal_string("UK"), tb.literal_string("Germany"), tb.literal_string("other")]) # ([tb.literal_string("USA"), tb.literal_string("UK"), tb.literal_string("Germany"), tb.literal_string("other")])
# tb.Node.add_property("country", country) # tb.Node.add_property("country", country)
graph = await b.ExtractContentGraph(content, mode="simple", baml_options={ "tb": config.baml_registry}) graph = await b.ExtractContentGraph(content, mode="simple", baml_options={ "tb": config.baml_registry})
return graph return graph