""" search_public.py 根据当前路径自动识别用户 ID，查询企业知识库： 1. 从当前工作目录路径中识别 workspace_xxx，提取 xxx 作为用户 ID 2. 通过 MySQL 查询 user_id → project_id 3. 加载对应企业项目的 FAISS 索引 4. 在企业知识库中语义搜索 用法： ~/faiss_env/bin/python search_public.py --query "查询内容" ~/faiss_env/bin/python search_public.py --query "查询内容" --top_k 5 --threshold 0.3 ~/faiss_env/bin/python search_public.py --list # 列出所有已建索引的企业项目 # 也可手动指定用户 ID（覆盖路径自动识别） ~/faiss_env/bin/python search_public.py --user_id 123 --query "查询内容" """ import os import sys import json import argparse import numpy as np # ===================== 配置 ===================== PUBLIC_DOC_DIR = "/home/admin/.openclaw/PublicDocument" INDEX_SUBDIR = ".faiss_index" MODEL_PATH = "/home/bgeSmallZh/" DEFAULT_TOP_K = 5 DEFAULT_THRESHOLD = 0.3 # MySQL 配置 MYSQL_CONFIG = { "host": "47.96.177.56", "port": 3306, "user": "yunying", "password": "123456", "database": "chat_gpt_new", "charset": "utf8mb4" } # ===================== 懒加载模型 ===================== _model = None def get_model(): global _model if _model is None: from sentence_transformers import SentenceTransformer _model = SentenceTransformer(MODEL_PATH) return _model # ===================== 从当前路径提取用户 ID ===================== def extract_user_id_from_cwd() -> str: """ 从当前工作目录路径中识别 workspace- 开头的部分，提取后面的用户 ID 例如：/home/admin/.openclaw/workspace-12345/xxx → 返回 "12345" """ import re cwd = os.getcwd() # 将路径按分隔符拆分，逐段匹配 workspace-xxx for part in cwd.replace("\\", "/").split("/"): match = re.match(r'^workspace-(.+)$', part) if match: return match.group(1) return None # ===================== 查询用户所属企业项目 ===================== def get_project_id_by_user(user_id: str) -> str: """ 通过 MySQL 查询用户所属的企业项目 ID SQL: SELECT project_id FROM gpt_user WHERE id = {user_id} """ import pymysql try: conn = pymysql.connect(**MYSQL_CONFIG) cursor = conn.cursor() cursor.execute("SELECT project_id FROM gpt_user WHERE id = %s", (user_id,)) result = cursor.fetchone() cursor.close() conn.close() if not result or result[0] is None: raise ValueError(f"用户 {user_id} 未关联任何企业项目（project_id 为空）") return str(result[0]) except Exception as e: raise RuntimeError(f"MySQL 查询失败: {str(e)}") # ===================== 加载企业项目索引 ===================== def load_index(project_id: str): import faiss project_dir = os.path.join(PUBLIC_DOC_DIR, f"Document-{project_id}") index_dir = os.path.join(project_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"企业项目 {project_id} 的索引不存在: {index_file}\n" f"请先运行: ~/faiss_env/bin/python build_public_index.py --project {project_id}" ) 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(user_id: str, query: str, top_k: int = DEFAULT_TOP_K, threshold: float = DEFAULT_THRESHOLD) -> dict: # 第一步：查询用户所属企业项目 try: project_id = get_project_id_by_user(user_id) except Exception as e: return { "status": "error", "message": str(e), "user_id": user_id } # 第二步：加载企业项目索引 try: index, meta = load_index(project_id) except FileNotFoundError as e: return { "status": "error", "message": str(e), "user_id": user_id, "project_id": project_id } # 第三步：向量化查询 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", "user_id": user_id, "project_id": project_id, "message": f"未找到相似度 >= {threshold} 的内容，最高得分: {top_score}", "results": [] } return { "status": "success", "user_id": user_id, "project_id": project_id, "query": query, "total": len(results), "results": results } # ===================== 列出所有已建索引的企业项目 ===================== def list_indexed_projects() -> list: if not os.path.exists(PUBLIC_DOC_DIR): return [] result = [] for entry in sorted(os.listdir(PUBLIC_DOC_DIR)): if not entry.startswith("Document-"): continue project_id = entry[len("Document-"):] project_dir = os.path.join(PUBLIC_DOC_DIR, entry) index_file = os.path.join(project_dir, INDEX_SUBDIR, "index.faiss") meta_file = os.path.join(project_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({ "project_id": project_id, "total_chunks": meta.get("total", 0), "doc_dir": meta.get("doc_dir", ""), "index_file": index_file }) return result # ===================== 主函数 ===================== def main(): parser = argparse.ArgumentParser(description="根据用户 ID 在企业知识库中语义搜索") parser.add_argument("--user_id", type=str, default=None, help="用户 ID（可选，不填则从当前路径自动识别 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="列出所有已建索引的企业项目") args = parser.parse_args() # 列出所有企业项目 if args.list: projects = list_indexed_projects() if not projects: print(json.dumps({"status": "empty", "message": "暂无已建索引的企业项目"}, ensure_ascii=False, indent=2)) else: print(json.dumps({"status": "success", "projects": projects}, ensure_ascii=False, indent=2)) return # 确定用户 ID：优先用 --user_id 参数，否则从当前路径自动识别 user_id = args.user_id if not user_id: user_id = extract_user_id_from_cwd() if user_id: print(json.dumps({"info": f"从当前路径自动识别用户 ID: {user_id}"}, ensure_ascii=False), file=sys.stderr) else: print(json.dumps({ "status": "error", "message": "未指定 --user_id，且当前路径中未找到 workspace-xxx 格式的目录，无法确定用户 ID" }, ensure_ascii=False, indent=2)) sys.exit(1) if not args.query: parser.error("请指定 --query 参数") try: result = search(user_id, args.query, top_k=args.top_k, threshold=args.threshold) print(json.dumps(result, ensure_ascii=False, indent=2)) if result["status"] == "error": 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()