vlm4ocr 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

vlm4ocr/vlm_engines.py

@@ -1,337 +1,27 @@
 import abc
-import importlib.util
-from typing import Any, List, Dict, Union, Generator
-import warnings
-import os
-import re
+from typing import List, Dict
 from PIL import Image
 from vlm4ocr.utils import image_to_base64
 from vlm4ocr.data_types import FewShotExample
-
-
-class VLMConfig(abc.ABC):
-    def __init__(self, **kwargs):
-        """
-        This is an abstract class to provide interfaces for VLM configuration. 
-        Children classes that inherts this class can be used in extrators and prompt editor.
-        Common VLM parameters: max_new_tokens, temperature, top_p, top_k, min_p.
-        """
-        self.params = kwargs.copy()
-
-    @abc.abstractmethod
-    def preprocess_messages(self, messages:List[Dict[str,str]]) -> List[Dict[str,str]]:
-        """
-        This method preprocesses the input messages before passing them to the VLM.
-
-        Parameters:
-        ----------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        
-        Returns:
-        -------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        """
-        return NotImplemented
-
-    @abc.abstractmethod
-    def postprocess_response(self, response:Union[str, Dict[str, str], Generator[str, None, None]]) -> Union[str, Generator[str, None, None]]:
-        """
-        This method postprocesses the VLM response after it is generated.
-
-        Parameters:
-        ----------
-        response : Union[str, Generator[str, None, None]]
-            the VLM response. Can be a string or a generator.
-        
-        Returns:
-        -------
-        response : str
-            the postprocessed VLM response
-        """
-        return NotImplemented
-
-
-class BasicVLMConfig(VLMConfig):
-    def __init__(self, max_new_tokens:int=2048, temperature:float=0.0, **kwargs):
-        """
-        The basic VLM configuration for most non-reasoning models.
-        """
-        super().__init__(**kwargs)
-        self.max_new_tokens = max_new_tokens
-        self.temperature = temperature
-        self.params["max_new_tokens"] = self.max_new_tokens
-        self.params["temperature"] = self.temperature
-
-    def preprocess_messages(self, messages:List[Dict[str,str]]) -> List[Dict[str,str]]:
-        """
-        This method preprocesses the input messages before passing them to the VLM.
-
-        Parameters:
-        ----------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        
-        Returns:
-        -------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        """
-        return messages
-
-    def postprocess_response(self, response:Union[str, Dict[str, str], Generator[str, None, None]]) -> Union[Dict[str, str], Generator[Dict[str, str], None, None]]:
-        """
-        This method postprocesses the VLM response after it is generated.
-
-        Parameters:
-        ----------
-        response : Union[str, Generator[str, None, None]]
-            the VLM response. Can be a string or a generator.
-        
-        Returns: Union[str, Generator[Dict[str, str], None, None]]
-            the postprocessed VLM response. 
-            if input is a generator, the output will be a generator {"type": "response", "data": <content>}.
-        """
-        if isinstance(response, str):
-            return {"response": response}
-        
-        elif isinstance(response, dict):
-            if "response" in response:
-                return response
-            else:
-                warnings.warn(f"Invalid response dict keys: {response.keys()}. Returning default empty dict.", UserWarning)
-                return {"response": ""}
-
-        def _process_stream():
-            for chunk in response:
-                if isinstance(chunk, dict):
-                    yield chunk
-                elif isinstance(chunk, str):
-                    yield {"type": "response", "data": chunk}
-
-        return _process_stream()
-    
-class ReasoningVLMConfig(VLMConfig):
-    def __init__(self, thinking_token_start="<think>", thinking_token_end="</think>", **kwargs):
-        """
-        The general configuration for reasoning vision models.
-        """
-        super().__init__(**kwargs)
-        self.thinking_token_start = thinking_token_start
-        self.thinking_token_end = thinking_token_end
-
-    def preprocess_messages(self, messages:List[Dict[str,str]]) -> List[Dict[str,str]]:
-        """
-        This method preprocesses the input messages before passing them to the VLM.
-
-        Parameters:
-        ----------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        
-        Returns:
-        -------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        """
-        return messages.copy()
-
-    def postprocess_response(self, response:Union[str, Dict[str, str], Generator[str, None, None]]) -> Union[Dict[str,str], Generator[Dict[str,str], None, None]]:
-        """
-        This method postprocesses the VLM response after it is generated.
-        1. If input is a string, it will extract the reasoning and response based on the thinking tokens.
-        2. If input is a dict, it should contain keys "reasoning" and "response". This is for inference engines that already parse reasoning and response.
-        3. If input is a generator, 
-            a. if the chunk is a dict, it should contain keys "type" and "data". This is for inference engines that already parse reasoning and response.
-            b. if the chunk is a string, it will yield dicts with keys "type" and "data" based on the thinking tokens.
-
-        Parameters:
-        ----------
-        response : Union[str, Generator[str, None, None]]
-            the VLM response. Can be a string or a generator.
-        
-        Returns:
-        -------
-        response : Union[str, Generator[str, None, None]]
-            the postprocessed LLM response as a dict {"reasoning": <reasoning>, "response": <content>}
-            if input is a generator, the output will be a generator {"type": <reasoning or response>, "data": <content>}.
-        """
-        if isinstance(response, str):
-            # get contents between thinking_token_start and thinking_token_end
-            pattern = f"{re.escape(self.thinking_token_start)}(.*?){re.escape(self.thinking_token_end)}"
-            match = re.search(pattern, response, re.DOTALL)
-            reasoning = match.group(1) if match else ""
-            # get response AFTER thinking_token_end
-            response = re.sub(f".*?{self.thinking_token_end}", "", response, flags=re.DOTALL).strip()
-            return {"reasoning": reasoning, "response": response}
-
-        elif isinstance(response, dict):
-            if "reasoning" in response and "response" in response:
-                return response
-            else:
-                warnings.warn(f"Invalid response dict keys: {response.keys()}. Returning default empty dict.", UserWarning)
-                return {"reasoning": "", "response": ""}
-
-        elif isinstance(response, Generator):
-            def _process_stream():
-                think_flag = False
-                buffer = ""
-                for chunk in response:
-                    if isinstance(chunk, dict):
-                        yield chunk
-
-                    elif isinstance(chunk, str):
-                        buffer += chunk
-                        # switch between reasoning and response
-                        if self.thinking_token_start in buffer:
-                            think_flag = True
-                            buffer = buffer.replace(self.thinking_token_start, "")
-                        elif self.thinking_token_end in buffer:
-                            think_flag = False
-                            buffer = buffer.replace(self.thinking_token_end, "")
-                        
-                        # if chunk is in thinking block, tag it as reasoning; else tag it as response
-                        if chunk not in [self.thinking_token_start, self.thinking_token_end]:
-                            if think_flag:
-                                yield {"type": "reasoning", "data": chunk}
-                            else:
-                                yield {"type": "response", "data": chunk}
-
-            return _process_stream()
-        
-        else:
-            warnings.warn(f"Invalid response type: {type(response)}. Returning default empty dict.", UserWarning)
-            return {"reasoning": "", "response": ""}
-
-
-class OpenAIReasoningVLMConfig(ReasoningVLMConfig):
-    def __init__(self, reasoning_effort:str="low", **kwargs):
-        """
-        The OpenAI "o" series configuration.
-        1. The reasoning effort is set to "low" by default.
-        2. The temperature parameter is not supported and will be ignored.
-        3. The system prompt is not supported and will be concatenated to the next user prompt.
-
-        Parameters:
-        ----------
-        reasoning_effort : str, Optional
-            the reasoning effort. Must be one of {"low", "medium", "high"}. Default is "low".
-        """
-        super().__init__(**kwargs)
-        if reasoning_effort not in ["low", "medium", "high"]:
-            raise ValueError("reasoning_effort must be one of {'low', 'medium', 'high'}.")
-
-        self.reasoning_effort = reasoning_effort
-        self.params["reasoning_effort"] = self.reasoning_effort
-
-        if "temperature" in self.params:
-            warnings.warn("Reasoning models do not support temperature parameter. Will be ignored.", UserWarning)
-            self.params.pop("temperature")
-
-    def preprocess_messages(self, messages:List[Dict[str,str]]) -> List[Dict[str,str]]:
-        """
-        Concatenate system prompts to the next user prompt.
-
-        Parameters:
-        ----------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        
-        Returns:
-        -------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        """
-        system_prompt_holder = ""
-        new_messages = []
-        for i, message in enumerate(messages):
-            # if system prompt, store it in system_prompt_holder
-            if message['role'] == 'system':
-                system_prompt_holder = message['content']
-            # if user prompt, concatenate it with system_prompt_holder
-            elif message['role'] == 'user':
-                if system_prompt_holder:
-                    new_message = {'role': message['role'], 'content': f"{system_prompt_holder} {message['content']}"}
-                    system_prompt_holder = ""
-                else:
-                    new_message = {'role': message['role'], 'content': message['content']}
-
-                new_messages.append(new_message)
-            # if assistant/other prompt, do nothing
-            else:
-                new_message = {'role': message['role'], 'content': message['content']}
-                new_messages.append(new_message)
-
-        return new_messages
-
-
-class MessagesLogger:
-    def __init__(self):
-        """
-        This class is used to log the messages for InferenceEngine.chat().
-        """
-        self.messages_log = []
-
-    def log_messages(self, messages : List[Dict[str,str]]):
-        """
-        This method logs the messages to a list.
-        """
-        self.messages_log.append(messages)
-
-    def get_messages_log(self) -> List[List[Dict[str,str]]]:
-        """
-        This method returns a copy of the current messages log
-        """
-        return self.messages_log.copy()
-    
-    def clear_messages_log(self):
-        """
-        This method clears the current messages log
-        """
-        self.messages_log.clear()
-
-
-class VLMEngine:
-    @abc.abstractmethod
-    def __init__(self, config:VLMConfig, **kwrs):
-        """
-        This is an abstract class to provide interfaces for VLM inference engines. 
-        Children classes that inherts this class can be used in extrators. Must implement chat() method.
-
-        Parameters:
-        ----------
-        config : VLMConfig
-            the VLM configuration. Must be a child class of VLMConfig.
-        """
-        return NotImplemented
-
-    @abc.abstractmethod
-    def chat(self, messages:List[Dict[str,str]], verbose:bool=False, stream:bool=False, 
-             messages_logger:MessagesLogger=None) -> Union[Dict[str, str], Generator[Dict[str, str], None, None]]:
-        """
-        This method inputs chat messages and outputs VLM generated text.
-
-        Parameters:
-        ----------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        verbose : bool, Optional
-            if True, VLM generated text will be printed in terminal in real-time.
-        stream : bool, Optional
-            if True, returns a generator that yields the output in real-time.
-        Messages_logger : MessagesLogger, Optional
-            the message logger that logs the chat messages.
-        """
-        return NotImplemented
-    
-    @abc.abstractmethod
-    def chat_async(self, messages:List[Dict[str,str]], messages_logger:MessagesLogger=None) -> Dict[str, str]:
-        """
-        The async version of chat method. Streaming is not supported.
-        """
-        return NotImplemented
-
+from llm_inference_engine.llm_configs import (
+    LLMConfig as VLMConfig, 
+    BasicLLMConfig as BasicVLMConfig, 
+    ReasoningLLMConfig as ReasoningVLMConfig, 
+    OpenAIReasoningLLMConfig as OpenAIReasoningVLMConfig
+)
+from llm_inference_engine.utils import MessagesLogger
+from llm_inference_engine.engines import (
+    InferenceEngine,
+    OllamaInferenceEngine,
+    OpenAICompatibleInferenceEngine,
+    VLLMInferenceEngine,
+    OpenRouterInferenceEngine,
+    OpenAIInferenceEngine,
+    AzureOpenAIInferenceEngine,
+)
+
+
+class VLMEngine(InferenceEngine):
     @abc.abstractmethod
     def get_ocr_messages(self, system_prompt:str, user_prompt:str, image:Image.Image, few_shot_examples:List[FewShotExample]=None) -> List[Dict[str,str]]:
         """
@@ -349,217 +39,9 @@ class VLMEngine:
             list of few-shot examples. 
         """
         return NotImplemented
-    
-    def _format_config(self) -> Dict[str, Any]:
-        """
-        This method format the VLM configuration with the correct key for the inference engine. 
-
-        Return : Dict[str, Any]
-            the config parameters.
-        """
-        return NotImplemented
-
-
-class OllamaVLMEngine(VLMEngine):
-    def __init__(self, model_name:str, num_ctx:int=8192, keep_alive:int=300, config:VLMConfig=None, **kwrs):
-        """
-        The Ollama inference engine.
-
-        Parameters:
-        ----------
-        model_name : str
-            the model name exactly as shown in >> ollama ls
-        num_ctx : int, Optional
-            context length that LLM will evaluate.
-        keep_alive : int, Optional
-            seconds to hold the LLM after the last API call.
-        config : LLMConfig
-            the LLM configuration. 
-        """
-        if importlib.util.find_spec("ollama") is None:
-            raise ImportError("ollama-python not found. Please install ollama-python (```pip install ollama```).")
-        
-        from ollama import Client, AsyncClient
-        self.client = Client(**kwrs)
-        self.async_client = AsyncClient(**kwrs)
-        self.model_name = model_name
-        self.num_ctx = num_ctx
-        self.keep_alive = keep_alive
-        self.config = config if config else BasicVLMConfig()
-        self.formatted_params = self._format_config()
-    
-    def _format_config(self) -> Dict[str, Any]:
-        """
-        This method format the LLM configuration with the correct key for the inference engine. 
-        """
-        formatted_params = self.config.params.copy()
-        if "max_new_tokens" in formatted_params:
-            formatted_params["num_predict"] = formatted_params["max_new_tokens"]
-            formatted_params.pop("max_new_tokens")
-
-        return formatted_params
-
-    def chat(self, messages:List[Dict[str,str]], verbose:bool=False, stream:bool=False, 
-             messages_logger:MessagesLogger=None) -> Union[Dict[str,str], Generator[Dict[str, str], None, None]]:
-        """
-        This method inputs chat messages and outputs VLM generated text.
-
-        Parameters:
-        ----------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        verbose : bool, Optional
-            if True, VLM generated text will be printed in terminal in real-time.
-        stream : bool, Optional
-            if True, returns a generator that yields the output in real-time.
-        Messages_logger : MessagesLogger, Optional
-            the message logger that logs the chat messages.
-
-        Returns:
-        -------
-        response : Union[Dict[str,str], Generator[Dict[str, str], None, None]]
-            a dict {"reasoning": <reasoning>, "response": <response>} or Generator {"type": <reasoning or response>, "data": <content>}
-        """
-        processed_messages = self.config.preprocess_messages(messages)
 
-        options={'num_ctx': self.num_ctx, **self.formatted_params}
-        if stream:
-            def _stream_generator():
-                response_stream = self.client.chat(
-                    model=self.model_name, 
-                    messages=processed_messages, 
-                    options=options,
-                    stream=True, 
-                    keep_alive=self.keep_alive
-                )
-                res = {"reasoning": "", "response": ""}
-                for chunk in response_stream:
-                    if hasattr(chunk.message, 'thinking') and chunk.message.thinking:
-                        content_chunk = getattr(getattr(chunk, 'message', {}), 'thinking', '')
-                        res["reasoning"] += content_chunk
-                        yield {"type": "reasoning", "data": content_chunk}
-                    else:
-                        content_chunk = getattr(getattr(chunk, 'message', {}), 'content', '')
-                        res["response"] += content_chunk
-                        yield {"type": "response", "data": content_chunk}
 
-                    if chunk.done_reason == "length":
-                        warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
-                
-                # Postprocess response
-                res_dict = self.config.postprocess_response(res)
-                # Write to messages log
-                if messages_logger:
-                    # replace images content with a placeholder "[image]" to save space
-                    for messages in processed_messages:
-                        if "images" in messages:
-                            messages["images"] = ["[image]" for _ in messages["images"]]
-
-                    processed_messages.append({"role": "assistant",
-                                                "content": res_dict.get("response", ""),
-                                                "reasoning": res_dict.get("reasoning", "")})
-                    messages_logger.log_messages(processed_messages)
-
-            return self.config.postprocess_response(_stream_generator())
-
-        elif verbose:
-            response = self.client.chat(
-                            model=self.model_name, 
-                            messages=processed_messages, 
-                            options=options,
-                            stream=True,
-                            keep_alive=self.keep_alive
-                        )
-            
-            res = {"reasoning": "", "response": ""}
-            phase = ""
-            for chunk in response:
-                if hasattr(chunk.message, 'thinking') and chunk.message.thinking:
-                    if phase != "reasoning":
-                        print("\n--- Reasoning ---")
-                        phase = "reasoning"
-
-                    content_chunk = getattr(getattr(chunk, 'message', {}), 'thinking', '')
-                    res["reasoning"] += content_chunk
-                else:
-                    if phase != "response":
-                        print("\n--- Response ---")
-                        phase = "response"
-                    content_chunk = getattr(getattr(chunk, 'message', {}), 'content', '')
-                    res["response"] += content_chunk
-
-                print(content_chunk, end='', flush=True)
-
-                if chunk.done_reason == "length":
-                    warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
-            print('\n')
-
-        else:
-            response = self.client.chat(
-                                model=self.model_name, 
-                                messages=processed_messages, 
-                                options=options,
-                                stream=False,
-                                keep_alive=self.keep_alive
-                            )
-            res = {"reasoning": getattr(getattr(response, 'message', {}), 'thinking', ''),
-                   "response": getattr(getattr(response, 'message', {}), 'content', '')}
-        
-            if response.done_reason == "length":
-                warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
-
-        # Postprocess response
-        res_dict = self.config.postprocess_response(res)
-        # Write to messages log
-        if messages_logger:
-            # replace images content with a placeholder "[image]" to save space
-            for messages in processed_messages:
-                if "images" in messages:
-                    messages["images"] = ["[image]" for _ in messages["images"]]
-
-            processed_messages.append({"role": "assistant", 
-                                    "content": res_dict.get("response", ""), 
-                                    "reasoning": res_dict.get("reasoning", "")})
-            messages_logger.log_messages(processed_messages)
-
-        return res_dict
-        
-
-    async def chat_async(self, messages:List[Dict[str,str]], messages_logger:MessagesLogger=None) -> Dict[str,str]:
-        """
-        Async version of chat method. Streaming is not supported.
-        """
-        processed_messages = self.config.preprocess_messages(messages)
-
-        response = await self.async_client.chat(
-                            model=self.model_name, 
-                            messages=processed_messages, 
-                            options={'num_ctx': self.num_ctx, **self.formatted_params},
-                            stream=False,
-                            keep_alive=self.keep_alive
-                        )
-        
-        res = {"reasoning": getattr(getattr(response, 'message', {}), 'thinking', ''),
-               "response": getattr(getattr(response, 'message', {}), 'content', '')}
-        
-        if response.done_reason == "length":
-            warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
-        # Postprocess response
-        res_dict = self.config.postprocess_response(res)
-        # Write to messages log
-        if messages_logger:
-            # replace images content with a placeholder "[image]" to save space
-            for messages in processed_messages:
-                if "images" in messages:
-                    messages["images"] = ["[image]" for _ in messages["images"]]
-
-            processed_messages.append({"role": "assistant", 
-                                        "content": res_dict.get("response", ""), 
-                                        "reasoning": res_dict.get("reasoning", "")})
-            messages_logger.log_messages(processed_messages)
-
-        return res_dict
-    
+class OllamaVLMEngine(OllamaInferenceEngine, VLMEngine):
     def get_ocr_messages(self, system_prompt:str, user_prompt:str, image:Image.Image, few_shot_examples:List[FewShotExample]=None) -> List[Dict[str,str]]:
         """
         This method inputs an image and returns the correesponding chat messages for the inference engine.
@@ -600,217 +82,7 @@ class OllamaVLMEngine(VLMEngine):
         return output_messages
 
 
-class OpenAICompatibleVLMEngine(VLMEngine):
-    def __init__(self, model:str, api_key:str, base_url:str, config:VLMConfig=None, **kwrs):
-        """
-        General OpenAI-compatible server inference engine.
-        https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html
-
-        For parameters and documentation, refer to https://platform.openai.com/docs/api-reference/introduction
-
-        Parameters:
-        ----------
-        model_name : str
-            model name as shown in the vLLM server
-        api_key : str
-            the API key for the vLLM server.
-        base_url : str
-            the base url for the vLLM server. 
-        config : LLMConfig
-            the LLM configuration.
-        """
-        if importlib.util.find_spec("openai") is None:
-            raise ImportError("OpenAI Python API library not found. Please install OpanAI (```pip install openai```).")
-        
-        from openai import OpenAI, AsyncOpenAI
-        from openai.types.chat import ChatCompletionChunk
-        self.ChatCompletionChunk = ChatCompletionChunk
-        super().__init__(config)
-        self.client = OpenAI(api_key=api_key, base_url=base_url, **kwrs)
-        self.async_client = AsyncOpenAI(api_key=api_key, base_url=base_url, **kwrs)
-        self.model = model
-        self.config = config if config else BasicVLMConfig()
-        self.formatted_params = self._format_config()
-
-    def _format_config(self) -> Dict[str, Any]:
-        """
-        This method format the VLM configuration with the correct key for the inference engine. 
-        """
-        formatted_params = self.config.params.copy()
-        if "max_new_tokens" in formatted_params:
-            formatted_params["max_completion_tokens"] = formatted_params["max_new_tokens"]
-            formatted_params.pop("max_new_tokens")
-
-        return formatted_params
-    
-
-    def _format_response(self, response: Any) -> Dict[str, str]:
-        """
-        This method format the response from OpenAI API to a dict with keys "type" and "data".
-
-        Parameters:
-        ----------
-        response : Any
-            the response from OpenAI-compatible API. Could be a dict, generator, or object.
-        """
-        if isinstance(response, self.ChatCompletionChunk):
-            chunk_text = getattr(response.choices[0].delta, "content", "")
-            if chunk_text is None:
-                chunk_text = ""
-            return {"type": "response", "data": chunk_text}
-
-        return {"response": getattr(response.choices[0].message, "content", "")}
-
-    def chat(self, messages:List[Dict[str,str]], verbose:bool=False, stream:bool=False, 
-             messages_logger:MessagesLogger=None) -> Union[Dict[str, str], Generator[Dict[str, str], None, None]]:
-        """
-        This method inputs chat messages and outputs LLM generated text.
-
-        Parameters:
-        ----------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        verbose : bool, Optional
-            if True, VLM generated text will be printed in terminal in real-time.
-        stream : bool, Optional
-            if True, returns a generator that yields the output in real-time.
-        messages_logger : MessagesLogger, Optional
-            the message logger that logs the chat messages.
-
-        Returns:
-        -------
-        response : Union[Dict[str,str], Generator[Dict[str, str], None, None]]
-            a dict {"reasoning": <reasoning>, "response": <response>} or Generator {"type": <reasoning or response>, "data": <content>}
-        """
-        processed_messages = self.config.preprocess_messages(messages)
-
-        if stream:
-            def _stream_generator():
-                response_stream = self.client.chat.completions.create(
-                                        model=self.model,
-                                        messages=processed_messages,
-                                        stream=True,
-                                        **self.formatted_params
-                                    )
-                res_text = ""
-                for chunk in response_stream:
-                    if len(chunk.choices) > 0:
-                        chunk_dict = self._format_response(chunk)
-                        yield chunk_dict
-
-                        res_text += chunk_dict["data"]
-                        if chunk.choices[0].finish_reason == "length":
-                            warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
-
-                # Postprocess response
-                res_dict = self.config.postprocess_response(res_text)
-                # Write to messages log
-                if messages_logger:
-                    # replace images content with a placeholder "[image]" to save space
-                    for messages in processed_messages:
-                        if "content" in messages and isinstance(messages["content"], list):
-                            for content in messages["content"]:
-                                if isinstance(content, dict) and content.get("type") == "image_url":
-                                    content["image_url"]["url"] = "[image]"
-
-                    processed_messages.append({"role": "assistant",
-                                                "content": res_dict.get("response", ""),
-                                                "reasoning": res_dict.get("reasoning", "")})
-                    messages_logger.log_messages(processed_messages)
-
-            return self.config.postprocess_response(_stream_generator())
-
-        elif verbose:
-            response = self.client.chat.completions.create(
-                model=self.model,
-                messages=processed_messages,
-                stream=True,
-                **self.formatted_params
-            )
-            res = {"reasoning": "", "response": ""}
-            phase = ""
-            for chunk in response:
-                if len(chunk.choices) > 0:
-                    chunk_dict = self._format_response(chunk)
-                    chunk_text = chunk_dict["data"]
-                    res[chunk_dict["type"]] += chunk_text
-                    if phase != chunk_dict["type"] and chunk_text != "":
-                        print(f"\n--- {chunk_dict['type'].capitalize()} ---")
-                        phase = chunk_dict["type"]
-
-                    print(chunk_text, end="", flush=True)
-                    if chunk.choices[0].finish_reason == "length":
-                        warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
-
-            print('\n')
-
-        else:
-            response = self.client.chat.completions.create(
-                model=self.model,
-                messages=processed_messages,
-                stream=False,
-                **self.formatted_params
-            )
-            res = self._format_response(response)
-
-            if response.choices[0].finish_reason == "length":
-                warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
-            
-        # Postprocess response
-        res_dict = self.config.postprocess_response(res)
-        # Write to messages log
-        if messages_logger:
-            # replace images content with a placeholder "[image]" to save space
-            for messages in processed_messages:
-                if "content" in messages and isinstance(messages["content"], list):
-                    for content in messages["content"]:
-                        if isinstance(content, dict) and content.get("type") == "image_url":
-                            content["image_url"]["url"] = "[image]"
-
-            processed_messages.append({"role": "assistant", 
-                                    "content": res_dict.get("response", ""), 
-                                    "reasoning": res_dict.get("reasoning", "")})
-            messages_logger.log_messages(processed_messages)
-
-        return res_dict
-    
-
-    async def chat_async(self, messages:List[Dict[str,str]], messages_logger:MessagesLogger=None) -> Dict[str,str]:
-        """
-        Async version of chat method. Streaming is not supported.
-        """
-        processed_messages = self.config.preprocess_messages(messages)
-
-        response = await self.async_client.chat.completions.create(
-            model=self.model,
-            messages=processed_messages,
-            stream=False,
-            **self.formatted_params
-        )
-        
-        if response.choices[0].finish_reason == "length":
-            warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
-
-        res = self._format_response(response)
-
-        # Postprocess response
-        res_dict = self.config.postprocess_response(res)
-        # Write to messages log
-        if messages_logger:
-            # replace images content with a placeholder "[image]" to save space
-            for messages in processed_messages:
-                if "content" in messages and isinstance(messages["content"], list):
-                    for content in messages["content"]:
-                        if isinstance(content, dict) and content.get("type") == "image_url":
-                            content["image_url"]["url"] = "[image]"
-
-            processed_messages.append({"role": "assistant", 
-                                        "content": res_dict.get("response", ""), 
-                                        "reasoning": res_dict.get("reasoning", "")})
-            messages_logger.log_messages(processed_messages)
-
-        return res_dict
-    
+class OpenAICompatibleVLMEngine(OpenAICompatibleInferenceEngine, VLMEngine):
     def get_ocr_messages(self, system_prompt:str, user_prompt:str, image:Image.Image, format:str='png', 
                          detail:str="high", few_shot_examples:List[FewShotExample]=None) -> List[Dict[str,str]]:
         """
@@ -879,274 +151,44 @@ class OpenAICompatibleVLMEngine(VLMEngine):
         return output_messages
 
 
-class VLLMVLMEngine(OpenAICompatibleVLMEngine):
-    def __init__(self, model:str, api_key:str="", base_url:str="http://localhost:8000/v1", config:VLMConfig=None, **kwrs):
-        """
-        vLLM OpenAI compatible server inference engine.
-        https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html
-
-        For parameters and documentation, refer to https://platform.openai.com/docs/api-reference/introduction
-
-        Parameters:
-        ----------
-        model_name : str
-            model name as shown in the vLLM server
-        api_key : str, Optional
-            the API key for the vLLM server.
-        base_url : str, Optional
-            the base url for the vLLM server. 
-        config : LLMConfig
-            the LLM configuration.
-        """
-        super().__init__(model, api_key, base_url, config, **kwrs)
-
-
-    def _format_response(self, response: Any) -> Dict[str, str]:
-        """
-        This method format the response from OpenAI API to a dict with keys "type" and "data".
-
-        Parameters:
-        ----------
-        response : Any
-            the response from OpenAI-compatible API. Could be a dict, generator, or object.
-        """
-        if isinstance(response, self.ChatCompletionChunk):
-            if hasattr(response.choices[0].delta, "reasoning_content") and getattr(response.choices[0].delta, "reasoning_content") is not None:
-                chunk_text = getattr(response.choices[0].delta, "reasoning_content", "")
-                if chunk_text is None:
-                    chunk_text = ""
-                return {"type": "reasoning", "data": chunk_text}
-            else:
-                chunk_text = getattr(response.choices[0].delta, "content", "")
-                if chunk_text is None:
-                    chunk_text = ""
-                return {"type": "response", "data": chunk_text}
-
-        return {"reasoning": getattr(response.choices[0].message, "reasoning_content", ""),
-                "response": getattr(response.choices[0].message, "content", "")}
-        
-
-class OpenRouterVLMEngine(OpenAICompatibleVLMEngine):
-    def __init__(self, model:str, api_key:str=None, base_url:str="https://openrouter.ai/api/v1", config:VLMConfig=None, **kwrs):
-        """
-        OpenRouter OpenAI-compatible server inference engine.
-
-        Parameters:
-        ----------
-        model_name : str
-            model name as shown in the vLLM server
-        api_key : str, Optional
-            the API key for the vLLM server. If None, will use the key in os.environ['OPENROUTER_API_KEY'].
-        base_url : str, Optional
-            the base url for the vLLM server. 
-        config : LLMConfig
-            the LLM configuration.
-        """
-        self.api_key = api_key
-        if self.api_key is None:
-            self.api_key = os.getenv("OPENROUTER_API_KEY")
-        super().__init__(model, self.api_key, base_url, config, **kwrs)
-
-    def _format_response(self, response: Any) -> Dict[str, str]:
-        """
-        This method format the response from OpenAI API to a dict with keys "type" and "data".
-
-        Parameters:
-        ----------
-        response : Any
-            the response from OpenAI-compatible API. Could be a dict, generator, or object.
-        """
-        if isinstance(response, self.ChatCompletionChunk):
-            if hasattr(response.choices[0].delta, "reasoning") and getattr(response.choices[0].delta, "reasoning") is not None:
-                chunk_text = getattr(response.choices[0].delta, "reasoning", "")
-                if chunk_text is None:
-                    chunk_text = ""
-                return {"type": "reasoning", "data": chunk_text}
-            else:
-                chunk_text = getattr(response.choices[0].delta, "content", "")
-                if chunk_text is None:
-                    chunk_text = ""
-                return {"type": "response", "data": chunk_text}
-
-        return {"reasoning": getattr(response.choices[0].message, "reasoning", ""),
-                "response": getattr(response.choices[0].message, "content", "")}
-
-
-class OpenAIVLMEngine(VLMEngine):
-    def __init__(self, model:str, config:VLMConfig=None, **kwrs):
-        """
-        The OpenAI API inference engine. Supports OpenAI models and OpenAI compatible servers:
-        - vLLM OpenAI compatible server (https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html)
-
-        For parameters and documentation, refer to https://platform.openai.com/docs/api-reference/introduction
-
-        Parameters:
-        ----------
-        model_name : str
-            model name as described in https://platform.openai.com/docs/models
-        config : VLMConfig, Optional
-            the VLM configuration. Must be a child class of VLMConfig.
-        """
-        if importlib.util.find_spec("openai") is None:
-            raise ImportError("OpenAI Python API library not found. Please install OpanAI (```pip install openai```).")
-        
-        from openai import OpenAI, AsyncOpenAI
-        self.client = OpenAI(**kwrs)
-        self.async_client = AsyncOpenAI(**kwrs)
-        self.model = model
-        self.config = config if config else BasicVLMConfig()
-        self.formatted_params = self._format_config()
-
-    def _format_config(self) -> Dict[str, Any]:
-        """
-        This method format the LLM configuration with the correct key for the inference engine. 
-        """
-        formatted_params = self.config.params.copy()
-        if "max_new_tokens" in formatted_params:
-            formatted_params["max_completion_tokens"] = formatted_params["max_new_tokens"]
-            formatted_params.pop("max_new_tokens")
-
-        return formatted_params
-
-    def chat(self, messages:List[Dict[str,str]], verbose:bool=False, stream:bool=False, messages_logger:MessagesLogger=None) -> Union[Dict[str, str], Generator[Dict[str, str], None, None]]:
-        """
-        This method inputs chat messages and outputs LLM generated text.
-
-        Parameters:
-        ----------
-        messages : List[Dict[str,str]]
-            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
-        verbose : bool, Optional
-            if True, VLM generated text will be printed in terminal in real-time.
-        stream : bool, Optional
-            if True, returns a generator that yields the output in real-time.
-        messages_logger : MessagesLogger, Optional
-            the message logger that logs the chat messages.
-
-        Returns:
-        -------
-        response : Union[Dict[str,str], Generator[Dict[str, str], None, None]]
-            a dict {"reasoning": <reasoning>, "response": <response>} or Generator {"type": <reasoning or response>, "data": <content>}
-        """
-        processed_messages = self.config.preprocess_messages(messages)
-
-        if stream:
-            def _stream_generator():
-                response_stream = self.client.chat.completions.create(
-                                        model=self.model,
-                                        messages=processed_messages,
-                                        stream=True,
-                                        **self.formatted_params
-                                    )
-                res_text = ""
-                for chunk in response_stream:
-                    if len(chunk.choices) > 0:
-                        chunk_text = chunk.choices[0].delta.content
-                        if chunk_text is not None:
-                            res_text += chunk_text
-                            yield chunk_text
-                        if chunk.choices[0].finish_reason == "length":
-                            warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
-
-                # Postprocess response
-                res_dict = self.config.postprocess_response(res_text)
-                # Write to messages log
-                if messages_logger:
-                    # replace images content with a placeholder "[image]" to save space
-                    for messages in processed_messages:
-                        if "content" in messages and isinstance(messages["content"], list):
-                            for content in messages["content"]:
-                                if isinstance(content, dict) and content.get("type") == "image_url":
-                                    content["image_url"]["url"] = "[image]"
-
-                    processed_messages.append({"role": "assistant",
-                                                "content": res_dict.get("response", ""),
-                                                "reasoning": res_dict.get("reasoning", "")})
-                    messages_logger.log_messages(processed_messages)
-
-            return self.config.postprocess_response(_stream_generator())
-
-        elif verbose:
-            response = self.client.chat.completions.create(
-                model=self.model,
-                messages=processed_messages,
-                stream=True,
-                **self.formatted_params
-            )
-            res = ''
-            for chunk in response:
-                if len(chunk.choices) > 0:
-                    if chunk.choices[0].delta.content is not None:
-                        res += chunk.choices[0].delta.content
-                        print(chunk.choices[0].delta.content, end="", flush=True)
-                    if chunk.choices[0].finish_reason == "length":
-                        warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
-
-            print('\n')
-
-        else:
-            response = self.client.chat.completions.create(
-                model=self.model,
-                messages=processed_messages,
-                stream=False,
-                **self.formatted_params
-            )
-            res = response.choices[0].message.content
-            
-        # Postprocess response
-        res_dict = self.config.postprocess_response(res)
-        # Write to messages log
-        if messages_logger:
-            # replace images content with a placeholder "[image]" to save space
-            for messages in processed_messages:
-                if "content" in messages and isinstance(messages["content"], list):
-                    for content in messages["content"]:
-                        if isinstance(content, dict) and content.get("type") == "image_url":
-                            content["image_url"]["url"] = "[image]"
-
-            processed_messages.append({"role": "assistant", 
-                                    "content": res_dict.get("response", ""), 
-                                    "reasoning": res_dict.get("reasoning", "")})
-            messages_logger.log_messages(processed_messages)
-
-        return res_dict
-    
-
-    async def chat_async(self, messages:List[Dict[str,str]], messages_logger:MessagesLogger=None) -> Dict[str,str]:
-        """
-        Async version of chat method. Streaming is not supported.
-        """
-        processed_messages = self.config.preprocess_messages(messages)
-
-        response = await self.async_client.chat.completions.create(
-            model=self.model,
-            messages=processed_messages,
-            stream=False,
-            **self.formatted_params
-        )
-        
-        if response.choices[0].finish_reason == "length":
-            warnings.warn("Model stopped generating due to context length limit.", RuntimeWarning)
-
-        res = response.choices[0].message.content
-        # Postprocess response
-        res_dict = self.config.postprocess_response(res)
-        # Write to messages log
-        if messages_logger:
-            # replace images content with a placeholder "[image]" to save space
-            for messages in processed_messages:
-                if "content" in messages and isinstance(messages["content"], list):
-                    for content in messages["content"]:
-                        if isinstance(content, dict) and content.get("type") == "image_url":
-                            content["image_url"]["url"] = "[image]"
-                            
-            processed_messages.append({"role": "assistant", 
-                                    "content": res_dict.get("response", ""), 
-                                    "reasoning": res_dict.get("reasoning", "")})
-            messages_logger.log_messages(processed_messages)
-
-        return res_dict
-    
+class VLLMVLMEngine(VLLMInferenceEngine, OpenAICompatibleVLMEngine):
+    """
+    vLLM OpenAI compatible server inference engine.
+    https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html
+
+    For parameters and documentation, refer to https://platform.openai.com/docs/api-reference/introduction
+
+    Parameters:
+    ----------
+    model_name : str
+        model name as shown in the vLLM server
+    api_key : str, Optional
+        the API key for the vLLM server.
+    base_url : str, Optional
+        the base url for the vLLM server. 
+    config : LLMConfig
+        the LLM configuration.
+    """
+    pass
+
+class OpenRouterVLMEngine(OpenRouterInferenceEngine, OpenAICompatibleVLMEngine):
+    """
+    OpenRouter OpenAI-compatible server inference engine.
+
+    Parameters:
+    ----------
+    model_name : str
+        model name as shown in the vLLM server
+    api_key : str, Optional
+        the API key for the vLLM server. If None, will use the key in os.environ['OPENROUTER_API_KEY'].
+    base_url : str, Optional
+        the base url for the vLLM server. 
+    config : LLMConfig
+        the LLM configuration.
+    """
+    pass
+
+class OpenAIVLMEngine(OpenAIInferenceEngine, VLMEngine):
     def get_ocr_messages(self, system_prompt:str, user_prompt:str, image:Image.Image, format:str='png', 
                          detail:str="high", few_shot_examples:List[FewShotExample]=None) -> List[Dict[str,str]]:
         """
@@ -1215,32 +257,20 @@ class OpenAIVLMEngine(VLMEngine):
         return output_messages
 
 
-class AzureOpenAIVLMEngine(OpenAIVLMEngine):
-    def __init__(self, model:str, api_version:str, config:VLMConfig=None, **kwrs):
-        """
-        The Azure OpenAI API inference engine.
-        For parameters and documentation, refer to 
-        - https://azure.microsoft.com/en-us/products/ai-services/openai-service
-        - https://learn.microsoft.com/en-us/azure/ai-services/openai/quickstart
-        
-        Parameters:
-        ----------
-        model : str
-            model name as described in https://platform.openai.com/docs/models
-        api_version : str
-            the Azure OpenAI API version
-        config : LLMConfig
-            the LLM configuration.
-        """
-        if importlib.util.find_spec("openai") is None:
-            raise ImportError("OpenAI Python API library not found. Please install OpanAI (```pip install openai```).")
-        
-        from openai import AzureOpenAI, AsyncAzureOpenAI
-        self.model = model
-        self.api_version = api_version
-        self.client = AzureOpenAI(api_version=self.api_version, 
-                                  **kwrs)
-        self.async_client = AsyncAzureOpenAI(api_version=self.api_version, 
-                                             **kwrs)
-        self.config = config if config else BasicVLMConfig()
-        self.formatted_params = self._format_config()
+class AzureOpenAIVLMEngine(AzureOpenAIInferenceEngine, OpenAIVLMEngine):
+    """
+    The Azure OpenAI API inference engine.
+    For parameters and documentation, refer to 
+    - https://azure.microsoft.com/en-us/products/ai-services/openai-service
+    - https://learn.microsoft.com/en-us/azure/ai-services/openai/quickstart
+    
+    Parameters:
+    ----------
+    model : str
+        model name as described in https://platform.openai.com/docs/models
+    api_version : str
+        the Azure OpenAI API version
+    config : LLMConfig
+        the LLM configuration.
+    """
+    pass

{vlm4ocr-0.4.0.dist-info → vlm4ocr-0.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vlm4ocr
-Version: 0.4.0
+Version: 0.4.1
 Summary: Python package and Web App for OCR with vision language models.
 License: MIT
 Author: Enshuo (David) Hsu
@@ -12,6 +12,7 @@ Classifier: Programming Language :: Python :: 3.12
 Provides-Extra: tesseract
 Requires-Dist: colorama (>=0.4.4)
 Requires-Dist: json-repair (>=0.30.0)
+Requires-Dist: llm-inference-engine (>=0.1.1,<0.2.0)
 Requires-Dist: pdf2image (>=1.16.0)
 Requires-Dist: pillow (>=10.0.0)
 Requires-Dist: pytesseract (>=0.3.13) ; extra == "tesseract"

{vlm4ocr-0.4.0.dist-info → vlm4ocr-0.4.1.dist-info}/RECORD

@@ -10,8 +10,8 @@ vlm4ocr/cli.py,sha256=qFFIynex4sQSmT9ryjO2fPrkWfMyy3Aefp0p99rD3lU,22741
 vlm4ocr/data_types.py,sha256=DAlMl6UsfajVs7tIVjl2E8hT8BQg2fcOW7SDG12uIaA,5922
 vlm4ocr/ocr_engines.py,sha256=aLoA5J29bbffJwoeHo_DgPr0EaR7Q6Z5PlSW6UhZ2hw,27044
 vlm4ocr/utils.py,sha256=nQhUskOze99wCVMKmvsen0dhq-9NdN4EPC_bdYfkjgA,13611
-vlm4ocr/vlm_engines.py,sha256=fcy8l4lfCD2K4BoguMY9tL8mBFX12uJtTbmpvX01MF8,54242
-vlm4ocr-0.4.0.dist-info/METADATA,sha256=LgxGesSSxykldjFJ3Vv621JzkklfidXYe698iIxl4Yk,710
-vlm4ocr-0.4.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vlm4ocr-0.4.0.dist-info/entry_points.txt,sha256=qzWUk_QTZ12cH4DLjjfqce89EAlOydD85dreRRZF3K4,44
-vlm4ocr-0.4.0.dist-info/RECORD,,
+vlm4ocr/vlm_engines.py,sha256=Rv8-QcOBJYgjjaBLMHfAanQr9aTbH7rSCtuf9dH3lTc,10298
+vlm4ocr-0.4.1.dist-info/METADATA,sha256=3Txes_WwZIVBKkla3-zhPdrUzq6VJDeWcH0WALJbvH4,763
+vlm4ocr-0.4.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vlm4ocr-0.4.1.dist-info/entry_points.txt,sha256=qzWUk_QTZ12cH4DLjjfqce89EAlOydD85dreRRZF3K4,44
+vlm4ocr-0.4.1.dist-info/RECORD,,