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

1from __future__ import annotations 

2 

3from datetime import UTC, datetime 

4 

5from lexigram import serialization as json 

6from lexigram.ai.rag.exceptions import RAGError 

7from lexigram.ai.rag.knowledge_graph.types import ( 

8 Entity, 

9 EntityType, 

10 Relationship, 

11 RelationshipType, 

12) 

13from lexigram.contracts import ( 

14 LLMClientProtocol, 

15) 

16from lexigram.contracts.ai.llm import ChatMessage, Role 

17from lexigram.logging import ( 

18 get_logger, 

19) 

20 

21logger = get_logger(__name__) 

22 

23 

24class EntityExtractor: 

25 def __init__( 

26 self, 

27 llm_client: LLMClientProtocol, 

28 entity_types: list[EntityType | str] | None = None, 

29 min_confidence: float = 0.5, 

30 ) -> None: 

31 self.llm_client = llm_client 

32 self.entity_types = entity_types 

33 self.min_confidence = min_confidence 

34 

35 async def extract(self, text: str) -> list[Entity]: 

36 # Call the LLM to extract entities as JSON and parse the result. 

37 try: 

38 result = await self.llm_client.complete( 

39 [ChatMessage(role=Role.USER, content=text)], 

40 ) 

41 if result.is_err(): 

42 raise result.unwrap_err() 

43 completion = result.unwrap() 

44 

45 parsed = json.loads(completion.content) 

46 entities: list[Entity] = [] 

47 for e in parsed: 

48 name = e.get("name") 

49 if not isinstance(name, str): 

50 continue 

51 

52 ent = Entity( 

53 name=name, 

54 type=e.get("type", EntityType.OTHER), 

55 properties=e.get("properties", {}), 

56 metadata={"extraction_time": datetime.now(UTC).isoformat()}, 

57 ) 

58 entities.append(ent) 

59 

60 return entities 

61 except json.JSONDecodeError as e: 

62 logger.exception( 

63 "Failed to parse entity extraction JSON. Content: %s", 

64 completion.content, 

65 ) 

66 return [] 

67 except Exception as e: 

68 logger.exception("Unexpected error during entity extraction") 

69 raise RAGError(f"Entity extraction failed: {e}") from e 

70 

71 

72class RelationshipExtractor: 

73 def __init__( 

74 self, 

75 llm_client: LLMClientProtocol, 

76 relationship_types: list[RelationshipType | str] | None = None, 

77 min_confidence: float = 0.5, 

78 ) -> None: 

79 self.llm_client = llm_client 

80 self.relationship_types = relationship_types 

81 self.min_confidence = min_confidence 

82 

83 async def extract( 

84 self, 

85 text: str, 

86 entities: list[Entity] | None = None, 

87 ) -> list[Relationship]: 

88 try: 

89 result = await self.llm_client.complete( 

90 [ChatMessage(role=Role.USER, content=text)], 

91 ) 

92 if result.is_err(): 

93 raise result.unwrap_err() 

94 completion = result.unwrap() 

95 

96 rels_data = json.loads(completion.content) 

97 relationships: list[Relationship] = [] 

98 for data in rels_data: 

99 if isinstance(data, dict) and "source" in data and "target" in data: 

100 rel = Relationship( 

101 source=data["source"], 

102 target=data["target"], 

103 type=data.get("type", RelationshipType.OTHER), 

104 confidence=data.get("confidence", 1.0), 

105 metadata={"extraction_time": datetime.now(UTC).isoformat()}, 

106 ) 

107 if rel.confidence >= self.min_confidence: 

108 relationships.append(rel) 

109 return relationships 

110 except json.JSONDecodeError: 

111 logger.exception( 

112 "Failed to parse relationship extraction JSON. Content: %s", 

113 completion.content, 

114 ) 

115 return [] 

116 except Exception as e: 

117 logger.exception("Unexpected error during relationship extraction") 

118 raise RAGError(f"Relationship extraction failed: {e}") from e