Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/knowledge_graph/types.py: 76%

1from __future__ import annotations 

2 

3from dataclasses import dataclass, field 

4from enum import StrEnum 

5from typing import Any 

6 

7 

8class EntityType(StrEnum): 

9 PERSON = "PERSON" 

10 ORGANIZATION = "ORGANIZATION" 

11 LOCATION = "LOCATION" 

12 EVENT = "EVENT" 

13 PRODUCT = "PRODUCT" 

14 CONCEPT = "CONCEPT" 

15 DATE = "DATE" 

16 OTHER = "OTHER" 

17 

18 

19class RelationshipType(StrEnum): 

20 WORKS_AT = "WORKS_AT" 

21 LOCATED_IN = "LOCATED_IN" 

22 PART_OF = "PART_OF" 

23 RELATED_TO = "RELATED_TO" 

24 CREATED_BY = "CREATED_BY" 

25 OCCURRED_AT = "OCCURRED_AT" 

26 IS_A = "IS_A" 

27 HAS_PROPERTY = "HAS_PROPERTY" 

28 OTHER = "OTHER" 

29 

30 

31@dataclass 

32class Entity: 

33 name: str 

34 type: EntityType | str = EntityType.OTHER 

35 properties: dict[str, Any] = field(default_factory=dict) 

36 metadata: dict[str, Any] = field(default_factory=dict) 

37 

38 def __hash__(self) -> int: 

39 return hash((self.name.lower(), self.type)) 

40 

41 def __eq__(self, other: object) -> bool: 

42 if not isinstance(other, Entity): 

43 return False 

44 return self.name.lower() == other.name.lower() and self.type == other.type 

45 

46 

47@dataclass 

48class Relationship: 

49 source: str 

50 target: str 

51 type: RelationshipType | str = RelationshipType.OTHER 

52 properties: dict[str, Any] = field(default_factory=dict) 

53 confidence: float = 1.0 

54 metadata: dict[str, Any] = field(default_factory=dict) 

55 

56 def __hash__(self) -> int: 

57 return hash((self.source.lower(), self.target.lower(), self.type)) 

58 

59 def __eq__(self, other: object) -> bool: 

60 if not isinstance(other, Relationship): 

61 return False 

62 return ( 

63 self.source.lower() == other.source.lower() 

64 and self.target.lower() == other.target.lower() 

65 and self.type == other.type 

66 ) 

67 

68 

69@dataclass 

70class GraphPath: 

71 entities: list[str] 

72 relationships: list[Relationship] 

73 length: int 

74 score: float = 1.0 

75 

76 def __repr__(self) -> str: 

77 if not self.entities: 

78 return "Empty path" 

79 parts = [self.entities[0]] 

80 for i, rel in enumerate(self.relationships): 

81 parts.append(f"--{rel.type}-->") 

82 parts.append(self.entities[i + 1]) 

83 return " ".join(parts)