@@ -180,25 +180,36 @@ async def _get_question_answer(row: GetPromptWithOutputsRow) -> Optional[Partial
180180
181181
182182def parse_question_answer (input_text : str ) -> str :
183- # given a string, detect if we have a pattern of "Context: xxx \n\nQuery: xxx" and strip it
184- pattern = r"^Context:.*?\n\n\s*Query:\s*(.*)$"
183+ # Remove the <environment_details>...</environment_details> pattern if present
184+ env_details_pattern = r"\n<environment_details>.*?</environment_details>"
185+ input_text = re .sub (env_details_pattern , "" , input_text , flags = re .DOTALL ).strip ()
185186
186- # Search using the regex pattern
187- match = re .search (pattern , input_text , re .DOTALL )
187+ # Check for the <task>...</task> pattern first
188+ task_pattern = r"^<task>(.*?)</task>"
189+ task_match = re .search (task_pattern , input_text , re .DOTALL )
188190
189- # If a match is found, return the captured group after "Query:"
190- if match :
191- return match .group (1 )
192- else :
193- return input_text
191+ if task_match :
192+ return task_match .group (1 ).strip ()
193+
194+ # If no <task>...</task>, check for "Context: xxx \n\nQuery: xxx"
195+ context_query_pattern = r"^Context:.*?\n\n\s*Query:\s*(.*)$"
196+ context_query_match = re .search (context_query_pattern , input_text , re .DOTALL )
197+
198+ if context_query_match :
199+ return context_query_match .group (1 ).strip ()
200+
201+ # If no pattern matches, return the original input text
202+ return input_text
194203
195204
196205def _clean_secrets_from_message (message : str ) -> str :
197206 pattern = re .compile (r"REDACTED<(\$?[^>]+)>" )
198207 return pattern .sub ("REDACTED_SECRET" , message )
199208
200209
201- def _group_partial_messages (pq_list : List [PartialQuestions ]) -> List [List [PartialQuestions ]]:
210+ def _group_partial_messages (
211+ pq_list : List [PartialQuestions ],
212+ ) -> List [List [PartialQuestions ]]: # noqa: C901
202213 """
203214 A PartialQuestion is an object that contains several user messages provided from a
204215 chat conversation. Example:
0 commit comments