405 lines
12 KiB
Python
405 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
|
import os
|
|
import os.path as osp
|
|
import json
|
|
import logging
|
|
import hashlib
|
|
import time
|
|
from typing import Optional
|
|
|
|
from langchain_community.vectorstores import FAISS
|
|
|
|
TOOLS_INFO_PATH = osp.join(osp.dirname(__file__), "dj_funcs_all.json")
|
|
CACHE_RETRIEVED_TOOLS_PATH = osp.join(osp.dirname(__file__), "cache_retrieve")
|
|
VECTOR_INDEX_CACHE_PATH = osp.join(osp.dirname(__file__), "vector_index_cache")
|
|
|
|
# Global variable to cache the vector store
|
|
_cached_vector_store: Optional[FAISS] = None
|
|
_cached_tools_info: Optional[list] = None
|
|
_cached_file_hash: Optional[str] = None
|
|
|
|
RETRIEVAL_PROMPT = """You are a professional tool retrieval assistant
|
|
responsible for filtering the top {limit} most relevant tools from a large
|
|
tool library based on user requirements. Execute the following steps:
|
|
|
|
# Requirement Analysis
|
|
Carefully read the user's [requirement description], extract core keywords,
|
|
functional objectives, usage scenarios, and technical requirements
|
|
(such as real-time performance, data types, industry domains, etc.).
|
|
|
|
# Tool Matching
|
|
Perform multi-dimensional matching based on the following tool attributes:
|
|
- Tool name and functional description
|
|
- Supported input/output formats
|
|
- Applicable industry or scenario tags
|
|
- Technical implementation principles
|
|
(API, local deployment, AI model types)
|
|
- Relevance ranking
|
|
|
|
# Use weighted scoring mechanism (example weights):
|
|
- Functional match (40%)
|
|
- Scenario compatibility (30%)
|
|
- Technical compatibility (20%)
|
|
- User rating/usage rate (10%)
|
|
|
|
# Deduplication and Optimization
|
|
Exclude the following low-quality results:
|
|
- Tools with duplicate functionality (keep only the best one)
|
|
- Tools that cannot meet basic requirements
|
|
- Tools missing critical parameter descriptions
|
|
|
|
# Constraints
|
|
- Strictly control output to a maximum of {limit} tools
|
|
- Refuse to speculate on unknown tool attributes
|
|
- Maintain accuracy of domain expertise
|
|
|
|
# Output Format
|
|
Return a JSON format TOP{limit} tool list containing:
|
|
[
|
|
{{
|
|
"rank": 1,
|
|
"tool_name": "Tool Name",
|
|
"description": "Core functionality summary",
|
|
"relevance_score": 98.7,
|
|
"key_match": ["Matching keywords/features"]
|
|
}}
|
|
]
|
|
Output strictly in JSON array format, and only output the JSON array format
|
|
tool list.
|
|
"""
|
|
|
|
|
|
def fast_text_encoder(text: str) -> str:
|
|
"""Fast encoding using xxHash algorithm"""
|
|
import xxhash
|
|
|
|
hasher = xxhash.xxh64(seed=0)
|
|
hasher.update(text.encode("utf-8"))
|
|
|
|
# Return 16-bit hexadecimal string
|
|
return hasher.hexdigest()
|
|
|
|
|
|
async def retrieve_ops_lm(user_query, limit=20):
|
|
"""Tool retrieval using language model - returns list of tool names"""
|
|
hash_id = fast_text_encoder(user_query + str(limit))
|
|
|
|
# Ensure cache directory exists
|
|
os.makedirs(CACHE_RETRIEVED_TOOLS_PATH, exist_ok=True)
|
|
|
|
cache_tools_path = osp.join(CACHE_RETRIEVED_TOOLS_PATH, f"{hash_id}.json")
|
|
if osp.exists(cache_tools_path):
|
|
with open(cache_tools_path, "r", encoding="utf-8") as f:
|
|
return json.loads(f.read())
|
|
|
|
if osp.exists(TOOLS_INFO_PATH):
|
|
with open(TOOLS_INFO_PATH, "r", encoding="utf-8") as f:
|
|
dj_func_info = json.loads(f.read())
|
|
tool_descriptions = [
|
|
f"{t['class_name']}: {t['class_desc']}" for t in dj_func_info
|
|
]
|
|
tools_string = "\n".join(tool_descriptions)
|
|
else:
|
|
from create_dj_func_info import dj_func_info
|
|
|
|
project_root = os.path.abspath(
|
|
os.path.join(os.path.dirname(__file__), ".."),
|
|
)
|
|
|
|
with open(
|
|
os.path.join(project_root, TOOLS_INFO_PATH),
|
|
"w",
|
|
encoding="utf-8",
|
|
) as f:
|
|
f.write(json.dumps(dj_func_info))
|
|
|
|
tool_descriptions = [
|
|
f"{t['class_name']}: {t['class_desc']}" for t in dj_func_info
|
|
]
|
|
tools_string = "\n".join(tool_descriptions)
|
|
|
|
from agentscope.model import DashScopeChatModel
|
|
from agentscope.message import Msg
|
|
from agentscope.formatter import DashScopeChatFormatter
|
|
|
|
model = DashScopeChatModel(
|
|
model_name="qwen-turbo",
|
|
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
|
stream=False,
|
|
)
|
|
|
|
formatter = DashScopeChatFormatter()
|
|
|
|
# Update retrieval prompt to use the specified limit
|
|
retrieval_prompt_with_limit = RETRIEVAL_PROMPT.format(limit=limit)
|
|
|
|
user_prompt = (
|
|
retrieval_prompt_with_limit
|
|
+ f"""
|
|
User requirement description:
|
|
{user_query}
|
|
|
|
Available tools:
|
|
{tools_string}
|
|
"""
|
|
)
|
|
|
|
msgs = [
|
|
Msg(name="user", role="user", content=user_prompt),
|
|
]
|
|
|
|
formatted_msgs = await formatter.format(msgs)
|
|
|
|
response = await model(formatted_msgs)
|
|
|
|
msg = Msg(name="assistant", role="assistant", content=response.content)
|
|
retrieved_tools_text = msg.get_text_content()
|
|
retrieved_tools = json.loads(retrieved_tools_text)
|
|
|
|
# Extract tool names and validate they exist
|
|
tool_names = []
|
|
for tool_info in retrieved_tools:
|
|
if not isinstance(tool_info, dict) or "tool_name" not in tool_info:
|
|
logging.warning(f"Invalid tool info format: {tool_info}")
|
|
continue
|
|
|
|
tool_name = tool_info["tool_name"]
|
|
|
|
# Verify tool exists in dj_func_info
|
|
tool_exists = any(t["class_name"] == tool_name for t in dj_func_info)
|
|
if not tool_exists:
|
|
logging.error(f"Tool not found: `{tool_name}`, skipping!")
|
|
continue
|
|
|
|
tool_names.append(tool_name)
|
|
|
|
# Cache the result
|
|
with open(cache_tools_path, "w", encoding="utf-8") as f:
|
|
json.dump(tool_names, f)
|
|
|
|
return tool_names
|
|
|
|
|
|
def _get_file_hash(file_path: str) -> str:
|
|
"""Get file content hash using SHA256"""
|
|
try:
|
|
with open(file_path, "rb") as f:
|
|
file_content = f.read()
|
|
return hashlib.sha256(file_content).hexdigest()
|
|
except (OSError, IOError):
|
|
return ""
|
|
|
|
|
|
def _load_cached_index() -> bool:
|
|
"""Load cached vector index from disk"""
|
|
global _cached_vector_store, _cached_tools_info, _cached_file_hash
|
|
|
|
try:
|
|
# Ensure cache directory exists
|
|
os.makedirs(VECTOR_INDEX_CACHE_PATH, exist_ok=True)
|
|
|
|
index_path = osp.join(VECTOR_INDEX_CACHE_PATH, "faiss_index")
|
|
metadata_path = osp.join(VECTOR_INDEX_CACHE_PATH, "metadata.json")
|
|
|
|
if not all(os.path.exists(p) for p in [index_path, metadata_path]):
|
|
return False
|
|
|
|
# Check if cached index matches current tools info file
|
|
with open(metadata_path, "r", encoding="utf-8") as f:
|
|
metadata = json.load(f)
|
|
|
|
cached_hash = metadata.get("tools_info_hash", "")
|
|
current_hash = _get_file_hash(TOOLS_INFO_PATH)
|
|
|
|
if current_hash != cached_hash:
|
|
return False
|
|
|
|
# Load cached data
|
|
from langchain_community.embeddings import DashScopeEmbeddings
|
|
|
|
embeddings = DashScopeEmbeddings(
|
|
dashscope_api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
|
model="text-embedding-v1",
|
|
)
|
|
|
|
_cached_vector_store = FAISS.load_local(
|
|
index_path,
|
|
embeddings,
|
|
allow_dangerous_deserialization=True,
|
|
)
|
|
|
|
_cached_file_hash = cached_hash
|
|
|
|
logging.info("Successfully loaded cached vector index")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logging.warning(f"Failed to load cached index: {e}")
|
|
return False
|
|
|
|
|
|
def _save_cached_index():
|
|
"""Save vector index to disk cache"""
|
|
global _cached_vector_store, _cached_file_hash
|
|
|
|
try:
|
|
# Ensure cache directory exists
|
|
os.makedirs(VECTOR_INDEX_CACHE_PATH, exist_ok=True)
|
|
|
|
index_path = osp.join(VECTOR_INDEX_CACHE_PATH, "faiss_index")
|
|
metadata_path = osp.join(VECTOR_INDEX_CACHE_PATH, "metadata.json")
|
|
|
|
# Save vector store
|
|
if _cached_vector_store:
|
|
_cached_vector_store.save_local(index_path)
|
|
|
|
# Save metadata
|
|
metadata = {
|
|
"tools_info_hash": _cached_file_hash,
|
|
"created_at": time.time(),
|
|
}
|
|
with open(metadata_path, "w", encoding="utf-8") as f:
|
|
json.dump(metadata, f)
|
|
|
|
logging.info("Successfully saved vector index to cache")
|
|
|
|
except Exception as e:
|
|
logging.error(f"Failed to save cached index: {e}")
|
|
|
|
|
|
def _build_vector_index():
|
|
"""Build and cache vector index"""
|
|
global _cached_vector_store, _cached_file_hash
|
|
|
|
with open(TOOLS_INFO_PATH, "r", encoding="utf-8") as f:
|
|
tools_info = json.loads(f.read())
|
|
|
|
tool_descriptions = [
|
|
f"{t['class_name']}: {t['class_desc']}" for t in tools_info
|
|
]
|
|
|
|
from langchain_community.embeddings import DashScopeEmbeddings
|
|
|
|
embeddings = DashScopeEmbeddings(
|
|
dashscope_api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
|
model="text-embedding-v1",
|
|
)
|
|
|
|
metadatas = [{"index": i} for i in range(len(tool_descriptions))]
|
|
vector_store = FAISS.from_texts(
|
|
tool_descriptions,
|
|
embeddings,
|
|
metadatas=metadatas,
|
|
)
|
|
|
|
# Cache the results
|
|
_cached_vector_store = vector_store
|
|
_cached_file_hash = _get_file_hash(TOOLS_INFO_PATH)
|
|
|
|
# Save to disk cache
|
|
_save_cached_index()
|
|
|
|
logging.info("Successfully built and cached vector index")
|
|
|
|
|
|
def retrieve_ops_vector(user_query, limit=20):
|
|
"""Tool retrieval using vector search with caching"""
|
|
global _cached_vector_store
|
|
|
|
# Try to load from cache first
|
|
if not _load_cached_index():
|
|
logging.info("Building new vector index...")
|
|
_build_vector_index()
|
|
|
|
# Perform similarity search
|
|
retrieved_tools = _cached_vector_store.similarity_search(
|
|
user_query,
|
|
k=limit,
|
|
)
|
|
retrieved_indices = [doc.metadata["index"] for doc in retrieved_tools]
|
|
|
|
with open(TOOLS_INFO_PATH, "r", encoding="utf-8") as f:
|
|
tools_info = json.loads(f.read())
|
|
|
|
# Extract tool names from retrieved indices
|
|
tool_names = []
|
|
for raw_idx in retrieved_indices:
|
|
tool_info = tools_info[raw_idx]
|
|
tool_names.append(tool_info["class_name"])
|
|
|
|
return tool_names
|
|
|
|
|
|
async def retrieve_ops(
|
|
user_query: str,
|
|
limit: int = 20,
|
|
mode: str = "auto",
|
|
) -> list:
|
|
"""
|
|
Tool retrieval with configurable mode
|
|
|
|
Args:
|
|
user_query: User query string
|
|
limit: Maximum number of tools to retrieve
|
|
mode: Retrieval mode - "llm", "vector", or "auto" (default: "auto")
|
|
- "llm": Use language model only
|
|
- "vector": Use vector search only
|
|
- "auto": Try LLM first, fallback to vector search on failure
|
|
|
|
Returns:
|
|
List of tool names
|
|
"""
|
|
if mode in ("llm", "auto"):
|
|
try:
|
|
return await retrieve_ops_lm(user_query, limit=limit)
|
|
except Exception as e:
|
|
logging.error(f"LLM retrieval failed: {str(e)}")
|
|
if mode != "auto":
|
|
return []
|
|
|
|
if mode in ("vector", "auto"):
|
|
try:
|
|
return retrieve_ops_vector(user_query, limit=limit)
|
|
except Exception as e:
|
|
logging.error(f"Vector retrieval failed: {str(e)}")
|
|
return []
|
|
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid mode: {mode}. Must be 'llm', 'vector', or 'auto'",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import asyncio
|
|
|
|
query = (
|
|
"Clean special characters from text and filter samples with "
|
|
+ "excessive length. Mask sensitive information and filter "
|
|
+ "unsafe content including adult/terror-related terms."
|
|
+ "Additionally, filter out small images, perform image "
|
|
+ "tagging, and remove duplicate images."
|
|
)
|
|
|
|
# Test different modes
|
|
print("=== Testing LLM mode ===")
|
|
tool_names_llm = asyncio.run(
|
|
retrieve_ops(query, limit=10, mode="llm"),
|
|
)
|
|
print("Retrieved tool names (LLM):")
|
|
print(tool_names_llm)
|
|
|
|
print("\n=== Testing Vector mode ===")
|
|
tool_names_vector = asyncio.run(
|
|
retrieve_ops(query, limit=10, mode="vector"),
|
|
)
|
|
print("Retrieved tool names (Vector):")
|
|
print(tool_names_vector)
|
|
|
|
print("\n=== Testing Auto mode (default) ===")
|
|
tool_names_auto = asyncio.run(
|
|
retrieve_ops(query, limit=10, mode="auto"),
|
|
)
|
|
print("Retrieved tool names (Auto):")
|
|
print(tool_names_auto)
|