Multimodal RAG

import pandas as pd
from docling_core.types.doc import DoclingDocument


def extract_tables(doc: DoclingDocument) -> list[tuple[int, pd.DataFrame, str]]:
    """(page_no, dataframe, caption) for every table, in reading order."""
    return [(t.prov[0].page_no, t.export_to_dataframe(doc=doc), t.caption_text(doc)) for t in doc.tables]


def stitch(tables: list[tuple[int, pd.DataFrame, str]]) -> list[tuple[int, pd.DataFrame, str]]:
    """Merge a table into the previous one when it continues on the next page."""
    merged: list[tuple[int, pd.DataFrame, str]] = []
    for page, df, caption in tables:
        if merged:
            prev_page, prev_df, prev_caption = merged[-1]
            same_header = list(df.columns) == list(prev_df.columns)
            headerless_continuation = (len(df.columns) == len(prev_df.columns)
                                       and all(str(c).isdigit() for c in df.columns))
            if page == prev_page + 1 and (same_header or headerless_continuation):
                df.columns = prev_df.columns
                merged[-1] = (page, pd.concat([prev_df, df], ignore_index=True), prev_caption)
                continue
        merged.append((page, df, caption))
    return merged


def to_chunk(page: int, df: pd.DataFrame, caption: str, source: str) -> dict:
    return {
        "text": (f"Table: {caption}\n\n" if caption else "") + df.to_markdown(index=False),
        "metadata": {"source": source, "page": page, "element_type": "Table",
                     "rows_json": df.to_dict(orient="records")},
    }

import base64
import io
import os

from docling_core.types.doc import DoclingDocument
from langchain.chat_models import init_chat_model
from langchain_core.documents import Document
from langchain_core.messages import HumanMessage

vision = init_chat_model(os.environ["VISION_MODEL"], temperature=0)

PROMPT = ("Describe this figure from technical documentation for search. State what it shows, every label and "
          "number you can read, and the conclusion it supports. Do not guess text you cannot read.")


def summarise_figures(doc: DoclingDocument, source: str, image_uri: callable) -> list[Document]:
    out = []
    for i, picture in enumerate(doc.pictures):
        image = picture.get_image(doc)                     # needs generate_picture_images=True at parse time
        if image is None:
            continue
        buf = io.BytesIO()
        image.save(buf, format="PNG")
        message = HumanMessage(content_blocks=[
            {"type": "text", "text": PROMPT + "\nCaption: " + picture.caption_text(doc)},
            {"type": "image", "base64": base64.b64encode(buf.getvalue()).decode(), "mime_type": "image/png"},
        ])
        summary = vision.invoke([message]).text
        prov = picture.prov[0]
        out.append(Document(page_content=summary, metadata={
            "source": source, "page": prov.page_no, "bbox": prov.bbox.as_tuple(), "element_type": "Picture",
            "image_uri": image_uri(source, i),             # stored in the object store, used as visual citation
            "parser": "vlm-summary",
        }))
    return out

package com.example.docsassistant.rag.multimodal;

import java.util.Map;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;
import org.springframework.util.MimeTypeUtils;

@Service
public class FigureSummarizer {

    private static final String PROMPT = """
            Describe this figure from technical documentation for search. State what it shows, every label and
            number you can read, and the conclusion it supports. Do not guess text you cannot read.
            Caption: {caption}""";

    private final ChatClient chatClient;

    public FigureSummarizer(ChatClient.Builder builder) {
        this.chatClient = builder.build();
    }

    public Document summarise(Resource png, String caption, String source, int page, String imageUri) {
        String summary = chatClient.prompt()
                .user(u -> u.text(PROMPT).param("caption", caption).media(MimeTypeUtils.IMAGE_PNG, png))
                .call()
                .content();
        return Document.builder()
                .text(summary)
                .metadata(Map.of("source", source, "page", page, "element_type", "Picture",
                        "image_uri", imageUri, "parser", "vlm-summary"))
                .build();
    }
}

package com.example.docsassistant.l4j.multimodal;

import java.util.Base64;

import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.message.ImageContent;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatModel;

public class FigureSummarizer {

    private static final String PROMPT = "Describe this figure from technical documentation for search. State what "
            + "it shows, every label and number you can read, and the conclusion it supports. "
            + "Do not guess text you cannot read. Caption: ";

    private final ChatModel visionModel;

    public FigureSummarizer(ChatModel visionModel) {
        this.visionModel = visionModel;
    }

    public TextSegment summarise(byte[] png, String caption, String source, int page, String imageUri) {
        UserMessage message = UserMessage.from(
                TextContent.from(PROMPT + caption),
                ImageContent.from(Base64.getEncoder().encodeToString(png), "image/png"));
        String summary = visionModel.chat(message).aiMessage().text();
        Metadata metadata = new Metadata()
                .put("source", source).put("page", page).put("element_type", "Picture")
                .put("image_uri", imageUri).put("parser", "vlm-summary");
        return TextSegment.from(summary, metadata);
    }
}

package com.example.docsassistant.rag.multimodal;

import java.util.Map;

import org.springframework.ai.audio.transcription.TranscriptionModel;
import org.springframework.ai.document.Document;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;

@Service
public class TranscriptIngestor {

    private final TranscriptionModel transcriptionModel;

    public TranscriptIngestor(TranscriptionModel transcriptionModel) {
        this.transcriptionModel = transcriptionModel;
    }

    public Document ingest(Resource recording, String source) {
        String transcript = transcriptionModel.transcribe(recording);
        return Document.builder()
                .text(transcript)
                .metadata(Map.of("source", source, "element_type", "Transcript"))
                .build();
    }
}