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

services/article-parser/src/app.py
196 строк6 KB
Бредихин Борис Андреевич
repredict problems and tasks
23 авг 2025, 14:36
23 авг 2025, 14:364ab4b27
import hashlib
import json
import os
import secrets

from fastapi import FastAPI, Response, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.openapi.docs import (
    get_swagger_ui_html,
    get_swagger_ui_oauth2_redirect_html,
)
from fastapi.responses import Response, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import pika
from sse_starlette import EventSourceResponse

from ai.graphrag import GraphRAG
from ai.rag_generation import RAGGenerator
from common import (
    articles_db,
    get_rabbitmq_connection,
    grobid_adapter,
    ollama_client,
    entities_db,
    paper_root,
)
from config import os_logger
from controllers.entity_extractor import EntityExtractor
from routes import clustering, graph
from schemas import (
    ArticleListResponseSchema,
    BulkProcessRequestSchema,
    ChatRequestSchema,
    ExtractEntitiesRequestSchema,
)
from views.article_summarizer import SummarizerView
from views.basic import IndexView

os_logger.info("Starting app...")

app = FastAPI(docs_url=None, redoc_url=None)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.mount("/static", StaticFiles(directory="static"), name="/static")

templates = Jinja2Templates("templates")

graphrag = GraphRAG(ollama_client, "templates/prompts/graphrag_prompt.txt")
rag_generator = RAGGenerator(
    ollama_client, articles_db, entities_db, graph_rag=graphrag
)


@app.get("/docs", include_in_schema=False)
async def custom_swagger_ui_html():
    return get_swagger_ui_html(
        openapi_url=app.openapi_url,
        title=app.title + " - Swagger UI",
        oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
        swagger_js_url="/static/swagger-ui-bundle.js",
        swagger_css_url="/static/swagger-ui.css",
    )


@app.get(app.swagger_ui_oauth2_redirect_url, include_in_schema=False)
async def swagger_ui_redirect():
    return get_swagger_ui_oauth2_redirect_html()


@app.get("/get_tei")
async def get_tei(pdf_url: str):
    return Response(
        grobid_adapter.get_tei_by_pdf_url(pdf_url, paper_root),
        media_type="application/tei+xml",
    )


@app.post("/get_tei")
async def get_tei(file: UploadFile):
    return Response(
        grobid_adapter.get_tei(file.file),
        media_type="application/tei+xml",
    )


@app.get("/tei/{hash}")
def get_tei_by_hash(hash: str):
    with open(os.path.join(paper_root, hash + ".tei.xml"), "r") as f:
        tei = f.read()

    return Response(tei, media_type="application/tei+xml")


@app.get("/hash_article")
async def get_hash_article(pdf_url: str):
    return {
        "hash": hashlib.sha1(pdf_url.encode()).hexdigest(),
    }


@app.post("/extract_entities")
async def extract_entities(data: ExtractEntitiesRequestSchema):
    entity_extractor = EntityExtractor(articles_db, entities_db)

    part = articles_db.db.document({"_id": data.part_id})
    if part.get("entities_extracted"):
        return JSONResponse(entity_extractor.load_entities(part))
    return JSONResponse(entity_extractor.extract_and_save(part, graphrag))


@app.get("/articles")
async def list_articles() -> ArticleListResponseSchema:
    article_collection = articles_db.get_vertex_collection("Article")
    articles = article_collection.all()
    result = [
        {
            "title": article["title"] or f"<<UNKNOWN>> {article['_id']}",
            "url": article.get("url", None),
            "id": article["_id"],
        }
        for article in articles
    ]
    return {
        "articles": result,
    }


prompts: dict[str, str] = {}


@app.post("/chat")
def chat(data: ChatRequestSchema):
    prompt_id = secrets.token_urlsafe(8)
    prompts[prompt_id] = rag_generator.build_prompt(
        data.prompt,
        simple_rag=data.simple_rag,
        graph_rag=data.graph_rag,
        article_id=data.article_id,
    )
    return JSONResponse({"promptId": prompt_id})


@app.post("/chat/stream/{prompt_id}")
async def chat_stream(prompt_id: str) -> EventSourceResponse:
    def stream(prompt: str):
        for i in ollama_client.chat(
            "llama3.1", [{"role": "user", "content": prompt}], stream=True
        ):
            print(i["message"]["content"], end="")
            yield {
                "data": i["message"]["content"],
                "event": "message",
            }

    prompt = prompts.pop(prompt_id)
    return EventSourceResponse(stream(prompt))  # llm.stream(request)


@app.post("/bulk-process")
async def bulk_process(data: BulkProcessRequestSchema):
    theme = data.theme_tag
    for pdf_url in data.pdf_urls:
        connection = get_rabbitmq_connection()
        channel = connection.channel()
        channel.queue_declare(queue="entity_extraction", durable=True)
        channel.basic_publish(
            "",
            "entity_extraction",
            json.dumps(
                {
                    "type": "process_article",
                    "params": pdf_url.model_dump() | {"theme": theme},
                }
            ),
            properties=pika.BasicProperties(delivery_mode=pika.DeliveryMode.Persistent),
        )


# @app.get("/", response_class=HTMLResponse)
# async def index(request: Request):
# return templates.TemplateResponse(request, "index.jinja2.html")
IndexView(templates).mount(app, "/")
SummarizerView(templates, grobid_adapter, paper_root, ollama_client, articles_db).mount(
    app, "/summarize", ["get", "post"]
)

app.include_router(graph.router, prefix="/graph", tags=["Graph"])
app.include_router(clustering.router, prefix="/clustering", tags=["Clustering"])