rag/document_parser.py - dev - zidanez/mcp-test

rag/document_parser.py
78 строк3 KB
zidanez
Add rag
24 мар 2025, 00:04
24 мар 2025, 00:042f0c5cc
import os
import logging
import dotenv
import glob
import typing as t
from langchain_core.messages import SystemMessage, HumanMessage, BaseMessage

from langchain_gigachat import GigaChat

# from pdf2image import convert_from_path


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

dotenv.load_dotenv()

DESCRIBING_SLIDE_PROMPT_TEMPLATE = """Тебе представлен документ бизнес-требования. Опиши ключевую информацию, которая \
необходима для бизнес-требования.
Описание должно быть конкретным и точным. Обрати особое внимание на графики, диаграммы \
или визуальные элементы, которые можно проанализировать. 
Если в разделе присутсвуют графики, диаграммы подробно опиши их.
Очень важно не упустить детали.
Ответь в формате Markdown."""

LLM = GigaChat(
    base_url="https://gigachat.devices.sberbank.ru/api/v1",
    credentials=os.getenv("GIGACHAT_API_KEY"),
    scope="GIGACHAT_API_PERS",
    model="GigaChat-2-Max",
    verify_ssl_certs=False,
    temperature=0.4,
    timeout=3000,
    max_tokens=10000
)


# def save_pdf_pages_as_images(pdf_path: str, output_dir: str) -> None:
#     if not os.path.exists(output_dir):
#         os.makedirs(output_dir)
#
#     images = convert_from_path(pdf_path)
#
#     for i, image in enumerate(images):
#         image_path = os.path.join(output_dir, f"page_{i + 1}.png")
#         image.save(image_path, "PNG")
#         print(f"Сохранена страница {i + 1} как {image_path}")
#

def document_summary(path: str) -> str:
    """На вход поучаем путь до файла БТ.\n\n
    Суммаризируем и отдаем результат суммаризации"""

    file = LLM.upload_file(open(path, "rb"))

    additional_kwargs = {"attachments": [file.id_]}
    llm_promt: t.List[BaseMessage] = [
        SystemMessage(DESCRIBING_SLIDE_PROMPT_TEMPLATE),
        HumanMessage(content="", additional_kwargs=additional_kwargs)
    ]

    response = LLM.invoke(llm_promt)
    return response.content


def examples_summarization(example_dir_path: str, output_dir_path: str):
    """Суммаризируем все файлы в `example_dir_path`
    и выводим итог суммаризации в `output_dir_path/{basename_file}.txt`"""

    for file in glob.glob(f"{example_dir_path}/*.pdf"):
        file_base_name: str = os.path.basename(file).replace('.pdf', '')
        summary: str = document_summary(file)

        with open(f"{output_dir_path}/{file_base_name}.txt", 'w+') as f:
            f.write(summary)


if __name__ == '__main__':
    examples_summarization("../bt_examples", "./bt_sumaries")