|
|
|
|
|
""" |
|
|
Created on Sat Oct 26 15:35:19 2024 |
|
|
|
|
|
@author: BM109X32G-10GPU-02 |
|
|
""" |
|
|
|
|
|
from langchain_community.embeddings import OllamaEmbeddings |
|
|
from langchain.tools import BaseTool |
|
|
from langchain_openai import ChatOpenAI |
|
|
from langchain_core.messages import HumanMessage, SystemMessage |
|
|
from langchain.base_language import BaseLanguageModel |
|
|
import base64 |
|
|
from io import BytesIO |
|
|
from PIL import Image |
|
|
import os |
|
|
|
|
|
|
|
|
def convert_to_base64(pil_image): |
|
|
buffered = BytesIO() |
|
|
pil_image.save(buffered, format="PNG") |
|
|
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") |
|
|
return img_str |
|
|
|
|
|
|
|
|
class Imageanalysis(BaseTool): |
|
|
name: str = "Imageanalysis" |
|
|
description: str = ( |
|
|
"Useful to answer questions according to the image, figure, diagram or graph. " |
|
|
"Useful to analysis the information in the image, figure, diagram or graph. " |
|
|
"Input query about image/figure/graph/diagram, return the response" |
|
|
) |
|
|
return_direct: bool = True |
|
|
llm: BaseLanguageModel = None |
|
|
path : str = None |
|
|
|
|
|
def __init__(self, path): |
|
|
super().__init__( ) |
|
|
self.llm = ChatOpenAI(model="gpt-5",api_key=os.getenv("OPENAI_API_KEY"), |
|
|
base_url=os.getenv("OPENAI_API_BASE")) |
|
|
self.path = path |
|
|
|
|
|
|
|
|
def _run(self, query ) -> str: |
|
|
try: |
|
|
pil_image = Image.open(self.path) |
|
|
rgb_im = pil_image.convert('RGB') |
|
|
image_b64 = convert_to_base64(pil_image) |
|
|
message = HumanMessage( |
|
|
content=[ |
|
|
{"type": "text", "text": query}, |
|
|
{ |
|
|
"type": "image_url", |
|
|
"image_url": {"url":f"data:image/jpeg;base64,{image_b64}"}, |
|
|
}, |
|
|
],) |
|
|
response = self.llm.invoke([message]) |
|
|
return response.content |
|
|
|
|
|
except Exception as e: |
|
|
return str(e) |
|
|
|
|
|
|
|
|
async def _arun(self, query) -> str: |
|
|
"""Use the tool asynchronously.""" |
|
|
raise NotImplementedError("this tool does not support async") |
|
|
|
|
|
|
