Create utils.py

src/txagent/utils.py

@@ -0,0 +1,94 @@
+import sys
+import json
+import hashlib
+import torch
+from typing import List
+
+
+def get_md5(input_str):
+    # Create an MD5 hash object
+    md5_hash = hashlib.md5()
+    md5_hash.update(input_str.encode('utf-8'))
+    return md5_hash.hexdigest()
+
+
+def tool_result_format(function_call_messages):
+    current_output = "\n\n<details>\n<summary> <strong>Verified Feedback from Tools</strong>, click to see details:</summary>\n\n"
+    for each_message in function_call_messages:
+        if each_message['role'] == 'tool':
+            try:
+                parsed = json.loads(each_message['content'])
+                tool_name = parsed.get("tool_name", "Unknown Tool")
+                tool_output = parsed.get("content", each_message['content'])
+                current_output += f"**🔧 Tool: {tool_name}**\n\n{tool_output}\n\n"
+            except Exception:
+                current_output += f"{each_message['content']}\n\n"
+    current_output += "</details>\n\n\n"
+    return current_output
+
+
+class NoRepeatSentenceProcessor:
+    def __init__(self, forbidden_sequences: List[List[int]], allowed_prefix_length: int):
+        self.allowed_prefix_length = allowed_prefix_length
+        self.forbidden_prefix_dict = {}
+        for seq in forbidden_sequences:
+            if len(seq) > allowed_prefix_length:
+                prefix = tuple(seq[:allowed_prefix_length])
+                next_token = seq[allowed_prefix_length]
+                self.forbidden_prefix_dict.setdefault(prefix, set()).add(next_token)
+
+    def __call__(self, token_ids: List[int], logits: torch.Tensor) -> torch.Tensor:
+        if len(token_ids) >= self.allowed_prefix_length:
+            prefix = tuple(token_ids[:self.allowed_prefix_length])
+            if prefix in self.forbidden_prefix_dict:
+                for token_id in self.forbidden_prefix_dict[prefix]:
+                    logits[token_id] = -float("inf")
+        return logits
+
+
+class ReasoningTraceChecker:
+    def __init__(self, question, conversation, init_index=None):
+        self.question = question.lower()
+        self.conversation = conversation
+        self.existing_thoughts = []
+        self.existing_actions = []
+        self.new_thoughts = []
+        self.new_actions = []
+        self.index = init_index if init_index is not None else 1
+
+    def check_conversation(self):
+        info = ''
+        current_index = self.index
+        for i in range(current_index, len(self.conversation)):
+            each = self.conversation[i]
+            self.index = i
+            if each['role'] == 'assistant':
+                thought = each['content']
+                actions = each['tool_calls']
+                good_status, current_info = self.check_repeat_thought(thought)
+                info += current_info
+                if not good_status:
+                    return False, info
+                good_status, current_info = self.check_repeat_action(actions)
+                info += current_info
+                if not good_status:
+                    return False, info
+        return True, info
+
+    def check_repeat_thought(self, thought):
+        if thought in self.existing_thoughts:
+            return False, "repeat_thought"
+        self.existing_thoughts.append(thought)
+        return True, ''
+
+    def check_repeat_action(self, actions):
+        if type(actions) != list:
+            actions = json.loads(actions)
+        for each_action in actions:
+            if 'call_id' in each_action:
+                del each_action['call_id']
+            each_action = json.dumps(each_action)
+            if each_action in self.existing_actions:
+                return False, "repeat_action"
+            self.existing_actions.append(each_action)
+        return True, ''