ForcePilot/backend/package/yuxi/knowledge/graphs/milvus_graph_service.py

from __future__ import annotations

import asyncio
import json
from typing import Any

from yuxi.knowledge.graphs.extractors import GraphExtractor, GraphExtractorFactory, normalize_extraction_result
from yuxi.knowledge.graphs.graph_utils import (
    build_graph_payload,
    compute_entity_id,
    compute_triple_id,
    cypher_merge_chunk,
    cypher_merge_entity_mention,
    cypher_merge_relation,
    normalize_entity_name,
)
from yuxi.knowledge.graphs.milvus_graph_vector_store import MilvusGraphVectorStore
from yuxi.knowledge.graphs.neo4j_utils import (
    Neo4jConnectionManager,
    get_shared_neo4j_connection,
    neo4j_read,
    neo4j_write,
    safe_neo4j_label,
)
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
from yuxi.repositories.knowledge_chunk_repository import KnowledgeChunkRepository
from yuxi.repositories.knowledge_graph_repository import KnowledgeGraphRepository
from yuxi.utils import logger
from yuxi.utils.datetime_utils import utc_isoformat

GRAPH_CONFIG_KEY = "graph_build_config"
GRAPH_TASK_TYPE = "knowledge_graph_index"


class MilvusGraphService:
    def __init__(
        self,
        *,
        db_id: str | None = None,
        kb_repo: KnowledgeBaseRepository | None = None,
        chunk_repo: KnowledgeChunkRepository | None = None,
        graph_repo: KnowledgeGraphRepository | None = None,
        graph_vector_store: MilvusGraphVectorStore | None = None,
        neo4j_connection: Neo4jConnectionManager | None = None,
    ):
        self.db_id = db_id
        self.kb_repo = kb_repo or KnowledgeBaseRepository()
        self.chunk_repo = chunk_repo or KnowledgeChunkRepository()
        self.graph_repo = graph_repo or KnowledgeGraphRepository()
        self._graph_vector_store = graph_vector_store
        self._connection = neo4j_connection

    @property
    def connection(self) -> Neo4jConnectionManager:
        if self._connection is None:
            self._connection = get_shared_neo4j_connection()
        return self._connection

    @property
    def graph_vector_store(self) -> MilvusGraphVectorStore:
        if self._graph_vector_store is None:
            self._graph_vector_store = MilvusGraphVectorStore()
        return self._graph_vector_store

    @property
    def driver(self):
        return self.connection.driver

    async def get_status(self, db_id: str, *, tasker: Any = None) -> dict[str, Any]:
        kb = await self._get_milvus_kb(db_id)
        params = dict(kb.additional_params or {})
        config = params.get(GRAPH_CONFIG_KEY) or {}
        total_chunks, pending_chunks, indexed_chunks, graph_counts = await asyncio.gather(
            self.chunk_repo.count_by_db_id(db_id),
            self.chunk_repo.count_graph_pending_by_db_id(db_id),
            self.chunk_repo.count_graph_indexed_by_db_id(db_id),
            self.graph_repo.count_by_db_id(db_id),
        )
        entity_count, relationship_count = graph_counts

        build_task_status = None
        build_task_progress = 0
        if tasker is not None:
            active_task = await tasker.find_task_by_payload(
                task_type=GRAPH_TASK_TYPE,
                payload_match={"db_id": db_id},
                statuses={"pending", "running"},
            )
            if active_task:
                build_task_status = active_task.status
                build_task_progress = round(active_task.progress)
            else:
                failed_task = await tasker.find_task_by_payload(
                    task_type=GRAPH_TASK_TYPE,
                    payload_match={"db_id": db_id},
                    statuses={"failed", "cancelled"},
                )
                if failed_task:
                    build_task_status = "failed"
                    build_task_progress = 0

        return {
            "db_id": db_id,
            "kb_type": kb.kb_type,
            "configured": bool(config),
            "locked": bool(config.get("locked")),
            "config": self._public_config(config),
            "total_chunks": total_chunks,
            "pending_chunks": pending_chunks,
            "indexed_chunks": indexed_chunks,
            "entity_count": entity_count,
            "relationship_count": relationship_count,
            "build_task_status": build_task_status,
            "build_task_progress": build_task_progress,
        }

    async def configure(
        self,
        db_id: str,
        extractor_type: str,
        extractor_options: dict[str, Any],
        created_by: str,
    ) -> dict:
        kb = await self._get_milvus_kb(db_id)
        additional_params = dict(kb.additional_params or {})
        existing_config = additional_params.get(GRAPH_CONFIG_KEY) or {}
        normalized_extractor_type = (extractor_type or "").lower()
        if existing_config.get("locked"):
            existing_extractor_type = (existing_config.get("extractor_type") or "").lower()
            if normalized_extractor_type != existing_extractor_type:
                raise ValueError("图谱抽取器类型已锁定，只能修改模型、Schema 等抽取参数")

        extractor_options = extractor_options or {}
        if normalized_extractor_type == "llm" and extractor_options.get("prompt"):
            raise ValueError("LLM 图谱抽取器不支持自定义完整 Prompt，请使用 schema 配置抽取约束")
        GraphExtractorFactory.create(normalized_extractor_type, extractor_options)
        config = {
            "locked": True,
            "extractor_type": normalized_extractor_type,
            "extractor_options": extractor_options or {},
            "created_at": existing_config.get("created_at") or utc_isoformat(),
            "created_by": existing_config.get("created_by") or created_by,
        }
        if existing_config.get("locked"):
            config["updated_at"] = utc_isoformat()
            config["updated_by"] = created_by
        additional_params[GRAPH_CONFIG_KEY] = config
        await self.kb_repo.update(db_id, {"additional_params": additional_params})
        return config

    async def build_pending_chunks(self, db_id: str, *, batch_size: int, context=None) -> dict[str, Any]:
        kb = await self._get_milvus_kb(db_id)
        config = self._get_locked_config(kb.additional_params or {})
        extractor_options = self._runtime_extractor_options(config)
        extractor = GraphExtractorFactory.create(config["extractor_type"], extractor_options)
        worker_count = self._get_worker_count(config)
        total_pending = await self.chunk_repo.count_graph_pending_by_db_id(db_id)
        processed = 0
        failed = 0
        failed_chunk_ids: set[str] = set()
        write_lock = asyncio.Lock()

        while True:
            if context is not None:
                await context.raise_if_cancelled()
            chunks = await self.chunk_repo.list_graph_pending_by_db_id(db_id, batch_size)
            unprocessed = [c for c in chunks if c.chunk_id not in failed_chunk_ids]
            if not unprocessed:
                break

            queue: asyncio.Queue[Any] = asyncio.Queue()
            for chunk in unprocessed:
                queue.put_nowait(chunk)

            async def worker() -> None:
                nonlocal processed, failed
                while True:
                    if context is not None:
                        await context.raise_if_cancelled()
                    try:
                        chunk = queue.get_nowait()
                    except asyncio.QueueEmpty:
                        return
                    try:
                        extraction_result = await self._get_chunk_extraction_result(db_id, chunk, extractor)
                        async with write_lock:
                            entities, triples = await asyncio.to_thread(
                                self.write_chunk_graph,
                                db_id,
                                chunk,
                                extraction_result,
                            )
                            await self.graph_repo.upsert_chunk_graph(
                                db_id=db_id,
                                file_id=chunk.file_id,
                                chunk_id=chunk.chunk_id,
                                entities=entities,
                                triples=triples,
                            )
                            await self.graph_vector_store.insert_missing_graph_records(
                                db_id=db_id,
                                embedding_model_spec=kb.embedding_model_spec,
                                entities=entities,
                                triples=triples,
                            )
                            await self.chunk_repo.mark_graph_indexed(
                                chunk.chunk_id,
                                ent_ids=[entity["entity_id"] for entity in entities],
                            )
                        processed += 1
                    except Exception as exc:
                        logger.error(f"Chunk 图谱构建失败 chunk_id={chunk.chunk_id}: {exc}")
                        failed_chunk_ids.add(chunk.chunk_id)
                        failed += 1
                    finally:
                        queue.task_done()

                    if context is not None:
                        completed = processed + failed
                        progress = 5.0 + min(90.0, completed / max(total_pending, 1) * 90.0)
                        await context.set_progress(progress, f"图谱构建 {completed}/{total_pending}，失败 {failed}")

            workers = [asyncio.create_task(worker()) for _ in range(min(worker_count, len(unprocessed)))]
            try:
                await asyncio.gather(*workers)
            except Exception:
                for task in workers:
                    task.cancel()
                await asyncio.gather(*workers, return_exceptions=True)
                raise

        remaining = await self.chunk_repo.count_graph_pending_by_db_id(db_id)
        return {"db_id": db_id, "success": processed, "failed": failed, "remaining": remaining}

    @staticmethod
    def _get_worker_count(config: dict[str, Any]) -> int:
        if (config.get("extractor_type") or "").lower() != "llm":
            return 1
        try:
            worker_count = int((config.get("extractor_options") or {}).get("concurrency_count") or 1)
        except (TypeError, ValueError):
            return 1
        return max(1, min(worker_count, 1000))

    @staticmethod
    def _runtime_extractor_options(config: dict[str, Any]) -> dict[str, Any]:
        options = dict(config.get("extractor_options") or {})
        options.pop("prompt", None)
        return options

    async def _get_chunk_extraction_result(self, db_id: str, chunk, extractor: GraphExtractor) -> dict[str, Any]:
        extractor_type = extractor.extractor_type
        if chunk.extraction_result:
            return normalize_extraction_result(chunk.extraction_result, extractor_type)

        extraction_result = await extractor.extract(
            chunk.content,
            chunk_metadata={
                "db_id": db_id,
                "chunk_id": chunk.chunk_id,
                "file_id": chunk.file_id,
                "chunk_index": chunk.chunk_index,
            },
        )
        normalized_result = normalize_extraction_result(extraction_result, extractor_type)
        await self.chunk_repo.update_extraction_result(chunk.chunk_id, normalized_result)
        return normalized_result

    def write_chunk_graph(
        self,
        db_id: str,
        chunk,
        normalized_result: dict[str, Any],
    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
        """将单个 chunk 的抽取结果写入 Neo4j。"""
        label = safe_neo4j_label(db_id)
        graph_payload = build_graph_payload(normalized_result)
        relation_extractor_type = graph_payload["metadata"].get("extractor_type", "unknown")
        entities = graph_payload["entities"]
        relations = graph_payload["relations"]
        entity_by_id = {entity["id"]: entity for entity in entities}
        entity_records = self._build_entity_records(db_id, entities)
        entity_record_by_local_id = {
            entity["id"]: record for entity, record in zip(entities, entity_records, strict=True)
        }
        triple_records = self._build_triple_records(db_id, relations, entity_record_by_local_id, graph_payload)
        content_preview = (chunk.content or "")[:300]

        # 预构建 Cypher 模板（同一 chunk 内复用）
        merge_chunk_cypher = cypher_merge_chunk(label)
        merge_entity_cypher = cypher_merge_entity_mention(label)
        merge_relation_cypher = cypher_merge_relation(label)

        def query(tx):
            # 1. MERGE Chunk 节点
            tx.run(
                merge_chunk_cypher,
                chunk_id=chunk.chunk_id,
                file_id=chunk.file_id,
                db_id=db_id,
                chunk_index=chunk.chunk_index,
                content_preview=content_preview,
                start_char_pos=chunk.start_char_pos,
                end_char_pos=chunk.end_char_pos,
            )

            # 2. MERGE Entity 节点 + Chunk→Entity (MENTIONS)
            for entity in entities:
                entity_record = entity_record_by_local_id[entity["id"]]
                tx.run(
                    merge_entity_cypher,
                    chunk_id=chunk.chunk_id,
                    file_id=chunk.file_id,
                    db_id=db_id,
                    entity_id=entity_record["entity_id"],
                    normalized_name=normalize_entity_name(entity["text"]),
                    entity_label=entity.get("label") or "Entity",
                    name=entity["text"],
                    attributes=json.dumps(entity.get("attributes") or [], ensure_ascii=False),
                )

            # 3. MERGE Entity→Entity (RELATION) 边
            for relation in relations:
                source = entity_by_id[relation["source"]]
                target = entity_by_id[relation["target"]]
                source_record = entity_record_by_local_id[relation["source"]]
                target_record = entity_record_by_local_id[relation["target"]]
                relation_type = relation.get("label") or "RELATED_TO"
                triple_id = compute_triple_id(
                    db_id,
                    source_record["normalized_name"],
                    source_record["label"],
                    relation_type,
                    target_record["normalized_name"],
                    target_record["label"],
                )
                tx.run(
                    merge_relation_cypher,
                    db_id=db_id,
                    chunk_id=chunk.chunk_id,
                    file_id=chunk.file_id,
                    source_name=normalize_entity_name(source["text"]),
                    source_label=source.get("label") or "Entity",
                    target_name=normalize_entity_name(target["text"]),
                    target_label=target.get("label") or "Entity",
                    relation_type=relation_type,
                    triple_id=triple_id,
                    text=relation["text"],
                    extractor_type=relation_extractor_type,
                )

        neo4j_write(self.driver, query)
        return entity_records, triple_records

    def _build_entity_records(self, db_id: str, entities: list[dict[str, Any]]) -> list[dict[str, Any]]:
        records = []
        for entity in entities:
            label = entity.get("label") or "Entity"
            normalized_name = normalize_entity_name(entity["text"])
            entity_id = compute_entity_id(db_id, normalized_name, label)
            records.append(
                {
                    "entity_id": entity_id,
                    "db_id": db_id,
                    "normalized_name": normalized_name,
                    "label": label,
                    "name": entity["text"],
                    "attributes": entity.get("attributes") or [],
                    "content": normalized_name,
                }
            )
        return records

    def _build_triple_records(
        self,
        db_id: str,
        relations: list[dict[str, Any]],
        entity_record_by_local_id: dict[str, dict[str, Any]],
        graph_payload: dict[str, Any],
    ) -> list[dict[str, Any]]:
        records = []
        seen_triple_ids: set[str] = set()
        extractor_type = graph_payload["metadata"].get("extractor_type", "unknown")
        for relation in relations:
            source_record = entity_record_by_local_id[relation["source"]]
            target_record = entity_record_by_local_id[relation["target"]]
            relation_type = relation.get("label") or "RELATED_TO"
            triple_id = compute_triple_id(
                db_id,
                source_record["normalized_name"],
                source_record["label"],
                relation_type,
                target_record["normalized_name"],
                target_record["label"],
            )
            if triple_id in seen_triple_ids:
                continue
            seen_triple_ids.add(triple_id)
            content = f"{source_record['normalized_name']} → {relation_type} → {target_record['normalized_name']}"
            records.append(
                {
                    "triple_id": triple_id,
                    "db_id": db_id,
                    "source_entity_id": source_record["entity_id"],
                    "target_entity_id": target_record["entity_id"],
                    "relation_type": relation_type,
                    "content": content,
                    "text": relation["text"],
                    "extractor_type": extractor_type,
                }
            )
        return records

    async def reset(self, db_id: str, *, clear_extraction_result: bool, clear_config: bool) -> dict[str, Any]:
        kb = await self._get_milvus_kb(db_id)
        await asyncio.to_thread(self.delete_graph, db_id)
        await self.graph_repo.delete_by_db_id(db_id)
        reset_chunks = await self.chunk_repo.reset_graph_state_by_db_id(db_id, clear_extraction_result)
        if clear_config:
            additional_params = dict(kb.additional_params or {})
            additional_params.pop(GRAPH_CONFIG_KEY, None)
            await self.kb_repo.update(db_id, {"additional_params": additional_params})
        return {
            "message": "图谱构建状态已重置",
            "status": "success",
            "reset_chunks": reset_chunks,
            "clear_extraction_result": clear_extraction_result,
            "clear_config": clear_config,
        }

    def delete_graph(self, db_id: str) -> None:
        label = safe_neo4j_label(db_id)

        def query(tx):
            tx.run(f"MATCH (n:MilvusKB:`{label}`) DETACH DELETE n")

        neo4j_write(self.driver, query)
        self.graph_vector_store.drop_graph_collections(db_id)

    async def delete_file_graph(self, db_id: str, file_id: str) -> None:
        orphan_entity_ids, orphan_triple_ids = await self.graph_repo.delete_file_references(file_id)
        await self.graph_vector_store.delete_graph_records(
            db_id,
            entity_ids=orphan_entity_ids,
            triple_ids=orphan_triple_ids,
        )
        await asyncio.to_thread(self._delete_file_graph_from_neo4j, db_id, file_id)

    def _delete_file_graph_from_neo4j(self, db_id: str, file_id: str) -> None:
        label = safe_neo4j_label(db_id)

        def query(tx):
            tx.run(
                f"""
                MATCH (:Chunk:MilvusKB:`{label}`)-[m:MENTIONS {{db_id: $db_id, file_id: $file_id}}]->
                    (:Entity:MilvusKB:`{label}`)
                DELETE m
                """,
                db_id=db_id,
                file_id=file_id,
            )
            tx.run(
                f"""
                MATCH (:Entity:MilvusKB:`{label}`)-[r:RELATION {{db_id: $db_id, file_id: $file_id}}]->
                    (:Entity:MilvusKB:`{label}`)
                DELETE r
                """,
                db_id=db_id,
                file_id=file_id,
            )
            tx.run(
                f"""
                MATCH (c:Chunk:MilvusKB:`{label}` {{db_id: $db_id, file_id: $file_id}})
                DETACH DELETE c
                """,
                db_id=db_id,
                file_id=file_id,
            )
            tx.run(
                f"""
                MATCH (e:Entity:MilvusKB:`{label}` {{db_id: $db_id}})
                WHERE NOT ()-[:MENTIONS]->(e)
                DETACH DELETE e
                """,
                db_id=db_id,
            )

        neo4j_write(self.driver, query)

    async def query_nodes(
        self,
        db_id: str | None = None,
        *,
        keyword: str = "",
        max_depth: int = 1,
        max_nodes: int = 50,
        exclude_chunk: bool = False,
    ) -> dict[str, Any]:
        effective_db_id = db_id or self.db_id
        if not effective_db_id:
            return {"nodes": [], "edges": []}

        label = safe_neo4j_label(effective_db_id)
        limit = max_nodes
        try:
            with self.driver.session() as session:
                result = session.run(
                    self._build_query(label, keyword, limit, max_depth, exclude_chunk),
                    keyword=keyword,
                    limit=limit,
                )
                return self._process_query_result(result, limit, effective_db_id, exclude_chunk)
        except Exception as e:
            logger.error(f"Milvus graph query failed: {e}")
            return {"nodes": [], "edges": []}

    async def query_seed_subgraph(
        self,
        db_id: str,
        *,
        entity_ids: list[str],
        max_nodes: int,
    ) -> dict[str, Any]:
        if not entity_ids:
            return {"nodes": [], "edges": []}
        label = safe_neo4j_label(db_id)
        cypher = f"""
        MATCH (seed:Entity:MilvusKB:`{label}`)
        WHERE seed.entity_id IN $entity_ids
        MATCH p = (seed)-[*1..2]-(n:MilvusKB:`{label}`)
        WITH p LIMIT $path_limit
        WITH collect(p) AS paths
        UNWIND paths AS node_path
        UNWIND nodes(node_path) AS node
        WITH paths, collect(DISTINCT node) AS graph_nodes
        UNWIND paths AS rel_path
        UNWIND relationships(rel_path) AS rel
        RETURN graph_nodes AS nodes, collect(DISTINCT rel) AS edges
        """
        try:
            with self.driver.session() as session:
                record = session.run(
                    cypher,
                    entity_ids=list(dict.fromkeys(entity_ids)),
                    path_limit=max(max_nodes, 1) * 4,
                ).single()
                if not record:
                    return {"nodes": [], "edges": []}
                return self._process_subgraph_record(record, max_nodes, db_id)
        except Exception as e:
            logger.error(f"Milvus seed subgraph query failed: {e}")
            return {"nodes": [], "edges": []}

    async def get_labels(self, db_id: str | None = None) -> list[str]:
        effective_db_id = db_id or self.db_id
        if not effective_db_id:
            return []
        label = safe_neo4j_label(effective_db_id)

        cypher = f"""
        MATCH (n:MilvusKB:`{label}`)
        UNWIND labels(n) AS node_label
        WITH DISTINCT node_label
        WHERE node_label <> 'MilvusKB' AND node_label <> $db_id
        RETURN node_label
        ORDER BY node_label
        """
        try:
            records = neo4j_read(self.driver, cypher, db_id=effective_db_id)
            return [record["node_label"] for record in records]
        except Exception as e:
            logger.error(f"Failed to get Milvus graph labels: {e}")
            return []

    async def get_stats(self, db_id: str | None = None) -> dict[str, Any]:
        effective_db_id = db_id or self.db_id
        if not effective_db_id:
            return {"total_nodes": 0, "total_edges": 0, "entity_types": []}
        label = safe_neo4j_label(effective_db_id)

        stats_cypher = f"""
        MATCH (n:MilvusKB:`{label}`)
        WITH count(n) AS node_count
        OPTIONAL MATCH (:MilvusKB:`{label}`)-[r]->(:MilvusKB:`{label}`)
        RETURN node_count, count(r) AS edge_count
        """
        label_cypher = f"""
        MATCH (n:Entity:MilvusKB:`{label}`)
        WITH n.label AS entity_label, count(*) AS count
        RETURN entity_label, count
        ORDER BY count DESC
        """
        try:
            with self.driver.session() as session:
                stats = session.run(stats_cypher).single()
                label_stats = session.run(label_cypher)
                return {
                    "total_nodes": stats["node_count"] if stats else 0,
                    "total_edges": stats["edge_count"] if stats else 0,
                    "entity_types": [{"type": row["entity_label"], "count": row["count"]} for row in label_stats],
                }
        except Exception as e:
            logger.error(f"Failed to get Milvus graph stats: {e}")
            return {"total_nodes": 0, "total_edges": 0, "entity_types": []}

    async def _get_milvus_kb(self, db_id: str):
        kb = await self.kb_repo.get_by_id(db_id)
        if kb is None:
            raise ValueError(f"知识库 {db_id} 不存在")
        if (kb.kb_type or "").lower() != "milvus":
            raise ValueError("仅 Milvus 知识库支持独立图谱构建")
        return kb

    def _get_locked_config(self, additional_params: dict[str, Any]) -> dict[str, Any]:
        config = additional_params.get(GRAPH_CONFIG_KEY) or {}
        if not config.get("locked"):
            raise ValueError("请先确认并锁定图谱抽取配置")
        if not config.get("extractor_type"):
            raise ValueError("图谱抽取配置缺少 extractor_type")
        return config

    def _public_config(self, config: dict[str, Any]) -> dict[str, Any] | None:
        if not config:
            return None
        return {
            "locked": bool(config.get("locked")),
            "extractor_type": config.get("extractor_type"),
            "extractor_options": self._runtime_extractor_options(config),
            "created_at": config.get("created_at"),
            "created_by": config.get("created_by"),
            "updated_at": config.get("updated_at"),
            "updated_by": config.get("updated_by"),
        }

    @staticmethod
    def _build_where(exclude_chunk: bool, keyword: str) -> str:
        clauses = []
        if exclude_chunk:
            clauses.append("NOT n:Chunk")
        if keyword and keyword != "*":
            clauses.append(
                "(toLower(coalesce(n.name, '')) CONTAINS toLower($keyword)"
                " OR toLower(coalesce(n.content_preview, '')) CONTAINS toLower($keyword)"
                " OR toLower(coalesce(n.chunk_id, '')) CONTAINS toLower($keyword))"
            )
        return "WHERE " + " AND ".join(clauses) if clauses else ""

    def _build_query(self, label: str, keyword: str, limit: int, max_depth: int, exclude_chunk: bool = False) -> str:
        where = self._build_where(exclude_chunk, keyword)
        m_exclude = " WHERE NOT m:Chunk" if exclude_chunk else ""

        if max_depth <= 0:
            return f"""
            MATCH (n:MilvusKB:`{label}`)
            {where}
            RETURN n AS h, null AS r, null AS t
            LIMIT $limit
            """

        return f"""
        MATCH (n:MilvusKB:`{label}`)
        {where}
        WITH n LIMIT $limit
        OPTIONAL MATCH (n)-[r]-(m:MilvusKB:`{label}`){m_exclude}
        RETURN n AS h, r AS r, m AS t
        LIMIT {limit * 10}
        """

    def _process_query_result(self, result, limit: int, db_id: str, exclude_chunk: bool = False) -> dict[str, Any]:
        nodes = []
        edges = []
        node_ids = set()
        edge_ids = set()

        for record in result:
            for key in ("h", "t"):
                raw_node = record.get(key)
                if raw_node is None:
                    continue
                node = self._normalize_node(raw_node, db_id)
                if not node or node["id"] in node_ids:
                    continue
                if exclude_chunk and node.get("type") == "Chunk":
                    continue
                nodes.append(node)
                node_ids.add(node["id"])
            raw_edge = record.get("r")
            if raw_edge is not None:
                edge = self._normalize_edge(raw_edge)
                if edge and edge["id"] not in edge_ids:
                    edges.append(edge)
                    edge_ids.add(edge["id"])
            if len(nodes) >= limit:
                break

        return {"nodes": nodes[:limit], "edges": edges[: limit * 2]}

    def _process_subgraph_record(self, record: Any, limit: int, db_id: str) -> dict[str, Any]:
        nodes = []
        edges = []
        node_ids = set()
        edge_ids = set()

        for raw_node in record.get("nodes") or []:
            node = self._normalize_node(raw_node, db_id)
            if not node or node["id"] in node_ids:
                continue
            nodes.append(node)
            node_ids.add(node["id"])
            if len(nodes) >= limit:
                break

        for raw_edge in record.get("edges") or []:
            edge = self._normalize_edge(raw_edge)
            if not edge or edge["id"] in edge_ids:
                continue
            if edge["source_id"] not in node_ids or edge["target_id"] not in node_ids:
                continue
            edges.append(edge)
            edge_ids.add(edge["id"])

        return {"nodes": nodes, "edges": edges}

    def _normalize_node(self, raw_node: Any, db_id: str | None = None) -> dict[str, Any]:
        if hasattr(raw_node, "element_id"):
            node_id = raw_node.element_id
            labels = list(raw_node.labels)
            properties = dict(raw_node.items())
        elif isinstance(raw_node, dict):
            node_id = raw_node.get("id") or raw_node.get("element_id")
            labels = raw_node.get("labels", [])
            properties = raw_node.get("properties") or {k: v for k, v in raw_node.items() if k not in {"id", "labels"}}
        else:
            return {}

        effective_db_id = db_id or self.db_id
        db_label = properties.get("db_id") or effective_db_id
        filtered_labels = [label for label in labels if label not in {"MilvusKB", db_label}]
        entity_type = "Chunk" if "Chunk" in labels else properties.get("label", "Entity")
        name = properties.get("name") or properties.get("content_preview") or properties.get("chunk_id") or "Unknown"
        return {
            "id": node_id,
            "name": name,
            "original_id": node_id,
            "type": entity_type,
            "labels": filtered_labels,
            "properties": properties,
            "normalized": {
                "name": name,
                "type": entity_type,
                "source": "milvus",
            },
            "graph_type": "milvus",
        }

    def _normalize_edge(self, raw_edge: Any) -> dict[str, Any]:
        if hasattr(raw_edge, "element_id"):
            edge_id = raw_edge.element_id
            edge_type = raw_edge.type
            source_id = raw_edge.start_node.element_id
            target_id = raw_edge.end_node.element_id
            properties = dict(raw_edge.items())
            edge_type = properties.get("type") or edge_type
        elif isinstance(raw_edge, dict):
            edge_id = raw_edge.get("id")
            edge_type = raw_edge.get("type")
            source_id = raw_edge.get("source_id")
            target_id = raw_edge.get("target_id")
            properties = raw_edge.get("properties", {})
        else:
            return {}

        return {
            "id": edge_id,
            "source_id": source_id,
            "target_id": target_id,
            "type": edge_type,
            "properties": properties,
            "normalized": {
                "type": edge_type,
                "direction": "directed",
            },
        }