| |
| """ |
| 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") |
| |
| |
