knowledge/service/kg_node_service.py

from sqlalchemy.orm import Session
from typing import Optional
from model.kg_node import KGNode
from db.session import get_db
import logging
from sqlalchemy.exc import IntegrityError

from utils import vectorizer
from utils.vectorizer import Vectorizer
from sqlalchemy import func
from service.kg_prop_service import KGPropService
from service.kg_edge_service import KGEdgeService
from cachetools import TTLCache
from cachetools.keys import hashkey
logger = logging.getLogger(__name__)
DISTANCE_THRESHOLD = 0.65
class KGNodeService:
    def __init__(self, db: Session):
        self.db = db

    _cache = TTLCache(maxsize=100000, ttl=60*60*24*30)

    def search_title_index(self, index: str, keywrod: str,category: str, top_k: int = 3,distance: float = 0.3) -> Optional[int]:
        cache_key = f"search_title_index_{index}:{keywrod}:{category}:{top_k}:{distance}"
        if cache_key in self._cache:
            return self._cache[cache_key]

        query_embedding = Vectorizer.get_embedding(keywrod)
        db = next(get_db())
        # 执行向量搜索
        results = (
            db.query(
                KGNode.id,
                KGNode.name,
                KGNode.category,
                KGNode.embedding.l2_distance(query_embedding).label('distance')
            )
            .filter(KGNode.status == 0)
            .filter(KGNode.category == category)
            #todo 是否能提高性能 改成余弦算法
            .filter(KGNode.embedding.l2_distance(query_embedding) <= distance)
            .order_by('distance').limit(top_k).all()
        )
        results = [
            {
                "id": node.id,
               "title": node.name,
               "text": node.category,
               "score": 2.0-node.distance
            }
                for node in results
            ]

        self._cache[cache_key] = results
        return results

    def paginated_search(self, search_params: dict) -> dict:
        load_props = search_params.get('load_props', False)
        prop_service = KGPropService(self.db)
        edge_service = KGEdgeService(self.db)
        keyword = search_params.get('keyword', '')
        category = search_params.get('category', None)
        page_no = search_params.get('pageNo', 1)
        distance = search_params.get('distance',DISTANCE_THRESHOLD)
        limit = search_params.get('limit', 10)

        if page_no < 1:
            page_no = 1
        if limit < 1:
            limit = 10

        embedding = Vectorizer.get_embedding(keyword)
        offset = (page_no - 1) * limit

        try:
            # 构建基础查询条件
            base_query = self.db.query(func.count(KGNode.id)).filter(
                KGNode.status == 0,
                KGNode.embedding.l2_distance(embedding) < distance
            )
            # 如果有category，则添加额外过滤条件
            if category:
                base_query = base_query.filter(KGNode.category == category)
            # 如果有knowledge_ids，则添加额外过滤条件
            if search_params.get('knowledge_ids'):
                total_count = base_query.filter(
                    KGNode.version.in_(search_params['knowledge_ids'])
                ).scalar()
            else:
                total_count = base_query.scalar()

            query = self.db.query(
                KGNode.id,
                KGNode.name,
                KGNode.category,
                KGNode.embedding.l2_distance(embedding).label('distance')
            )            
            query = query.filter(KGNode.status == 0)
            #category有值时，过滤掉category不等于category的节点
            if category:
                query = query.filter(KGNode.category == category)
            if search_params.get('knowledge_ids'):
                query = query.filter(KGNode.version.in_(search_params['knowledge_ids']))
            query = query.filter(KGNode.embedding.l2_distance(embedding) < distance)
            results = query.order_by('distance').offset(offset).limit(limit).all()
            #将results相同distance的category=疾病的放在前面
            results = sorted(results, key=lambda x: (x.distance, not x.category == '疾病'))
            return {
                'records': [{
                    'id': r.id,
                    'name': r.name,
                    'category': r.category,
                    'props': prop_service.get_props_by_ref_id(r.id) if load_props else [],
                    #'edges':edge_service.get_edges_by_nodes(r.id, r.id,False) if load_props else [],
                    'distance': round(r.distance, 3)
                } for r in results],
                'pagination': {
                    'total': total_count,
                    'pageNo': page_no,
                    'limit': limit,
                    'totalPages': (total_count + limit - 1) // limit
                }
            
            }
        except Exception as e:
            logger.error(f"分页查询失败: {str(e)}")
            raise e

    def create_node(self, node_data: dict):
        try:
            existing = self.db.query(KGNode).filter(
                KGNode.name == node_data['name'],
                KGNode.category == node_data['category'],
                KGNode.version == node_data.get('version')
            ).first()
            
            if existing:
                raise ValueError("Node already exists")

            new_node = KGNode(**node_data)
            self.db.add(new_node)
            self.db.commit()
            return new_node

        except IntegrityError as e:
            self.db.rollback()
            logger.error(f"创建节点失败: {str(e)}")
            raise ValueError("Database integrity error")

    def get_node(self, node_id: int):
        if node_id is None:
            raise ValueError("Node ID is required")
        
        cache_key = f"get_node_{node_id}"
        if cache_key in self._cache:
            return self._cache[cache_key]
        
        node = self.db.query(KGNode).filter(KGNode.id == node_id, KGNode.status == 0).first()
        
        if not node:
            raise ValueError("Node not found")
        
        node_data = {
            'id': node.id,
            'name': node.name,
            'category': node.category,
            'version': node.version
        }
        self._cache[cache_key] = node_data
        return node_data

    def update_node(self, node_id: int, update_data: dict):
        node = self.db.query(KGNode).get(node_id)
        if not node:
            raise ValueError("Node not found")

        try:
            for key, value in update_data.items():
                setattr(node, key, value)
            self.db.commit()
            return node
        except Exception as e:
            self.db.rollback()
            logger.error(f"更新节点失败: {str(e)}")
            raise ValueError("Update failed")

    def delete_node(self, node_id: int):
        node = self.db.query(KGNode).get(node_id)
        if not node:
            raise ValueError("Node not found")

        try:
            self.db.delete(node)
            self.db.commit()
            return None
        except Exception as e:
            self.db.rollback()
            logger.error(f"删除节点失败: {str(e)}")
            raise ValueError("Delete failed")

    def batch_process_er_nodes(self):
        batch_size = 200
        offset = 0

        while True:
            try:
                #下面的查询语句，增加根据id排序，防止并发问题
                nodes = self.db.query(KGNode).filter(
                    #KGNode.version == 'ER',
                    KGNode.embedding == None
                ).order_by(KGNode.id).offset(offset).limit(batch_size).all()

                if not nodes:
                    break

                updated_nodes = []
                for node in nodes:
                    if not node.embedding:
                        embedding = Vectorizer.get_embedding(node.name)
                        node.embedding = embedding
                        updated_nodes.append(node)
                if updated_nodes:
                    self.db.commit()

                offset += batch_size
            except Exception as e:
                self.db.rollback()
                print(f"批量处理ER节点失败: {str(e)}")
                raise ValueError("Batch process failed")