Pacama95/chatbot_agent

1from pydantic import BaseModel, Field2from langchain.tools import BaseTool3from PIL import Image4from transformers import pipeline5import requests6 7from typing import Type, Final, Any8 9DOCUMENT_QUESTION_MODEL: Final[str] = 'impira/layoutlm-document-qa'10 11class DocumentQuestionAnsweringInput(BaseModel):12    query: str = Field(description='Your question about the document')13    document_url: str = Field(description='URL to the document.')14 15class DocumentQuestionAnsweringTool(BaseTool):16 17    name: str = "document_question_answering_tool"18    description: str = """19        Use this tool when asked to answer questions about documents.20        This tool is designed to answer basic factual questions based on the contents of a text-based document accessible via a public URL. It supports simple information retrieval tasks such as:21            - Identifying key facts (e.g., “What is the date of the agreement?”)22            - Locating named entities (e.g., “Who signed the contract?”)23            - Extracting short text spans or direct answers from the document (e.g., “What is the total cost?”)24    """25    args_schema: Type[BaseModel] = DocumentQuestionAnsweringInput26 27    def __init__(self, **kwargs: Any) -> None:28        super().__init__()29 30    def _run(self, document_url: str, query: str) -> str:31        """32            Perform basic visual question answering on documents accessible via a public URL.33            Args:34                document_url (str): The URL to the document to analyze.35                query (str): Your question about the provided document.36 37            Returns:38                str: An answer for the query based on the document content39        """40        pipe = pipeline("document-question-answering", model="impira/layoutlm-document-qa")41 42        # Load document from URL43        document = Image.open(requests.get(document_url, stream=True).raw)44 45        result = pipe(image=document, question=query)[0]46 47        return f"Answer: '{result['answer']}' - (Precission: {result['score']})"48