31 lines
707 B
Python
31 lines
707 B
Python
import asyncio
|
|
|
|
from pydantic import BaseModel
|
|
from typing import List
|
|
from cognee.infrastructure.llm.LLMGateway import LLMGateway
|
|
|
|
|
|
class MiniEntity(BaseModel):
|
|
name: str
|
|
type: str
|
|
|
|
|
|
class MiniGraph(BaseModel):
|
|
nodes: List[MiniEntity]
|
|
|
|
|
|
async def main():
|
|
system_prompt = (
|
|
"Extract entities as nodes with name and type. "
|
|
"Use concise, literal values present in the text."
|
|
)
|
|
|
|
text = "Apple develops iPhone; Audi produces the R8."
|
|
|
|
result = await LLMGateway.acreate_structured_output(text, system_prompt, MiniGraph)
|
|
print(result)
|
|
# MiniGraph(nodes=[MiniEntity(name='Apple', type='Organization'), ...])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|