import re
import httpx
from openai import OpenAI

class VLLMModel():
    def __init__(self, model_name, ip):
        self.client = OpenAI(
            base_url=f"http://{ip}:8000/v1",
            api_key="empty",
        )
        self.model_name = model_name

    def _post_process_result(self, result):
        """如果是mimo模型，则去除<think>...</think>内容"""
        if "mimo" in self.model_name.lower():
            result = re.sub(r"<think>.*?</think>", "", result, flags=re.DOTALL)
        return result.strip()

    def chat(self, user_prompt, image_paths):
        image_contents = [
            {
                "type": "image_url",
                "image_url": {"url": f"file://{img}"}
            } for img in image_paths
        ]
        messages = [{
            "role": "user",
            "content": image_contents + [{"type": "text", "text": user_prompt}]
        }]
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages
        )
        result = response.choices[0].message.content.strip()
        return self._post_process_result(result)
