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

refactor: remove literal support

Browse source
This commit is contained in:
Igor Ilic 2025-09-09 17:27:14 +02:00
parent 11aefadb81
commit 4258f2e92b
1 changed files with 14 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/acreate_structured_output.py
View file

@ -55,9 +55,9 @@ def create_dynamic_baml_type(tb, baml_model, pydantic_model):
# ------------------------------------------------------------------ # ------------------------------------------------------------------
# 3. Dict / Map ------------------------------------------------------- # 3. Dict / Map -------------------------------------------------------
# ------------------------------------------------------------------ # ------------------------------------------------------------------
def _is_enum_subclass(tp) -> bool: def _is_enum_subclass(key_type) -> bool:
"""Guarded issubclass – returns False when tp is not a class.""" """Guarded issubclass – returns False when tp is not a class."""
return isinstance(tp, type) and issubclass(tp, Enum) return isinstance(key_type, type) and issubclass(key_type, Enum)
if origin in (dict,): if origin in (dict,):
key_type, value_type = args or (str, object) key_type, value_type = args or (str, object)
@ -71,13 +71,7 @@ def create_dynamic_baml_type(tb, baml_model, pydantic_model):
) )
# ------------------------------------------------------------------ # ------------------------------------------------------------------
# 4. Literal ---------------------------------------------------------- # 4. Enum -------------------------------------------------------------
# ------------------------------------------------------------------
if origin is Literal:
return tb.union(*(tb.literal(v) for v in args))
# ------------------------------------------------------------------
# 5. Enum -------------------------------------------------------------
# ------------------------------------------------------------------ # ------------------------------------------------------------------
if _is_enum_subclass(field_type): if _is_enum_subclass(field_type):
enum_builder = tb.add_enum(field_type.__name__) enum_builder = tb.add_enum(field_type.__name__)
@ -86,7 +80,7 @@ def create_dynamic_baml_type(tb, baml_model, pydantic_model):
return enum_builder.type() return enum_builder.type()
# ------------------------------------------------------------------ # ------------------------------------------------------------------
# 6. Nested Pydantic model ------------------------------------------- # 5. Nested Pydantic model -------------------------------------------
# ------------------------------------------------------------------ # ------------------------------------------------------------------
from pydantic import BaseModel # local import from pydantic import BaseModel # local import
@ -173,6 +167,15 @@ if __name__ == "__main__":
loop = asyncio.new_event_loop() loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop) asyncio.set_event_loop(loop)
try: try:
loop.run_until_complete(acreate_structured_output("TEST", "THIS IS A TEST", str)) 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: finally:
loop.run_until_complete(loop.shutdown_asyncgens()) loop.run_until_complete(loop.shutdown_asyncgens())