""" search_workspace.py 在指定 workspace 的 FAISS 索引中查询文档内容。 每个 workspace 完全隔离，只在当前 workspace 的文档中搜索。 workspace 名称优先从当前工作目录路径自动识别（workspace-xxx 格式）， 也可通过 --workspace 参数手动指定。 用法： ~/faiss_env/bin/python search_workspace.py --query "查询内容" ~/faiss_env/bin/python search_workspace.py --query "查询内容" --top_k 5 --threshold 0.3 ~/faiss_env/bin/python search_workspace.py --workspace workspace-abc --query "查询内容" ~/faiss_env/bin/python search_workspace.py --list # 列出所有已建索引的 workspace """ import os import sys import json import argparse import numpy as np # ===================== 配置 ===================== OPENCLAW_DIR = "/home/admin/.openclaw" INDEX_SUBDIR = ".faiss_index" MODEL_PATH = "/home/bgeSmallZh/" DEFAULT_TOP_K = 5 DEFAULT_THRESHOLD = 0.3 # ===================== 从当前路径提取 workspace 名称 ===================== def extract_workspace_from_cwd() -> str: """ 从当前工作目录路径中识别 workspace- 开头的部分，返回完整 workspace 名称 例如：/home/admin/.openclaw/workspace-abc123/scripts → 返回 "workspace-abc123" """ import re cwd = os.getcwd() for part in cwd.replace("\\", "/").split("/"): if re.match(r'^workspace-.+$', part): return part return None # ===================== 懒加载模型 ===================== _model = None def get_model(): global _model if _model is None: from sentence_transformers import SentenceTransformer _model = SentenceTransformer(MODEL_PATH) return _model # ===================== 加载索引 ===================== def load_index(workspace_name: str): import faiss ws_dir = os.path.join(OPENCLAW_DIR, workspace_name) index_dir = os.path.join(ws_dir, INDEX_SUBDIR) index_file = os.path.join(index_dir, "index.faiss") meta_file = os.path.join(index_dir, "metadata.json") if not os.path.exists(index_file): raise FileNotFoundError( f"索引不存在: {index_file}\n" f"请先运行: ~/faiss_env/bin/python build_workspace_index.py --workspace {workspace_name}" ) index = faiss.read_index(index_file) with open(meta_file, "r", encoding="utf-8") as f: meta = json.load(f) return index, meta # ===================== 搜索 ===================== def search(workspace_name: str, query: str, top_k: int = DEFAULT_TOP_K, threshold: float = DEFAULT_THRESHOLD) -> dict: index, meta = load_index(workspace_name) # 向量化查询 model = get_model() q_vec = model.encode(query, normalize_embeddings=True).astype(np.float32).reshape(1, -1) scores, indices = index.search(q_vec, top_k) results = [] for score, idx in zip(scores[0], indices[0]): if idx == -1: continue score = round(float(score), 4) if score < threshold: continue record = meta["records"][idx] results.append({ "score": score, "file_name": record["file_name"], "file_path": record["file_path"], "chunk_index": record["chunk_index"], "total_chunks": record["total_chunks"], "preview": record["text"][:200] + ("..." if len(record["text"]) > 200 else "") }) if not results: top_score = round(float(scores[0][0]), 4) if len(scores[0]) > 0 and indices[0][0] != -1 else 0.0 return { "status": "no_match", "workspace": workspace_name, "message": f"未找到相似度 >= {threshold} 的内容，最高得分: {top_score}", "results": [] } return { "status": "success", "workspace": workspace_name, "query": query, "total": len(results), "results": results } # ===================== 列出所有已建索引的 workspace ===================== def list_indexed_workspaces() -> list: if not os.path.exists(OPENCLAW_DIR): return [] result = [] for entry in sorted(os.listdir(OPENCLAW_DIR)): if not entry.startswith("workspace-"): continue ws_dir = os.path.join(OPENCLAW_DIR, entry) index_file = os.path.join(ws_dir, INDEX_SUBDIR, "index.faiss") meta_file = os.path.join(ws_dir, INDEX_SUBDIR, "metadata.json") if os.path.exists(index_file) and os.path.exists(meta_file): with open(meta_file, "r", encoding="utf-8") as f: meta = json.load(f) result.append({ "workspace": entry, "total_chunks": meta.get("total", 0), "doc_dir": meta.get("doc_dir", ""), "index_file": index_file }) return result # ===================== 主函数 ===================== def main(): parser = argparse.ArgumentParser(description="在指定 workspace 的文档中语义搜索") parser.add_argument("--workspace", type=str, default=None, help="workspace 名称（可选，不填则从当前路径自动识别 workspace-xxx）") parser.add_argument("--query", type=str, default=None, help="查询内容") parser.add_argument("--top_k", type=int, default=DEFAULT_TOP_K, help=f"返回结果数量，默认 {DEFAULT_TOP_K}") parser.add_argument("--threshold", type=float, default=DEFAULT_THRESHOLD,help=f"最低相似度阈值，默认 {DEFAULT_THRESHOLD}") parser.add_argument("--list", action="store_true", help="列出所有已建索引的 workspace") args = parser.parse_args() # 列出所有 workspace if args.list: workspaces = list_indexed_workspaces() if not workspaces: print(json.dumps({"status": "empty", "message": "暂无已建索引的 workspace"}, ensure_ascii=False, indent=2)) else: print(json.dumps({"status": "success", "workspaces": workspaces}, ensure_ascii=False, indent=2)) return # 确定 workspace：优先用 --workspace 参数，否则从当前路径自动识别 workspace_name = args.workspace if not workspace_name: workspace_name = extract_workspace_from_cwd() if workspace_name: print(json.dumps({"info": f"从当前路径自动识别 workspace: {workspace_name}"}, ensure_ascii=False), file=sys.stderr) else: print(json.dumps({ "status": "error", "message": "未指定 --workspace，且当前路径中未找到 workspace-xxx 格式的目录，无法确定 workspace" }, ensure_ascii=False, indent=2)) sys.exit(1) if not args.query: parser.error("请指定 --query 参数") try: result = search(workspace_name, args.query, top_k=args.top_k, threshold=args.threshold) print(json.dumps(result, ensure_ascii=False, indent=2)) except FileNotFoundError as e: print(json.dumps({"status": "error", "message": str(e)}, ensure_ascii=False, indent=2)) sys.exit(1) except Exception as e: print(json.dumps({"status": "error", "message": f"查询失败: {str(e)}"}, ensure_ascii=False, indent=2)) sys.exit(1) if __name__ == "__main__": main()