Final_Assignment_Template / tools /answer_question_from_file.py
AlbertoFor's picture
Added comments
731825e
raw
history blame
2.57 kB
# Importing necessary libraries and modules
from langchain_core.tools.base import BaseTool
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
from langchain_google_genai import ChatGoogleGenerativeAI
from pydantic import PrivateAttr
import os
from dotenv import load_dotenv
import whisper
import base64
load_dotenv(".env", override=True) # Loading environment variables
AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT") # Fetching Azure OpenAI endpoint from environment
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
OPENAI_API_VERSION = os.getenv("OPENAI_API_VERSION_GEN", "2023-12-01-preview") # Default API version
# AZURE_OPENAI_DEPLOYMENT_NAME will be used as the 'model' for API calls
AZURE_OPENAI_DEPLOYMENT_NAME = "gpt-4.1"
# Defining the AnswerQuestionFromFileTool class which extends BaseTool
class AnswerQuestionFromFileTool(BaseTool):
name: str = "answer_question_from_file_tool"
description: str = """
This tool allows you to answer a question taking into account information that were provided inside a file.
You must provide the file in b64 when processing here.
Args:
The question that needs to be answered.
The file extension of the file that is being processed.
"""
_llm = PrivateAttr()
def __init__(self):
# Initializing the AnswerQuestionFromFileTool
super().__init__()
self._llm = ChatGoogleGenerativeAI( # Setting up the LLM with specific parameters
model="gemini-2.0-flash",
temperature=0)
def _run(self, question: str, file_name: str, file_extension: str) -> str:
with open(file_name, "rb") as f:
file = f.read()
if file_extension in ["png", "jpg"]:
encoded_file = base64.b64encode(file).decode("utf-8")
message = {"type": "image_url", "image_url": f"data:image/png;base64,{encoded_file}"}
elif file_extension == "pdf":
encoded_file = base64.b64encode(file).decode("utf-8")
message = {"type": "image_url",
"image_url": f"data:application/pdf;base64,{encoded_file}"
}
else:
message = {"type": "text", "text": "The file is not supported."}
message_local = HumanMessage(
content=[
{"type": "text", "text": question + "\nLet's think step by step."},
message,
]
)
response = self._llm.invoke([message_local])
return response