llm_client.py

import os
import json
import time
import queue
import base64
import urllib3
import requests
import threading
from PIL import Image
from io import BytesIO
from enum import Enum, auto
from urllib.parse import urlparse
from dataclasses import dataclass, field
from argparse import ArgumentParser
from typing import Callable, List, Optional

@dataclass
class Endpoint:
    store_name: str                             # Name of the endpoint, used to store the model in the solutions
    model_name: str                             # Model name that is used in the api request
    key: str = field(repr=False)                # API key (if required); never include it in logs/repr
    url: str                                    # URL of the endpoint
    _context_size: Optional[float] = None       # kilo-number of tokens
    _publication_date: Optional[str] = None     # ISO-short date, like 2025-09-19
    _quantization_level: Optional[int] = None   # number of bits per weight
    
    def get_dict(self) -> dict:
        return {
            "store_name": self.store_name,
            "model_name": self.model_name,
            "key": self.key,
            "url": self.url,
            "_context_size": self._context_size,
            "_publication_date": self._publication_date,
            "_quantization_level": self._quantization_level
        }


def load_endpoint_file(
    path: str,
    store_name: Optional[str] = None,
    model_name: Optional[str] = None,
) -> Endpoint:
    """Load an endpoint file, filling omitted names from command-line values."""
    with open(path, "r", encoding="utf-8") as file:
        endpoint_data = json.load(file)

    endpoint_data = dict(endpoint_data)
    if store_name is not None:
        endpoint_data.setdefault("store_name", store_name)
    if model_name is not None:
        endpoint_data.setdefault("model_name", model_name)

    missing = [
        name
        for name in ("key", "url", "store_name", "model_name")
        if name not in endpoint_data
        or (name != "key" and not endpoint_data.get(name))
    ]
    if missing:
        missing_options = [
            f"--{name}" for name in missing if name in ("store_name", "model_name")
        ]
        hint = (
            f"; provide {' and '.join(missing_options)} on the command line"
            if missing_options
            else ""
        )
        raise ValueError(
            f"Endpoint file {path} is missing required "
            f"{', '.join(missing)}{hint}."
        )

    allowed_fields = {
        "store_name",
        "model_name",
        "key",
        "url",
        "_context_size",
        "_publication_date",
        "_quantization_level",
    }
    return Endpoint(
        **{key: value for key, value in endpoint_data.items() if key in allowed_fields}
    )


def get_llm_url_stub(endpoint: Endpoint) -> str:
    """Get the base URL for the LLM API"""
    return urllib3.util.url.parse_url(endpoint.url)._replace(path='').url

def get_openai_models_url(endpoint: Endpoint) -> str:
    """Derive the models URL without discarding a provider-specific API prefix."""
    parsed = urlparse(endpoint.url)
    path = parsed.path.rstrip("/")
    chat_suffix = "/chat/completions"
    if path.endswith(chat_suffix):
        path = path[:-len(chat_suffix)] + "/models"
    else:
        path = path + "/models"
    return parsed._replace(path=path, params="", query="", fragment="").geturl()

def ollama_api_delete(endpoint: dict) -> bool:
    api_base = get_llm_url_stub(endpoint)
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
    try:
        # first unload the model
        response = requests.post(f"{api_base}/api/generate", verify=False,
                                 json={"model": endpoint.model_name, "keep_alive": 0})
        # then delete the model
        response = requests.request("DELETE", f"{api_base}/api/delete", verify=False,
                                    headers={'Accept': 'application/json', 'Content-Type': 'application/json'},
                                    json={"model": endpoint.model_name})
        return response.status_code == 200
    except requests.RequestException:
        return False

def openai_api_list(endpoint) -> dict:
    # Read model list from an openai-api-compatible endpoint (/v1/models).
    # The ollama endpoint now returns names different from the model list on the console (only lowercase match):
    # to check existence of a model, compare only lowercase. Do not call this method directly, use ensure_model_available instead.

    # We intentionally allow self-signed dev servers
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

    try:
        headers = {}
        if getattr(endpoint, "key", None):
            headers["Authorization"] = f"Bearer {endpoint.key}"

        resp = requests.get(
            get_openai_models_url(endpoint),
            headers=headers,
            verify=False,
            timeout=5,
        )
        resp.raise_for_status()
        data = resp.json()

        # Map to the same return shape; details are unknown here
        return {
            m["id"]: {
                "parameter_size": "unknown",
                "quantization_level": "unknown",
            }
            for m in data.get("data", [])
        }
    except (requests.RequestException, ValueError, KeyError, TypeError):
        # Both strategies failed; return empty so caller can handle it
        return {}


def is_ollama_endpoint(endpoint: Endpoint) -> bool:
    """Return whether the server exposes Ollama's native API."""
    api_base = get_llm_url_stub(endpoint)
    try:
        response = requests.get(f"{api_base}/api/version", verify=False, timeout=5)
        response.raise_for_status()
        data = response.json()
        return isinstance(data, dict) and isinstance(data.get("version"), str)
    except (requests.RequestException, ValueError, TypeError):
        return False


def ensure_model_available(endpoint: Endpoint, attempts: int = 3, fail_if_unavailable: bool = False) -> bool:
    api_base = get_llm_url_stub(endpoint)
    ollama_endpoint = is_ollama_endpoint(endpoint)
    for attempt in range(1, attempts + 1):
        models = openai_api_list(endpoint)
        # the endpoint now returns names different from the model listing on console; we must make a case insensitive match:
        models = {k.lower(): v for k, v in models.items()}
        if endpoint.model_name.lower() in models: return True
        print(f"Model availability check failed for {endpoint.model_name} on {api_base} (attempt {attempt}/{attempts}).")
        if not ollama_endpoint:
            print(
                f"{api_base} is not an Ollama server; skipping model pull and "
                "letting the chat request validate the model."
            )
            return True
        ollama_pull(endpoint)
        time.sleep(1)

    if fail_if_unavailable:
        raise RuntimeError(
            f"Could not verify endpoint {endpoint.url} with model {endpoint.model_name}. "
            "Check that /v1/models is reachable or that the model id is correct."
        )

    print(
        f"Model {endpoint.model_name} is still not verified on {api_base} after {attempts} attempts. "
        "Continuing and letting the chat request validate the model."
    )
    return False


def ensure_model_listed(endpoint: Endpoint) -> None:
    ensure_model_available(endpoint, attempts=3, fail_if_unavailable=False)

def ollama_pull(endpoint: Endpoint) -> dict:
    # Try to pull the model from the endpoint. If that does not work, we simply return.
    # Failure can be due to the model already being present, network issues, etc.,
    # we try to move on anyway.

    # pull the model if it is not available
    api_base = get_llm_url_stub(endpoint)
    print(f"Model {endpoint.model_name} is not available on server {api_base}. Pulling the model...")
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
    try:
        response = requests.request("POST", f"{api_base}/api/pull", verify=False,
                                    headers={'Accept': 'application/json', 'Content-Type': 'application/json'},
                                    json={"model": endpoint.model_name, "stream": False})
        if response.status_code != 200:
            print(f"Failed to pull model {endpoint.model_name} from server {api_base}. Status code: {response.status_code}, Response: {response.text}")
            return endpoint
        data = response.json()
        if data.get("error", False):
            print(f"Error pulling Model {endpoint.model_name} from server {api_base}.")
        else:
            print(f"Model {endpoint.model_name} is now available on server {api_base}.")
    except requests.RequestException as e:
        print(f"Error during model pull request: {e}")
    return endpoint

def hex2base64(hex_string) -> str:
    return base64.b64encode(bytes.fromhex(hex_string)).decode('utf-8')


def _extract_reasoning_tokens(usage: dict) -> Optional[int]:
    if not isinstance(usage, dict):
        return None
    detail_candidates = [
        usage.get("completion_tokens_details"),
        usage.get("output_tokens_details"),
        usage.get("reasoning"),
        usage.get("details"),
    ]
    for details in detail_candidates:
        if not isinstance(details, dict):
            continue
        for key in ("reasoning_tokens", "reasoning"):
            value = details.get(key)
            if isinstance(value, int):
                return value
    value = usage.get("reasoning_tokens")
    if isinstance(value, int):
        return value
    return None


def _normalize_usage(usage: dict, fallback_total_tokens: int = 0) -> dict:
    if not isinstance(usage, dict):
        usage = {}

    prompt_tokens = usage.get("prompt_tokens")
    if not isinstance(prompt_tokens, int):
        prompt_tokens = None

    completion_tokens = usage.get("completion_tokens")
    if not isinstance(completion_tokens, int):
        completion_tokens = usage.get("output_tokens")
    if not isinstance(completion_tokens, int):
        completion_tokens = None

    total_tokens = usage.get("total_tokens")
    if not isinstance(total_tokens, int):
        total_tokens = None

    reasoning_tokens = _extract_reasoning_tokens(usage)

    if total_tokens is None and prompt_tokens is not None and completion_tokens is not None:
        total_tokens = prompt_tokens + completion_tokens
    if completion_tokens is None and total_tokens is not None and prompt_tokens is not None:
        completion_tokens = total_tokens - prompt_tokens
    if total_tokens is None:
        total_tokens = fallback_total_tokens

    return {
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "reasoning_tokens": reasoning_tokens,
        "total_tokens": total_tokens,
    }


RATE_LIMIT_RETRY_SECONDS = 10


def _post_with_rate_limit_retry(endpoint: Endpoint, **request_kwargs):
    """Retry chat requests after a fixed cooldown when the provider returns 429."""
    while True:
        response = requests.post(endpoint.url, **request_kwargs)
        if response.status_code != 429:
            return response

        print(
            f"Rate limited by {get_llm_url_stub(endpoint)} (HTTP 429). "
            f"Waiting {RATE_LIMIT_RETRY_SECONDS} seconds before retrying..."
        )
        response.close()
        time.sleep(RATE_LIMIT_RETRY_SECONDS)


def openai_api_chat(
    endpoint: Endpoint,
    prompt: str = 'Hello World',
    base64_image: str = None,
    temperature: float = 0.0,
    max_tokens: int = 32768, # thats large and it requires that you set the context length in llm to 65536
    stream: bool = True,
    system_message: str = "You are a helpful assistant",
    tools: list = None,
    response_format: dict = None,
    return_response_json: bool = False,
    think = False,
    no_think = False
) -> tuple:
    """
    Function to interact with the LLM API for chat completions.
    
    Args:
        endpoint (dict): Dictionary containing endpoint information.
        prompt (str): The prompt to send to the model.
        base64_image (str): Base64 encoded image string (optional).
        temperature (float): Temperature for randomness in response.
        max_tokens (int): Maximum number of tokens for the response.
        
    Returns:
        tuple: A tuple containing the model's response, total tokens used, and tokens per second.
    """

    # Disable SSL warnings
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

    # Prepare the API endpoint URL
    stoptokens = ["[/INST]", "<|im_end|>", "<|end_of_turn|>", "<|eot_id|>", "<|end_header_id|>", "<EOS_TOKEN>", "</s>", "<|end|>"]

    # Set headers and payload
    headers = {'Content-Type': 'application/json', 'Accept': 'application/json'}
    if endpoint.key:
        headers['Authorization'] = f'Bearer {endpoint.key}'

    modelname = endpoint.model_name
    messages = []
    messages.append({"role": "system", "content": system_message})
    
    # special requirements of certain models
    if modelname.startswith("o1") or modelname.startswith("gpt-o1"): temperature = 1.0
    #if modelname.startswith("qwen3"): temperature = 0.6
    #if modelname.startswith("qwen3.5"): temperature = 0.7
    if modelname.startswith("4o") or modelname.startswith("gpt-4o") or modelname.startswith("gpt-3.5"):
        # reduce number of stoptokes to 4
        stoptokens = ["[/INST]", "<|im_end|>", "<|end_of_turn|>", "<|eot_id|>"]

    if base64_image:
        image_type = "jpeg"
        #base64_magic = {"/9j/": "jpeg", "iVBO": "png", "Qk": "bmp", "R0lG": "gif", "SUkq": "tiff", "SUkr": "tiff", "TU0A": "tiff", "GkXf": "webp", "UklG": "webp"}
        base64_magic = {"/9j/": "jpeg", "iVBO": "png", "R0lG": "gif"} # only jpeg and png are allowed as data type; however all of the types above (but gif!) are supported by the API
        for magic, itype in base64_magic.items():
            if base64_image.startswith(magic):
                #print(f"Detected {itype} image")
                image_type = itype
                break
        # If this is a gif we must convert it to png
        if image_type == "gif":
            #print("Converting gif to png")
            image = Image.open(BytesIO(base64.b64decode(base64_image)))
            png_image = BytesIO()
            image.save(png_image, format="PNG")
            base64_image = base64.b64encode(png_image.getvalue()).decode('utf-8')
            image_type = "png"

        # Add the image to the message
        image_url_object = {"url": f"data:image/{image_type};base64,{base64_image}"}
        usermessage = {"role": "user", "content": [
            {"type": "text", "text": prompt},
            {"type": "image_url", "image_url": image_url_object}
        ]}
    else:
        usermessage = {"role": "user", "content": prompt}
    messages.append(usermessage)

    if modelname.startswith("o1") or modelname.startswith("4o"):
        stoptokens = []

    payload = {
        "model": modelname,
        "messages": messages,
        # "response_format": { "type": "text" }, # do NOT set this, some engines do not accept "text" as response format because it is not valid
        "temperature": temperature, # ollama default: 0.8
        "top_k": 20,    # reduces the probability of generating nonsense: high = more diverse, low = more focused; ollama default: 40
        "top_p": 0.8,  # works together with top_k: high = more diverse, low = more focused; ollama default: 0.9
        "min_p": 0,     # alternative to top_p: p is minimum probability for a token to be considered; ollama default: 0.0
        "presence_penalty": 1.5,
        "stream": stream
    }
    if tools:
        payload["tools"] = tools
    if response_format:
        payload["response_format"] = response_format
    if stream:
        payload["stream_options"] = {"include_usage": True}
    if len(stoptokens) > 0 and not modelname.startswith("o4"):
        payload["stop"] = stoptokens
    if modelname.startswith("o1") or modelname.startswith("o4"):
        payload["max_completion_tokens"] = max_tokens
    else:
        payload["max_tokens"] = max_tokens
    modelname_lower = modelname.lower()
    if no_think:
        payload["enable_thinking"] = False
        payload["reasoning_effort"] = "none"

    # use the endpoints array as failover mechanism
    response = None
    text_chunks = []
    usage = None
    thinking_not_suppressed = False
    read_timeout = 600 # seconds
    token_count = 0
    parsed_url = urlparse(endpoint.url)
    host = parsed_url.hostname or ""
    c0 = host[0] if host else "."
    #print(f"Calling model in strem mode: {stream}, payload: {json.dumps(payload)}")
    try:
        t0 = time.time()
        response = _post_with_rate_limit_retry(
            endpoint,
            headers=headers,
            json=payload,
            verify=False,
            stream=stream,
            timeout=(60, read_timeout) # (connect_timeout, read_timeout))
        )
        #print(f"Response status: {response.status_code}")
        response.raise_for_status()
        #print(f"Response headers: {response.headers}")
        if stream:
            #print("Response (stream): ", end="", flush=True)
            timeouttime = t0 + read_timeout
            for line in response.iter_lines(decode_unicode=True):
                if time.time() > timeouttime: break # we simply silently terminate the stream after the timeout
                if not line: continue
                #print(line)
                if line.startswith("data: "):
                    payload_line = line[len("data: "):].strip()
                    if payload_line == "[DONE]":
                        print() # end progress line
                        break
                    try:
                        evt = json.loads(payload_line)
                        evt_usage = evt.get("usage")
                        if isinstance(evt_usage, dict):
                            usage = evt_usage
                        choices = evt.get("choices", [])
                        if choices:
                            delta = choices[0].get("delta", {})
                            reasoning_token = delta.get("reasoning")
                            if no_think and not thinking_not_suppressed and isinstance(reasoning_token, str) and reasoning_token.strip():
                                thinking_not_suppressed = True
                                print(f"WARNING - THINKING NOT SUPPRESSED on {endpoint.url}")
                            #print(delta)
                            if "content" in delta:
                                # delta may have attributes content or reasoning. Take whatever is non-empty
                                token = delta.get("content") or reasoning_token
                                if token:
                                    text_chunks.append(token)
                                    token_count += 1
                                    #print(token, end="", flush=True)
                                    if token_count % 100 == 0:
                                        print(c0, end="", flush=True) # print a dot for each 10 tokens to show progress
                    except Exception:
                        pass # robust against json parse errors
        t1 = time.time()
    except requests.exceptions.ReadTimeout as e:
        raise Exception(f"Read timeout while calling {endpoint.url} (timeout=600s). "
                        f"The model may be slow or the server overloaded.") from e
    except requests.exceptions.RequestException as e:
        # print(f"Failed to access api: {e}")
        # Get the error message from the response
        body = ""
        if hasattr(e, "response") and e.response is not None:
            try:
                body = (e.response.text or "")[:800].replace("\n", " ")
            except Exception:
                body = ""
        raise Exception(f"API request failed to {endpoint.url}: {e} | Body: {body}") from e

    # Parse the response
    try:
        if stream:
            answer = "".join(text_chunks).strip()
            usage_summary = _normalize_usage(usage, fallback_total_tokens=len(text_chunks))
            total_tokens = usage_summary["total_tokens"]
            token_per_second = 0.0 if (t1 - t0) <= 0 else total_tokens / (t1 - t0)
            if not answer: print(f"Empty streamed response from the API at {endpoint.url}")
            response_json = None
        else:
            ctype = response.headers.get('Content-Type', '')
            text  = response.text or ''
            if not text.strip():
                raise Exception(f"Empty response body (status {response.status_code}) from {endpoint.url}")

            if 'json' not in ctype.lower():
                # possibly a html error page
                snippet = text[:800].replace('\n',' ')
                raise Exception(f"Non-JSON response (status {response.status_code}, Content-Type {ctype}): {snippet}")

            data = response.json()
            usage_summary = _normalize_usage(data.get('usage', {}))
            total_tokens = usage_summary["total_tokens"]
            token_per_second = total_tokens / (t1 - t0)
            #print(f"Total tokens: {total_tokens}, tokens per second: {token_per_second:.2f}")
            choices = data.get('choices', [])
            if len(choices) == 0:
                raise Exception("No response from the API: " + str(data))
            message = choices[0].get('message', {})
            reasoning_token = message.get("reasoning")
            if no_think and isinstance(reasoning_token, str) and reasoning_token.strip():
                thinking_not_suppressed = True
                print(f"WARNING - THINKING NOT SUPPRESSED on {endpoint.url}")
            answer = message.get('content', '')
            response_json = data
        if return_response_json:
            return answer, total_tokens, token_per_second, usage_summary, t1 - t0, response_json
        return answer, total_tokens, token_per_second, usage_summary, t1 - t0
    except json.JSONDecodeError as e:
        raise Exception(f"Failed to parse JSON response from the API: {e}")

busy_waiting_time = 1 # seconds


@dataclass
class Task:
    """
    Dataclass for task.
    This dataclass is used to represent a task that needs to be processed by a server.
    Each task has an ID and a dictionary of data that contains the actual task data.
    """
    id: str             # Unique identifier for the task
    description: str    # a short description of the task (for logging)
    prompt: str         # the prompt to be sent to the model
    base64_image: str   # the base64 encoded image to be sent to the model
    response_processing: Callable[['Response'], None] # a function to process the result
    think: bool = False         # use thinking settings
    no_think: bool = False      # use non-thinking settings

@dataclass
class Response:
    """
    Dataclass for response.
    This dataclass is used to represent a response from the server.
    Each response has a task ID and the result of the task processing.
    """
    task: Task
    result: str         # The result of the task processing
    total_tokens: int   # Total tokens used in the response
    token_per_second: float # Tokens per second used in the response
    duration_seconds: float = 0.0
    prompt_tokens: Optional[int] = None
    completion_tokens: Optional[int] = None
    reasoning_tokens: Optional[int] = None

@dataclass
class Server:
    """
    Dataclass for server.
    This dataclass is used to represent a server that can process tasks.
    Each server has an ID, an endpoint (URL), and a status that indicates whether the server is available or busy.
    The status is used to track the current task being processed by the server.
    """
    endpoint: Endpoint  # The endpoint of the server
    current_task: Task = None  # Track the current task being processed

class LoadBalancer:
    """
    LoadBalancer class for managing task distribution across multiple servers.
    This class is responsible for distributing tasks to available servers and managing their status.
    - It uses a queue to manage tasks and a list of servers to distribute the load.
    - It will only assign tasks to servers that are AVAILABLE.
    - It implements backpressure to prevent overloading the servers.
    - It will wait for a server to become available before assigning a new task.
    - It will also retry failed tasks after a short delay.
    - The status of each server is updated as tasks are assigned and completed.
    """
    def __init__(self, max_queue_size: int = 1000):
        self.servers = []
        self.task_queue = queue.Queue[Task](maxsize=max_queue_size)
        self.available_servers = queue.Queue[Server]()
        self.lock = threading.Lock()
        
    def add_server(self, server: Server):
        """Add a server to the load balancer"""
        self.servers.append(server)
        self.available_servers.put(server)
        print(f"Server {get_llm_url_stub(server.endpoint)} added to load balancer.")

    def add_task(self, task: Task):
        """Add a task to the processing queue with backpressure"""
        try:
            self.task_queue.put(task, block=True, timeout=1)
            return True
        except queue.Full:
            print("Task queue full - applying backpressure")
            return False
    
    def mark_server_available(self, server: Server):
        """Mark a server as available for new tasks"""
        with self.lock:
            server.current_task = None
            self.available_servers.put(server)
    
    def get_available_server(self, timeout: float = 10.0) -> Optional[Server]:
        """Get the next available server with timeout"""
        try:
            # Remove and return the next available server from the queue.
            # The only way the server gets available is when the task is finished
            # and the task assignes its server back to the available_servers.
            return self.available_servers.get(timeout=timeout)
        except queue.Empty:
            return None
    
    def assign_task_to_server(self, task: Task, server: Server):
        """Assign task to server and mark it as busy"""
        with self.lock:
            server.current_task = task
        
        threading.Thread(
            target=self.process_task_remote,
            args=(server,),
            daemon=True
        ).start()
    
    def process_task_remote(self, server: Server):
        """Process task on remote server"""
        task = server.current_task
        endpoint = server.endpoint
        try:
            #print(f"Processing task ID {task.id} on server {server.endpoint} with model {task.model}")
            answer, total_tokens, token_per_second, usage_summary, duration_seconds = openai_api_chat(
                endpoint,
                task.prompt, 
                base64_image=task.base64_image,
                think = task.think,
                no_think = task.no_think
            )
            # Call the response processing function
            response = Response(
                task,
                answer,
                total_tokens,
                token_per_second,
                duration_seconds=duration_seconds,
                prompt_tokens=usage_summary.get("prompt_tokens"),
                completion_tokens=usage_summary.get("completion_tokens"),
                reasoning_tokens=usage_summary.get("reasoning_tokens"),
            )
            task.response_processing(response)
            print(f"Processed {task.description}, on {server.endpoint.url} with model {endpoint.model_name} in {duration_seconds:.2f} seconds with {total_tokens} tokens ({token_per_second:.2f} tokens/sec)")
            
            # mark server available
            self.mark_server_available(server)
                
        except Exception as e:
            # write a stack trace to std out
            import traceback
            traceback.print_exc()
            # Log the error and mark server available
            error_msg = f"Failed to process task ID {task.id} on {server.endpoint}: {str(e)}"
            if hasattr(e, 'response'):
                try:
                    error_details = e.response.json()
                    error_msg += f" | API Response: {error_details}"
                except:
                    error_msg += f" | Raw Response: {e.response.text}"
            print(error_msg)
            # make server available again
            self.mark_server_available(server)
    
    def start_distribution(self):
        """Start the task distribution process"""
        def distributor():
            while True:
                task = self.task_queue.get()
                assigned = False
                
                while not assigned:
                    server = self.get_available_server()
                    if server:
                        self.assign_task_to_server(task, server)
                        assigned = True
                    else:
                        # All servers busy, wait and try again
                        time.sleep(busy_waiting_time)
                
                self.task_queue.task_done()
        
        # Start distributor thread
        threading.Thread(target=distributor, daemon=True).start()
    
    def wait_completion(self):
        """Wait for all tasks to be processed"""
        self.task_queue.join()
        # Wait for all servers to finish their current tasks
        print("Waiting for all servers to finish processing...")
        while any(s.current_task != None for s in self.servers):
            time.sleep(busy_waiting_time)
            print("Still waiting for servers to finish...")
            # print out the current status of all servers
            for server in self.servers:
                if server.current_task:
                    print(f"Server {server.endpoint.url} - Current task ID: {server.current_task.id}")

        print("All servers finished processing.")

def main():
    from llm_model_test import test_vision
    parser = ArgumentParser(description="Testing the LLM API.")
    parser.add_argument('--api_base', required=False, default='http://localhost:11434', help='API base URL for the LLM, default is http://localhost:11434')
    parser.add_argument('--endpoint', required=False, default='', help='Name of an <endpoint>.json file in the endpoints directory')
    parser.add_argument('--store_name', help='Storage name when the endpoint file omits store_name')
    parser.add_argument('--model_name', help='API model name when the endpoint file omits model_name')
    parser.add_argument('--model', required=False, default='llama3.2:latest', help='Name of the model to use, default is llama3.2:latest')
    parser.add_argument('--image', required=False, default=None, help='path to an image that shall be processed')
    parser.add_argument('--think', action='store_true', help='forward a "think" flag to compatible backends')
    parser.add_argument('--no_think', action='store_true', help='forward a "no_think" flag to compatible backends')
    
    # parse the arguments
    args = parser.parse_args()
    api_base = args.api_base.split(",") if "," in args.api_base else [args.api_base]
    endpoint_name = args.endpoint
    model_name = args.model
    image_path = args.image
    think = args.think
    no_think = args.no_think

    # load the endpoint file
    endpoints:List[Endpoint] = []
    if endpoint_name:
        print(f"Using endpoint {endpoint_name}")
        endpoint_path = os.path.join('endpoints', f"{endpoint_name}.json")
        print(f"Using endpoint file {endpoint_path}")
        if not os.path.exists(endpoint_path):
            raise Exception(f"Endpoint file {endpoint_path} does not exist.")
        endpoints = [
            load_endpoint_file(
                endpoint_path,
                store_name=args.store_name,
                model_name=args.model_name,
            )
        ]
    else:
        endpoints = [
            Endpoint(store_name=model_name, model_name=model_name, key="",
                     url=f"{api_stub}/v1/chat/completions") for api_stub in api_base
        ]

    # test if the endpoint is a multimodal model
    if test_vision(endpoints[0]):
        print("Endpoint is a multimodal model.")
    else:
        print("Endpoint is not a multimodal model.")

    # load the image, if a path is given
    base64_image = None
    if image_path:
        with open(image_path, "rb") as image_file:
            base64_image = base64.b64encode(image_file.read()).decode('utf-8')

    # access the LLM API
    models_dict = openai_api_list(endpoints[0])
    for (model, attr) in models_dict.items():
        print(f"Model: {model}: {attr}")
    try:
        if base64_image:
            answer, total_tokens, token_per_second, usage_summary, duration_seconds = openai_api_chat(
                endpoints[0],
                prompt="what is in the image",
                base64_image=base64_image,
                think=think,
                no_think=no_think,
            )
        else:
            answer, total_tokens, token_per_second, usage_summary, duration_seconds = openai_api_chat(
                endpoints[0],
                think=think,
                no_think=no_think,
            )
    except Exception as e:
        answer = f"Error: {str(e)}"
    print(answer)

if __name__ == "__main__":
    main()