pvanand commited on
Commit
49af561
·
verified ·
1 Parent(s): ffe1325

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +9 -11
main.py CHANGED
@@ -132,24 +132,22 @@ def update_db(user_id, conversation_id, message, response):
132
  # Utility Functions
133
  # ============================================================================
134
 
135
- def extract_data_from_tag(input_string, tag, invert=False):
136
- # Create the regex pattern
137
  pattern = f'<{tag}.*?>(.*?)</{tag}>'
138
-
139
- # Find all matches
140
  matches = re.findall(pattern, input_string, re.DOTALL)
141
 
142
  if invert:
143
- # If inverted, replace all matches with empty string
144
- out = re.sub(pattern, '', input_string, flags=re.DOTALL)
145
- return out.strip()
 
 
146
  else:
147
- # If matches are found, return them joined by newlines
148
  if matches:
149
- out = '\n'.join(match.strip() for match in matches)
150
- return out
151
  else:
152
- return input_string
 
153
 
154
  def calculate_tokens(msgs):
155
  return sum(len(encoding.encode(str(m))) for m in msgs)
 
132
  # Utility Functions
133
  # ============================================================================
134
 
135
+ def extract_data_from_tag(input_string, tag, invert=False, return_full_if_not_found=False):
 
136
  pattern = f'<{tag}.*?>(.*?)</{tag}>'
 
 
137
  matches = re.findall(pattern, input_string, re.DOTALL)
138
 
139
  if invert:
140
+ if matches:
141
+ out = re.sub(pattern, '', input_string, flags=re.DOTALL)
142
+ return out.strip()
143
+ else:
144
+ return input_string if return_full_if_not_found else ""
145
  else:
 
146
  if matches:
147
+ return '\n'.join(match.strip() for match in matches)
 
148
  else:
149
+ return input_string if return_full_if_not_found else ""
150
+
151
 
152
  def calculate_tokens(msgs):
153
  return sum(len(encoding.encode(str(m))) for m in msgs)