Source code for graphragzen.entity_extraction.llm_output_structures

from typing import List

from pydantic import BaseModel, Field



[docs]
class ExtractedNode(BaseModel):
    type: str = "node"
    name: str
    category: str
    description: str = Field(
        description="Short explanation as to why you think the source node and the target node are related to each other"  # noqa: E501
    )




[docs]
class ExtractedEdge(BaseModel):
    type: str = "edge"
    source: str
    target: str
    description: str = Field(
        description="Short explanation as to why you think the source node and the target node are related to each other"  # noqa: E501
    )
    weight: float




[docs]
class ExtractedEntities(BaseModel):
    extracted_nodes: List[ExtractedNode]
    extracted_edges: List[ExtractedEdge]