services/article-parser/src/controllers/entity_extractor.py - angular - boris00043/semantic-papers

services/article-parser/src/controllers/entity_extractor.py
111 строк4 KB
Бредихин Борис Андреевич
fix
02 авг 2025, 16:21
02 авг 2025, 16:21aa93d6e
from arango.result import Result
from ai.graphrag import GraphRAG
from db.arango import ArangoDatabase


class EntityExtractor:
    def __init__(
        self, articles_db: ArangoDatabase, entities_db: ArangoDatabase
    ) -> None:
        self._articles_db = articles_db
        self._entities_db = entities_db

        self.entity = entities_db.get_vertex_collection("Entity")
        self.semantic_link = entities_db.get_edge_collection(
            "semantic_link", ["Entity"], ["Entity"]
        )
        self.summary_of = articles_db.get_edge_collection(
            "summary_of", ["Summary"], ["ArticlePart", "Article"]
        )
        self.part_of = articles_db.get_edge_collection(
            "part_of", ["ArticlePart"], ["Article"]
        )

    def load_entities(self, part: Result):
        semantic_links = self.semantic_link.find(
            {
                "article_part": part["_id"],
            }
        )
        ids = set()
        ents = {"entities": [], "relationships": []}
        for sl in semantic_links:
            e1 = self.entity.find({"_id": sl["_to"]}).pop()
            e2 = self.entity.find({"_id": sl["_from"]}).pop()
            link_props = sl.copy()
            link_props.pop("_key")
            link_props["id"] = link_props["_id"]
            for e in [e1, e2]:
                if e["_id"] not in ids:
                    ents["entities"].append(
                        {"id": e["_id"], "label": e["node_type"], "properties": e}
                    )
                    ids.add(e["_id"])
            ents["relationships"].append(
                {
                    "type": sl["type"],
                    "start_node_id": sl["_from"],
                    "end_node_id": sl["_to"],
                    "properties": link_props,
                }
            )
        return ents

    def extract_and_save(self, part: Result, graphrag: GraphRAG):
        ents = graphrag.to_graph(part["content"])
        if len(ents) != 2:
            return {}

        part_id: str = part["_id"]
        if part_id.startswith("Summary"):
            part_id = self.summary_of.find({"_from": part_id}).pop()["_to"]

        article_id = (
            part_id
            if part_id == "Article"
            else (self.part_of.find({"_from": part_id}).pop()["_to"])
        )

        node_id_to_db_id: dict[str, str] = {}

        for node in ents["entities"]:
            try:
                if not str.isdigit(node["id"][0]):
                    node["properties"]["id"] = node["id"]
            except TypeError:
                continue

            try:
                node_obj = self.entity.find(
                    node["properties"] | {"node_type": node["label"]}
                ).pop()
            except:
                node_obj = self.entity.insert(
                    node["properties"] | {"node_type": node["label"]}
                )
            node_id_to_db_id[node["id"]] = node_obj["_id"]

        link_kw = "relationships"
        if not link_kw in ents.keys():
            link_kw = "edges"
        for link in ents[link_kw]:
            try:
                self.semantic_link.insert(
                    {
                        "_from": node_id_to_db_id[link["start_node_id"]],
                        "_to": node_id_to_db_id[link["end_node_id"]],
                        "type": (
                            link["type"]
                            if "type" in link.keys()
                            else link["relationship"]
                        ),
                        "article_part": part["_id"],
                        "article": article_id,
                    }
                    | link.get("properties", {})
                )
            except KeyError:
                pass

        self._articles_db.db.update_document(part | {"entities_extracted": True})
        return ents