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 

13from ...utilities.llm_utils import invoke_llm_sync 

14 

15 

16class ConstraintAnalyzer: 

17 """Analyzes queries to extract constraints.""" 

18 

19 def __init__(self, model: BaseChatModel): 

20 """Initialize the constraint analyzer.""" 

21 self.model = model 

22 

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

24 """Extract constraints from a query.""" 

25 prompt = f""" 

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

27 

28Question: {query} 

29 

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

31 

32Examples: 

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

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

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

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

37 

38For each constraint, identify: 

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

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

413. Value: The specific relationship to verify 

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

43 

44Format your response as: 

45CONSTRAINT_1: 

46Type: [type] 

47Description: [description] 

48Value: [value] 

49Weight: [0.0-1.0] 

50 

51CONSTRAINT_2: 

52Type: [type] 

53Description: [description] 

54Value: [value] 

55Weight: [0.0-1.0] 

56 

57Focus on answer verification, not query parsing. 

58""" 

59 

60 content = get_llm_response_text(invoke_llm_sync(self.model, prompt)) 

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