[MODIFY] Code optimization

api/app/repositories/neo4j/entity_repository.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 """实体仓储模块
 
 本模块提供实体节点的数据访问功能。
@@ -7,7 +6,7 @@ Classes:
     EntityRepository: 实体仓储，管理ExtractedEntityNode的CRUD操作
 """
 
-from typing import List, Optional, Dict
+from typing import List, Dict
 from datetime import datetime
 
 from app.repositories.neo4j.base_neo4j_repository import BaseNeo4jRepository
@@ -49,9 +48,13 @@ class EntityRepository(BaseNeo4jRepository[ExtractedEntityNode]):
         # 处理datetime字段
         if isinstance(n.get('created_at'), str):
             n['created_at'] = datetime.fromisoformat(n['created_at'])
-        if n.get('expired_at') and isinstance(n['expired_at'], str):
+        if n.get('expired_at') and isinstance(n.get('expired_at'), str):
             n['expired_at'] = datetime.fromisoformat(n['expired_at'])
         
+        # 确保aliases字段存在且为列表
+        if 'aliases' not in n or n['aliases'] is None:
+            n['aliases'] = []
+        
         return ExtractedEntityNode(**n)
     
     async def find_by_type(self, entity_type: str, limit: int = 100) -> List[ExtractedEntityNode]:
@@ -66,274 +69,4 @@ class EntityRepository(BaseNeo4jRepository[ExtractedEntityNode]):
         """
         return await self.find({"entity_type": entity_type}, limit=limit)
     
-    async def find_by_group_id(self, group_id: str, limit: int = 100) -> List[ExtractedEntityNode]:
-        """根据group_id查询实体
-        
-        Args:
-            group_id: 组ID
-            limit: 返回结果的最大数量
-            
-        Returns:
-            List[ExtractedEntityNode]: 实体列表
-        """
-        return await self.find({"group_id": group_id}, limit=limit)
-    
-    async def find_by_name(
-        self,
-        name: str,
-        group_id: Optional[str] = None,
-        limit: int = 100
-    ) -> List[ExtractedEntityNode]:
-        """根据名称查询实体
-        
-        支持模糊匹配（CONTAINS）。
-        
-        Args:
-            name: 实体名称
-            group_id: 可选的组ID过滤
-            limit: 返回结果的最大数量
-            
-        Returns:
-            List[ExtractedEntityNode]: 实体列表
-        """
-        where_clause = "n.name CONTAINS $name"
-        if group_id:
-            where_clause += " AND n.group_id = $group_id"
-        
-        query = f"""
-        MATCH (n:{self.node_label})
-        WHERE {where_clause}
-        RETURN n
-        LIMIT $limit
-        """
-        
-        params = {"name": name, "limit": limit}
-        if group_id:
-            params["group_id"] = group_id
-        
-        results = await self.connector.execute_query(query, **params)
-        return [self._map_to_entity(r) for r in results]
-    
-    async def find_related_entities(
-        self,
-        entity_id: str,
-        relation_type: Optional[str] = None,
-        limit: int = 100
-    ) -> List[ExtractedEntityNode]:
-        """查询相关实体
-        
-        查询与指定实体有关系的其他实体。
-        
-        Args:
-            entity_id: 实体ID
-            relation_type: 可选的关系类型过滤
-            limit: 返回结果的最大数量
-            
-        Returns:
-            List[ExtractedEntityNode]: 相关实体列表
-        """
-        if relation_type:
-            query = """
-            MATCH (e1:ExtractedEntity {id: $entity_id})-[r:RELATES_TO {relation_type: $relation_type}]->(e2:ExtractedEntity)
-            RETURN e2 as n
-            LIMIT $limit
-            """
-            results = await self.connector.execute_query(
-                query,
-                entity_id=entity_id,
-                relation_type=relation_type,
-                limit=limit
-            )
-        else:
-            query = """
-            MATCH (e1:ExtractedEntity {id: $entity_id})-[r:RELATES_TO]->(e2:ExtractedEntity)
-            RETURN e2 as n
-            LIMIT $limit
-            """
-            results = await self.connector.execute_query(
-                query,
-                entity_id=entity_id,
-                limit=limit
-            )
-        
-        return [self._map_to_entity(r) for r in results]
-    
-    async def search_by_embedding(
-        self,
-        embedding: List[float],
-        group_id: Optional[str] = None,
-        limit: int = 10,
-        min_score: float = 0.7
-    ) -> List[Dict]:
-        """基于向量相似度搜索实体
-        
-        使用余弦相似度计算查询向量与实体名称向量的相似度。
-        
-        Args:
-            embedding: 查询向量
-            group_id: 可选的组ID过滤
-            limit: 返回结果的最大数量
-            min_score: 最小相似度分数阈值
-            
-        Returns:
-            List[Dict]: 包含实体和相似度分数的字典列表
-                每个字典包含: entity (ExtractedEntityNode), score (float)
-        """
-        where_clause = "n.name_embedding IS NOT NULL"
-        if group_id:
-            where_clause += " AND n.group_id = $group_id"
-        
-        query = f"""
-        MATCH (n:{self.node_label})
-        WHERE {where_clause}
-        WITH n, gds.similarity.cosine(n.name_embedding, $embedding) AS score
-        WHERE score > $min_score
-        RETURN n, score
-        ORDER BY score DESC
-        LIMIT $limit
-        """
-        
-        params = {
-            "embedding": embedding,
-            "min_score": min_score,
-            "limit": limit
-        }
-        if group_id:
-            params["group_id"] = group_id
-        
-        results = await self.connector.execute_query(query, **params)
-        
-        return [
-            {
-                "entity": self._map_to_entity(r),
-                "score": r.get("score", 0.0)
-            }
-            for r in results
-        ]
-    
-    async def find_by_statement_id(self, statement_id: str) -> List[ExtractedEntityNode]:
-        """根据陈述句ID查询实体
-        
-        查询从指定陈述句中提取的所有实体。
-        
-        Args:
-            statement_id: 陈述句ID
-            
-        Returns:
-            List[ExtractedEntityNode]: 实体列表
-        """
-        return await self.find({"statement_id": statement_id})
-    
-    async def find_strong_entities(
-        self,
-        group_id: str,
-        limit: int = 100
-    ) -> List[ExtractedEntityNode]:
-        """查询强连接的实体
-        
-        Args:
-            group_id: 组ID
-            limit: 返回结果的最大数量
-            
-        Returns:
-            List[ExtractedEntityNode]: 强连接的实体列表
-        """
-        return await self.find(
-            {"group_id": group_id, "connect_strength": "Strong"},
-            limit=limit
-        )
-    
-    async def get_entity_count_by_type(self, group_id: str) -> Dict[str, int]:
-        """统计各类型实体的数量
-        
-        Args:
-            group_id: 组ID
-            
-        Returns:
-            Dict[str, int]: 实体类型到数量的映射
-        """
-        query = """
-        MATCH (n:ExtractedEntity {group_id: $group_id})
-        RETURN n.entity_type as entity_type, count(n) as count
-        ORDER BY count DESC
-        """
-        results = await self.connector.execute_query(query, group_id=group_id)
-        return {r["entity_type"]: r["count"] for r in results}
-    
-    async def find_by_config_id(
-        self,
-        config_id: str,
-        limit: int = 100
-    ) -> List[ExtractedEntityNode]:
-        """根据config_id查询实体
-        
-        Args:
-            config_id: 配置ID
-            limit: 返回结果的最大数量
-            
-        Returns:
-            List[ExtractedEntityNode]: 实体列表
-        """
-        return await self.find({"config_id": config_id}, limit=limit)
-    
-    async def search_by_embedding_with_config(
-        self,
-        embedding: List[float],
-        config_id: Optional[str] = None,
-        group_id: Optional[str] = None,
-        limit: int = 10,
-        min_score: float = 0.7
-    ) -> List[Dict]:
-        """基于向量相似度搜索实体,可选择按config_id过滤
-        
-        使用余弦相似度计算查询向量与实体名称向量的相似度。
-        支持按config_id过滤结果,确保只返回使用特定配置处理的实体。
-        
-        Args:
-            embedding: 查询向量
-            config_id: 可选的配置ID过滤
-            group_id: 可选的组ID过滤
-            limit: 返回结果的最大数量
-            min_score: 最小相似度分数阈值
-            
-        Returns:
-            List[Dict]: 包含实体和相似度分数的字典列表
-                每个字典包含: entity (ExtractedEntityNode), score (float)
-        """
-        # 构建查询条件
-        where_clauses = ["n.name_embedding IS NOT NULL"]
-        params = {
-            "embedding": embedding,
-            "min_score": min_score,
-            "limit": limit
-        }
-        
-        if config_id:
-            where_clauses.append("n.config_id = $config_id")
-            params["config_id"] = config_id
-        
-        if group_id:
-            where_clauses.append("n.group_id = $group_id")
-            params["group_id"] = group_id
-        
-        where_str = " AND ".join(where_clauses)
-        
-        query = f"""
-        MATCH (n:{self.node_label})
-        WHERE {where_str}
-        WITH n, gds.similarity.cosine(n.name_embedding, $embedding) AS score
-        WHERE score > $min_score
-        RETURN n, score
-        ORDER BY score DESC
-        LIMIT $limit
-        """
-        
-        results = await self.connector.execute_query(query, **params)
-        
-        return [
-            {
-                "entity": self._map_to_entity(r),
-                "score": r.get("score", 0.0)
-            }
-            for r in results
-        ]
+