Skip to content

feat: support pdf files in Prompt #452

Open
@konrad-czarnota-ds

Description

@konrad-czarnota-ds

Feature description

I would like to be able to attach .pdf files to prompts in ragbits and use them when chatting with LLM. Here's a quick and dirty solution that I did to make it work that needs to be polished:

import base64
from pydantic import BaseModel
from ragbits.core.prompt import Prompt


class InstructionExtractionPromptInput(BaseModel):
    pdf_path: str


class InstructionExtractionPromptOutput(BaseModel):
    instructions: str


class InstructionExtractionPrompt(Prompt[InstructionExtractionPromptInput, InstructionExtractionPromptOutput]):
    user_prompt = """
    """

    image_input_fields = ["pdf_path"]

    @staticmethod
    def _create_message_with_image(image: str | bytes) -> dict:
        if type(image) == str and image.endswith(".pdf"):
            with open(image, "rb") as f:
                data = f.read()

            base64_string = base64.b64encode(data).decode("utf-8")

            return{
                        "type": "file",
                        "file": {
                            "filename": "instruction.pdf",
                            "file_data": f"data:application/pdf;base64,{base64_string}",
                        }
                    }
        else:
            return super()._create_message_with_image(image)

Motivation

LiteLLM supports files passing through API, I think ragbits should too as this might be necessary to handle different documents

Additional context

No response

Metadata

Metadata

Assignees

Labels

featureNew feature or request

Type

No type

Projects

Status

In review

Relationships

None yet

Development

No branches or pull requests

Issue actions