Coverage for src/local_deep_research/advanced_search_system/constraints/constraint_analyzer.py: 100%

1""" 

2Constraint analyzer for extracting constraints from queries. 

3""" 

4 

5import re 

6from typing import List 

7 

8from langchain_core.language_models import BaseChatModel 

9from loguru import logger 

10 

11from .base_constraint import Constraint, ConstraintType 

12 

13 

14class ConstraintAnalyzer: 

15 """Analyzes queries to extract constraints.""" 

16 

17 def __init__(self, model: BaseChatModel): 

18 """Initialize the constraint analyzer.""" 

19 self.model = model 

20 

21 def extract_constraints(self, query: str) -> List[Constraint]: 

22 """Extract constraints from a query.""" 

23 prompt = f""" 

24Generate constraints to verify if an answer candidate correctly answers this question. 

25 

26Question: {query} 

27 

28Create constraints that would help verify if a proposed answer is correct. Focus on the RELATIONSHIP between the question and answer, not just query analysis. 

29 

30Examples: 

31- "Which university did Alice study at?" → "Alice studied at this university" 

32- "What year was the company founded?" → "The company was founded in this year" 

33- "Who invented the device?" → "This person invented the device" 

34- "Where is the building located?" → "The building is located at this place" 

35 

36For each constraint, identify: 

371. Type: property, name_pattern, event, statistic, temporal, location, comparison, existence 

382. Description: What relationship must hold between question and answer 

393. Value: The specific relationship to verify 

404. Weight: How critical this constraint is (0.0-1.0) 

41 

42Format your response as: 

43CONSTRAINT_1: 

44Type: [type] 

45Description: [description] 

46Value: [value] 

47Weight: [0.0-1.0] 

48 

49CONSTRAINT_2: 

50Type: [type] 

51Description: [description] 

52Value: [value] 

53Weight: [0.0-1.0] 

54 

55Focus on answer verification, not query parsing. 

56""" 

57 

58 response = self.model.invoke(prompt) 

59 content = response.content 

60 

61 constraints = [] 

62 current_constraint = {} 

63 constraint_id = 1 

64 

65 for line in content.strip().split("\n"): 

66 line = line.strip() 

67 

68 if line.startswith("CONSTRAINT_"): 

69 if current_constraint and all( 

70 k in current_constraint 

71 for k in ["type", "description", "value"] 

72 ): 

73 constraint = Constraint( 

74 id=f"c{constraint_id}", 

75 type=self._parse_constraint_type( 

76 current_constraint["type"] 

77 ), 

78 description=current_constraint["description"], 

79 value=current_constraint["value"], 

80 weight=self._parse_weight( 

81 current_constraint.get("weight", 1.0) 

82 ), 

83 ) 

84 constraints.append(constraint) 

85 constraint_id += 1 

86 current_constraint = {} 

87 elif ":" in line: 

88 key, value = line.split(":", 1) 

89 key = key.strip().lower() 

90 value = value.strip() 

91 if key in ["type", "description", "value", "weight"]: 

92 current_constraint[key] = value 

93 

94 # Don't forget the last constraint 

95 if current_constraint and all( 

96 k in current_constraint for k in ["type", "description", "value"] 

97 ): 

98 constraint = Constraint( 

99 id=f"c{constraint_id}", 

100 type=self._parse_constraint_type(current_constraint["type"]), 

101 description=current_constraint["description"], 

102 value=current_constraint["value"], 

103 weight=self._parse_weight( 

104 current_constraint.get("weight", 1.0) 

105 ), 

106 ) 

107 constraints.append(constraint) 

108 

109 logger.info(f"Extracted {len(constraints)} constraints from query") 

110 return constraints 

111 

112 def _parse_constraint_type(self, type_str: str) -> ConstraintType: 

113 """Parse constraint type from string.""" 

114 type_map = { 

115 "property": ConstraintType.PROPERTY, 

116 "name_pattern": ConstraintType.NAME_PATTERN, 

117 "event": ConstraintType.EVENT, 

118 "statistic": ConstraintType.STATISTIC, 

119 "temporal": ConstraintType.TEMPORAL, 

120 "location": ConstraintType.LOCATION, 

121 "comparison": ConstraintType.COMPARISON, 

122 "existence": ConstraintType.EXISTENCE, 

123 } 

124 return type_map.get(type_str.lower(), ConstraintType.PROPERTY) 

125 

126 def _parse_weight(self, weight_value) -> float: 

127 """Parse weight value to float, handling text annotations. 

128 

129 Args: 

130 weight_value: String or numeric weight value, possibly with text annotations 

131 

132 Returns: 

133 float: Parsed weight value 

134 """ 

135 if isinstance(weight_value, (int, float)): 

136 return float(weight_value) 

137 if isinstance(weight_value, str): 

138 # Extract the first number from the string 

139 match = re.search(r"(\d+(\.\d+)?)", weight_value) 

140 if match: 

141 return float(match.group(1)) 

142 return 1.0 # Default weight