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

services/article-parser/src/controllers/article.py
263 строки7 KB
boris
Connect article to several themes
16 авг 2025, 10:17
16 авг 2025, 10:17b206b55
import json
import secrets
from typing import Any
import uuid

import pika
from qdrant_client.models import Distance, PointStruct, VectorParams
import tqdm

from ai.summarizer import Summarizer
from common import (
    articles_db,
    get_rabbitmq_connection,
    grobid_adapter,
    llm,
    ollama_client,
    paper_root,
)
from db.qdrant import get_qdrant_client
from grobid import TEIArticle
from models import arango_models as am


def process_article(
    title: str,
    href: str,
    theme: str,
    *,
    embeddings_model: str = "mxbai-embed-large:latest",
):
    try:
        tei = grobid_adapter.get_tei_by_pdf_url(href, paper_root)
    except:
        print(f"Error processing article {title}")
        return
    tei_article = TEIArticle(tei, grobid_adapter.llm, title)
    qdrant_client = get_qdrant_client()
    summarizer = Summarizer(ollama_client)
    connection = get_rabbitmq_connection()
    channel = connection.channel()
    channel.queue_declare(queue="entity_extraction", durable=True)

    if not qdrant_client.collection_exists("Article"):
        qdrant_client.create_collection(
            "Article",
            vectors_config=VectorParams(size=1024, distance=Distance.COSINE),
        )

    try:
        theme_obj: dict[Any, Any] = am.theme_collection.find({"title": theme}).pop()
    except:
        theme_obj = articles_db.graph.insert_vertex(
            "Theme",
            {
                "title": theme,
            },
        )
    if articles_db.get_vertex_collection("Article").find({"url": href}).count() > 0:
        art = articles_db.get_vertex_collection("Article").find({"url": href}).pop()
        if (
            am.theme_to_article.find(
                {"_from": theme_obj["_id"], "_to": art["_id"]}
            ).count()
            == 0
        ):
            am.theme_to_article.insert(
                {
                    "_from": theme_obj["_id"],
                    "_to": art["_id"],
                }
            )
        return
    art = articles_db.graph.insert_vertex(
        "Article",
        {
            "_key": secrets.token_urlsafe(8),
            "title": tei_article.title,
            "url": href,
        },
    )

    am.theme_to_article.insert(
        {
            "_from": theme_obj["_id"],
            "_to": art["_id"],
        }
    )

    for keyword in tei_article.keywords:
        try:
            kw_obj = am.keyword_collection.find({"title": keyword}).pop()
        except:
            kw_obj = am.articles_db.graph.insert_vertex(
                "Keyword",
                {
                    "title": keyword,
                },
            )
        am.article_has_keyword.insert(
            {
                "_from": art["_id"],
                "_to": kw_obj["_id"],
            }
        )
    for author in tei_article.authors:
        try:
            author_obj = am.keyword_collection.find(
                {"forename": author.forename, "surname": author.surname}
            ).pop()
        except:
            author_obj = articles_db.graph.insert_vertex(
                "Author",
                {
                    "forename": author.forename,
                    "surname": author.surname,
                },
            )
        am.author_of.insert(
            {
                "_from": author_obj["_id"],
                "_to": art["_id"],
            }
        )
    previous_part = None
    summaries = []

    for document_part in tqdm.tqdm(tei_article.body):
        part_content = "\n".join(document_part.content)
        part_content_embeddings = ollama_client.embeddings(
            embeddings_model, part_content
        )["embedding"]

        article_part_obj = articles_db.graph.insert_vertex(
            "ArticlePart",
            {"title": document_part.title, "content": part_content},
        )

        try:
            qdrant_client.upsert(
                "Article",
                points=[
                    PointStruct(
                        id=str(uuid.uuid4()),
                        vector=part_content_embeddings,
                        payload={
                            "external_id": article_part_obj["_id"],
                            "article_id": art["_id"],
                        },
                    )
                ],
            )
        except:
            pass
        if previous_part:
            am.previous_part.insert(
                {
                    "_key": secrets.token_urlsafe(8),
                    "_from": previous_part["_id"],
                    "_to": article_part_obj["_id"],
                }
            )
        previous_part = article_part_obj
        am.part_of.insert(
            {
                "_key": secrets.token_urlsafe(8),
                "_from": article_part_obj["_id"],
                "_to": art["_id"],
            }
        )

        content = [
            summarizer.summarize(c)
            for c in tqdm.tqdm(document_part.content, leave=False)
        ]

        summary_obj = am.articles_db.graph.insert_vertex(
            "Summary",
            {
                "content": "\n".join(content),
            },
        )

        for c in content:
            summary_embeddings = ollama_client.embed(embeddings_model, c).embeddings[0]

            try:
                qdrant_client.upsert(
                    "Article",
                    points=[
                        PointStruct(
                            id=str(uuid.uuid4()),
                            vector=summary_embeddings,
                            payload={
                                "external_id": summary_obj["_id"],
                                "article_id": art["_id"],
                            },
                        )
                    ],
                )
            except:
                pass

        summaries.append(
            {
                "title": document_part.title,
                "content": content,
                "id": summary_obj["_id"],
            }
        )

        channel.basic_publish(
            "",
            "entity_extraction",
            json.dumps({"type": "extract_entity", "part_id": summary_obj["_id"]}),
            properties=pika.BasicProperties(delivery_mode=pika.DeliveryMode.Persistent),
        )

        am.summary_of.insert(
            {
                "_key": secrets.token_urlsafe(8),
                "_from": summary_obj["_id"],
                "_to": article_part_obj["_id"],
            }
        )

    short_summary = summarizer.summarize(
        "\n\n".join("\n".join(i["content"]) for i in summaries)
    )
    short_summary_embeddings = ollama_client.embeddings(
        embeddings_model, short_summary
    )["embedding"]
    summary_obj = am.articles_db.graph.insert_vertex(
        "Summary", {"_key": secrets.token_urlsafe(8), "content": short_summary}
    )

    qdrant_client.upsert(
        "Article",
        points=[
            PointStruct(
                id=str(uuid.uuid4()),
                vector=short_summary_embeddings,
                payload={
                    "external_id": summary_obj["_id"],
                    "article_id": art["_id"],
                },
            )
        ],
    )

    am.summary_of.insert(
        {
            "_key": secrets.token_urlsafe(8),
            "_from": summary_obj["_id"],
            "_to": art["_id"],
        }
    )

    channel.basic_publish(
        "",
        "entity_extraction",
        json.dumps({"type": "extract_problems", "article_id": art["_id"]}),
        properties=pika.BasicProperties(delivery_mode=pika.DeliveryMode.Persistent),
    )