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

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 ...utilities.search_utilities import remove_think_tags 

12from .base_constraint import Constraint, ConstraintType 

13 

14 

15class ConstraintAnalyzer: 

16 """Analyzes queries to extract constraints.""" 

17 

18 def __init__(self, model: BaseChatModel): 

19 """Initialize the constraint analyzer.""" 

20 self.model = model 

21 

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

23 """Extract constraints from a query.""" 

24 prompt = f""" 

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

26 

27Question: {query} 

28 

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

30 

31Examples: 

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

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

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

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

36 

37For each constraint, identify: 

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

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

403. Value: The specific relationship to verify 

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

42 

43Format your response as: 

44CONSTRAINT_1: 

45Type: [type] 

46Description: [description] 

47Value: [value] 

48Weight: [0.0-1.0] 

49 

50CONSTRAINT_2: 

51Type: [type] 

52Description: [description] 

53Value: [value] 

54Weight: [0.0-1.0] 

55 

56Focus on answer verification, not query parsing. 

57""" 

58 

59 response = self.model.invoke(prompt) 

60 content = remove_think_tags(response.content) 

61 

62 constraints = [] 

63 current_constraint = {} 

64 constraint_id = 1 

65 

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

67 line = line.strip() 

68 

69 if line.startswith("CONSTRAINT_"): 

70 if current_constraint and all( 

71 k in current_constraint 

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

73 ): 

74 constraint = Constraint( 

75 id=f"c{constraint_id}", 

76 type=self._parse_constraint_type( 

77 current_constraint["type"] 

78 ), 

79 description=current_constraint["description"], 

80 value=current_constraint["value"], 

81 weight=self._parse_weight( 

82 current_constraint.get("weight", 1.0) 

83 ), 

84 ) 

85 constraints.append(constraint) 

86 constraint_id += 1 

87 current_constraint = {} 

88 elif ":" in line: 

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

90 key = key.strip().lower() 

91 value = value.strip() 

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

93 current_constraint[key] = value 

94 

95 # Don't forget the last constraint 

96 if current_constraint and all( 

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

98 ): 

99 constraint = Constraint( 

100 id=f"c{constraint_id}", 

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

102 description=current_constraint["description"], 

103 value=current_constraint["value"], 

104 weight=self._parse_weight( 

105 current_constraint.get("weight", 1.0) 

106 ), 

107 ) 

108 constraints.append(constraint) 

109 

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

111 return constraints 

112 

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

114 """Parse constraint type from string.""" 

115 type_map = { 

116 "property": ConstraintType.PROPERTY, 

117 "name_pattern": ConstraintType.NAME_PATTERN, 

118 "event": ConstraintType.EVENT, 

119 "statistic": ConstraintType.STATISTIC, 

120 "temporal": ConstraintType.TEMPORAL, 

121 "location": ConstraintType.LOCATION, 

122 "comparison": ConstraintType.COMPARISON, 

123 "existence": ConstraintType.EXISTENCE, 

124 } 

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

126 

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

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

129 

130 Args: 

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

132 

133 Returns: 

134 float: Parsed weight value 

135 """ 

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

137 return float(weight_value) 

138 if isinstance(weight_value, str): 

139 # Extract the first number from the string 

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

141 if match: 

142 return float(match.group(1)) 

143 return 1.0 # Default weight