37 lines
983 B
Python
37 lines
983 B
Python
|
from typing import List, Optional, Dict, Any
|
||
|
|
from pydantic import BaseModel, Field
|
||
|
|
from enum import Enum
|
||
|
|
|
||
|
|
class NodeType(str, Enum):
|
||
|
|
FUNCTION = "Function"
|
||
|
|
CLASS = "Class"
|
||
|
|
FILE = "File"
|
||
|
|
VARIABLE = "Variable"
|
||
|
|
NAMESPACE = "Namespace"
|
||
|
|
|
||
|
|
class EdgeType(str, Enum):
|
||
|
|
CALLS = "CALLS"
|
||
|
|
CONTAINS = "CONTAINS"
|
||
|
|
EXTENDS = "EXTENDS"
|
||
|
|
IMPLEMENTS = "IMPLEMENTS"
|
||
|
|
HAS_PARAMETER = "HAS_PARAMETER"
|
||
|
|
RETURNS = "RETURNS"
|
||
|
|
|
||
|
|
class GraphNode(BaseModel):
|
||
|
|
id: str
|
||
|
|
type: NodeType
|
||
|
|
name: str
|
||
|
|
signature: Optional[str] = None
|
||
|
|
file_path: Optional[str] = None
|
||
|
|
start_line: Optional[int] = None
|
||
|
|
end_line: Optional[int] = None
|
||
|
|
summary: Optional[str] = None
|
||
|
|
logic_flow: Optional[str] = None
|
||
|
|
embedding: Optional[List[float]] = None
|
||
|
|
raw_attributes: Dict[str, Any] = Field(default_factory=dict)
|
||
|
|
|
||
|
|
class GraphEdge(BaseModel):
|
||
|
|
source_id: str
|
||
|
|
target_id: str
|
||
|
|
type: EdgeType
|
||
|
|
properties: Dict[str, Any] = Field(default_factory=dict)
|