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 ...utilities.json_utils import get_llm_response_text 

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 content = get_llm_response_text(self.model.invoke(prompt)) 

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