Create toolrag.py

src/txagent/toolrag.py

@@ -0,0 +1,67 @@
+import os
+import json
+import torch
+from sentence_transformers import SentenceTransformer
+from .utils import get_md5
+
+
+class ToolRAGModel:
+    def __init__(self, rag_model_name):
+        self.rag_model_name = rag_model_name
+        self.rag_model = None
+        self.tool_desc_embedding = None
+        self.tool_name = None
+        self.tool_embedding_path = None
+        self.load_rag_model()
+
+    def load_rag_model(self):
+        self.rag_model = SentenceTransformer(self.rag_model_name)
+        self.rag_model.max_seq_length = 4096
+        self.rag_model.tokenizer.padding_side = "right"
+
+    def load_tool_desc_embedding(self, toolbox):
+        self.tool_name, _ = toolbox.refresh_tool_name_desc(enable_full_desc=True)
+        all_tools_str = [json.dumps(each) for each in toolbox.prepare_tool_prompts(toolbox.all_tools)]
+        md5_value = get_md5(str(all_tools_str))
+        print("Computed MD5 for tool embedding:", md5_value)
+
+        self.tool_embedding_path = os.path.join(
+            os.path.dirname(__file__),
+            self.rag_model_name.split("/")[-1] + f"_tool_embedding_{md5_value}.pt"
+        )
+
+        if os.path.exists(self.tool_embedding_path):
+            try:
+                self.tool_desc_embedding = torch.load(self.tool_embedding_path, map_location="cpu")
+                assert len(self.tool_desc_embedding) == len(toolbox.all_tools), \
+                    "Tool count mismatch with loaded embeddings."
+                print("\033[92mLoaded cached tool_desc_embedding.\033[0m")
+                return
+            except Exception as e:
+                print(f"⚠️ Failed loading cached embeddings: {e}")
+                self.tool_desc_embedding = None
+
+        print("\033[93mGenerating new tool_desc_embedding...\033[0m")
+        self.tool_desc_embedding = self.rag_model.encode(
+            all_tools_str, prompt="", normalize_embeddings=True
+        )
+
+        torch.save(self.tool_desc_embedding, self.tool_embedding_path)
+        print(f"\033[92mSaved new tool_desc_embedding to {self.tool_embedding_path}\033[0m")
+
+    def rag_infer(self, query, top_k=5):
+        torch.cuda.empty_cache()
+        queries = [query]
+        query_embeddings = self.rag_model.encode(
+            queries, prompt="", normalize_embeddings=True
+        )
+        if self.tool_desc_embedding is None:
+            raise RuntimeError("❌ tool_desc_embedding is not initialized. Did you forget to call load_tool_desc_embedding()?")
+
+        scores = self.rag_model.similarity(
+            query_embeddings, self.tool_desc_embedding
+        )
+        top_k = min(top_k, len(self.tool_name))
+        top_k_indices = torch.topk(scores, top_k).indices.tolist()[0]
+        top_k_tool_names = [self.tool_name[i] for i in top_k_indices]
+        return top_k_tool_names