#!/usr/bin/env python3 """ Entity Aligner - Phase 3 使用 embedding 进行实体对齐 """ import json import os from dataclasses import dataclass import httpx import numpy as np # API Keys KIMI_API_KEY = os.getenv("KIMI_API_KEY", "") KIMI_BASE_URL = os.getenv("KIMI_BASE_URL", "https://api.kimi.com/coding") @dataclass class EntityEmbedding: entity_id: str name: str definition: str embedding: list[float] class EntityAligner: """实体对齐器 - 使用 embedding 进行相似度匹配""" def __init__(self, similarity_threshold: float = 0.85): self.similarity_threshold = similarity_threshold self.embedding_cache: dict[str, list[float]] = {} def get_embedding(self, text: str) -> list[float] | None: """ 使用 Kimi API 获取文本的 embedding Args: text: 输入文本 Returns: embedding 向量或 None """ if not KIMI_API_KEY: return None # 检查缓存 cache_key = hash(text) if cache_key in self.embedding_cache: return self.embedding_cache[cache_key] try: response = httpx.post( f"{KIMI_BASE_URL}/v1/embeddings", headers={ "Authorization": f"Bearer {KIMI_API_KEY}", "Content-Type": "application/json", }, json={"model": "k2p5", "input": text[:500]}, # 限制长度 timeout=30.0, ) response.raise_for_status() result = response.json() embedding = result["data"][0]["embedding"] self.embedding_cache[cache_key] = embedding return embedding except Exception as e: print(f"Embedding API failed: {e}") return None def compute_similarity(self, embedding1: list[float], embedding2: list[float]) -> float: """ 计算两个 embedding 的余弦相似度 Args: embedding1: 第一个向量 embedding2: 第二个向量 Returns: 相似度分数 (0-1) """ vec1 = np.array(embedding1) vec2 = np.array(embedding2) # 余弦相似度 dot_product = np.dot(vec1, vec2) norm1 = np.linalg.norm(vec1) norm2 = np.linalg.norm(vec2) if norm1 == 0 or norm2 == 0: return 0.0 return float(dot_product / (norm1 * norm2)) def get_entity_text(self, name: str, definition: str = "") -> str: """ 构建用于 embedding 的实体文本 Args: name: 实体名称 definition: 实体定义 Returns: 组合文本 """ if definition: return f"{name}: {definition}" return name def find_similar_entity( self, project_id: str, name: str, definition: str = "", exclude_id: str | None = None, threshold: float | None = None, ) -> object | None: """ 查找相似的实体 Args: project_id: 项目 ID name: 实体名称 definition: 实体定义 exclude_id: 要排除的实体 ID threshold: 相似度阈值 Returns: 相似的实体或 None """ if threshold is None: threshold = self.similarity_threshold try: from db_manager import get_db_manager db = get_db_manager() except ImportError: return None # 获取项目的所有实体 entities = db.get_all_entities_for_embedding(project_id) if not entities: return None # 获取查询实体的 embedding query_text = self.get_entity_text(name, definition) query_embedding = self.get_embedding(query_text) if query_embedding is None: # 如果 embedding API 失败，回退到简单匹配 return self._fallback_similarity_match(entities, name, exclude_id) best_match = None best_score = threshold for entity in entities: if exclude_id and entity.id == exclude_id: continue # 获取实体的 embedding entity_text = self.get_entity_text(entity.name, entity.definition) entity_embedding = self.get_embedding(entity_text) if entity_embedding is None: continue # 计算相似度 similarity = self.compute_similarity(query_embedding, entity_embedding) if similarity > best_score: best_score = similarity best_match = entity return best_match def _fallback_similarity_match( self, entities: list[object], name: str, exclude_id: str | None = None ) -> object | None: """ 回退到简单的相似度匹配（不使用 embedding） Args: entities: 实体列表 name: 查询名称 exclude_id: 要排除的实体 ID Returns: 最相似的实体或 None """ name_lower = name.lower() # 1. 精确匹配 for entity in entities: if exclude_id and entity.id == exclude_id: continue if entity.name.lower() == name_lower: return entity if entity.aliases and name_lower in [a.lower() for a in entity.aliases]: return entity # 2. 包含匹配 for entity in entities: if exclude_id and entity.id == exclude_id: continue if name_lower in entity.name.lower() or entity.name.lower() in name_lower: return entity return None def batch_align_entities( self, project_id: str, new_entities: list[dict], threshold: float | None = None ) -> list[dict]: """ 批量对齐实体 Args: project_id: 项目 ID new_entities: 新实体列表 [{"name": "...", "definition": "..."}] threshold: 相似度阈值 Returns: 对齐结果列表 [{"new_entity": {...}, "matched_entity": {...}, "similarity": 0.9}] """ if threshold is None: threshold = self.similarity_threshold results = [] for new_ent in new_entities: matched = self.find_similar_entity( project_id, new_ent["name"], new_ent.get("definition", ""), threshold=threshold ) result = { "new_entity": new_ent, "matched_entity": None, "similarity": 0.0, "should_merge": False, } if matched: # 计算相似度 query_text = self.get_entity_text(new_ent["name"], new_ent.get("definition", "")) matched_text = self.get_entity_text(matched.name, matched.definition) query_emb = self.get_embedding(query_text) matched_emb = self.get_embedding(matched_text) if query_emb and matched_emb: similarity = self.compute_similarity(query_emb, matched_emb) result["matched_entity"] = { "id": matched.id, "name": matched.name, "type": matched.type, "definition": matched.definition, } result["similarity"] = similarity result["should_merge"] = similarity >= threshold results.append(result) return results def suggest_entity_aliases(self, entity_name: str, entity_definition: str = "") -> list[str]: """ 使用 LLM 建议实体的别名 Args: entity_name: 实体名称 entity_definition: 实体定义 Returns: 建议的别名列表 """ if not KIMI_API_KEY: return [] prompt = f"""为以下实体生成可能的别名或简称： 实体名称：{entity_name} 定义：{entity_definition} 请返回 JSON 格式的别名列表： {{"aliases": ["别名1", "别名2", "别名3"]}} 只返回 JSON，不要其他内容。""" try: response = httpx.post( f"{KIMI_BASE_URL}/v1/chat/completions", headers={ "Authorization": f"Bearer {KIMI_API_KEY}", "Content-Type": "application/json", }, json={ "model": "k2p5", "messages": [{"role": "user", "content": prompt}], "temperature": 0.3, }, timeout=30.0, ) response.raise_for_status() result = response.json() content = result["choices"][0]["message"]["content"] import re json_match = re.search(r"\{{.*?\}}", content, re.DOTALL) if json_match: data = json.loads(json_match.group()) return data.get("aliases", []) except Exception as e: print(f"Alias suggestion failed: {e}") return [] # 简单的字符串相似度计算（不使用 embedding） def simple_similarity(str1: str, str2: str) -> float: """ 计算两个字符串的简单相似度 Args: str1: 第一个字符串 str2: 第二个字符串 Returns: 相似度分数 (0-1) """ if str1 == str2: return 1.0 if not str1 or not str2: return 0.0 # 转换为小写 s1 = str1.lower() s2 = str2.lower() # 包含关系 if s1 in s2 or s2 in s1: return 0.8 # 计算编辑距离相似度 from difflib import SequenceMatcher return SequenceMatcher(None, s1, s2).ratio() if __name__ == "__main__": # 测试 aligner = EntityAligner() # 测试 embedding test_text = "Kubernetes 容器编排平台" embedding = aligner.get_embedding(test_text) if embedding: print(f"Embedding dimension: {len(embedding)}") print(f"First 5 values: {embedding[:5]}") else: print("Embedding API not available") # 测试相似度计算 emb1 = [1.0, 0.0, 0.0] emb2 = [0.9, 0.1, 0.0] sim = aligner.compute_similarity(emb1, emb2) print(f"Similarity: {sim:.4f}")