File size: 1,871 Bytes
3ec78dd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
from typing import Optional, Dict, List
from dataclasses import dataclass, field, asdict
from .Entity import Entity
@dataclass
class Node:
id: str = ''
name: str = ''
node_type: str = ''
description: Optional[str] = None
declared_entities: List[dict] = field(default_factory=list) # Classes, functions, variables
called_entities: List[str] = field( default_factory=list) # Classes, functions, variables, but also external libraries
def dict(self):
return {k: str(v) for k, v in asdict(self).items()}
@dataclass
class DirectoryNode(Node):
path: str = ''
node_type: str = 'directory'
@dataclass
class FileNode(Node):
path: str = ''
content: str = ''
node_type: str = 'file'
language : str = ''
@dataclass
class ChunkNode(FileNode):
node_type: str = 'chunk'
order_in_file: int = field(default_factory=int)
embedding : list = None
def get_field_to_embed(self) -> Optional[str]:
# Use description if available, otherwise fall back to content
# This ensures we always have something meaningful to embed
if self.description and self.description.strip():
return self.description
return self.content
@dataclass
class EntityNode(Node):
entity_type: str = ''
declaring_chunk_ids: List[str] = field(default_factory=list)
calling_chunk_ids: List[str] = field(default_factory=list)
aliases: List[str] = field(default_factory=list) # All possible aliases for this entity
node_type: str = 'entity'
def __post_init__(self):
# Use entity_name (stored in name field) as the id if id is not set
if not self.id and self.name:
self.id = self.name
def dict(self):
return {k: str(v) for k, v in asdict(self).items()}
def get_field_to_embed(self) -> Optional[str]:
return self.name |