import uuid import time import asyncio import weakref from concurrent.futures import ThreadPoolExecutor from .base_model_backend import BaseModelBackend from .chat_schemas import ChatCompletionRequest, Message, ContentItem from client.llm_client import LLMClient import aiohttp import base64 import logging from fastapi import HTTPException from typing import Union, List from services.memory_check import MemoryChecker import tiktoken class LlmClientBackend(BaseModelBackend): def __init__(self, model_config): super().__init__(model_config) self._client_pool = [] self._active_clients = {} self._pool_lock = asyncio.Lock() self.logger = logging.getLogger("api.llm") self.MAX_CONTEXT_LENGTH = model_config.get("max_context_length", 128) self.POOL_SIZE = model_config.get("pool_size", 2) self._inference_executor = ThreadPoolExecutor(max_workers=self.POOL_SIZE) self._active_tasks = weakref.WeakSet() self.memory_checker = MemoryChecker( host=self.config["host"], port=self.config["port"] ) self.tokenizer = tiktoken.get_encoding("cl100k_base") async def _parse_content(self, content: Union[str, List[ContentItem]], base64_images: list) -> str: text_parts = [] if isinstance(content, list): for item in content: if item.type == "text" and item.text: text_parts.append(item.text.strip()) elif item.type == "image_url" and item.image_url: url = item.image_url.get("url", "") if url.startswith("data:image"): base64_data = url.split(",", 1)[1] base64_images.append(base64_data) else: base64_str = await self.download_image(url) if base64_str: base64_images.append(base64_str) else: text_parts.append(str(content).strip()) return " ".join(text_parts).strip() async def _get_client(self, request): try: await asyncio.wait_for(self._pool_lock.acquire(), timeout=30.0) start_time = time.time() timeout = 30.0 retry_interval = 3 while True: if self._client_pool: client = self._client_pool.pop() self.logger.debug(f"Reusing client from pool | ID:{id(client)}") return client if len(self._active_clients) < self.POOL_SIZE: break for task in self._active_tasks: task.cancel() # Will interrupt the activated client inference self._pool_lock.release() await asyncio.sleep(retry_interval) await asyncio.wait_for(self._pool_lock.acquire(), timeout=timeout - (time.time() - start_time)) if "memory_required" in self.config: await self.memory_checker.check_memory(self.config["memory_required"]) self.logger.debug("Creating new LLM client") client = LLMClient( host=self.config["host"], port=self.config["port"] ) self._active_clients[id(client)] = client system_content = next( (m.content for m in request.messages if m.role == "system"), self.config.get("system_prompt", "You are a helpful assistant") ) parsed_prompt = await self._parse_content(system_content, []) loop = asyncio.get_event_loop() await loop.run_in_executor( None, lambda: client.setup( self.config["object"], { "model": self.config["model_name"], "response_format": self.config["response_format"], "input": self.config["input"], "enoutput": True, "max_token_len": request.max_tokens, "temperature": request.temperature, "top_p": request.top_p, "prompt": parsed_prompt, "b_video": True } ) ) return client except asyncio.TimeoutError: raise RuntimeError("Server busy, please try again later.") finally: if self._pool_lock.locked(): self._pool_lock.release() async def _release_client(self, client): async with self._pool_lock: self._client_pool.append(client) self.logger.debug(f"Returned client to pool | ID:{id(client)}") async def close(self): for task in self._active_tasks: task.cancel() if self._active_tasks: await asyncio.wait(self._active_tasks, timeout=2) for client in self._client_pool: client.exit() self._client_pool.clear() self._active_clients.clear() self._inference_executor.shutdown(wait=False) async def inference_stream(self, query: str, base64_images: list, request: ChatCompletionRequest): client = await self._get_client(request) task = asyncio.current_task() self._active_tasks.add(task) try: self.logger.debug(f"Starting inference | ClientID:{id(client)} Query length:{len(query)}") loop = asyncio.get_event_loop() for i, image_data in enumerate(base64_images): client.send_jpeg(image_data, object_type="vlm.jpeg.base64") sync_gen = client.inference_stream( query, object_type="llm.utf-8" ) while True: if task.cancelled(): client.stop_inference() break def get_next(): try: return next(sync_gen) except StopIteration: return None chunk = await loop.run_in_executor( self._inference_executor, get_next ) if chunk is None: break yield chunk except asyncio.CancelledError: self.logger.warning("Inference task cancelled, stopping...") client.stop_inference() raise except Exception as e: self.logger.error(f"Inference error: {str(e)}") yield f"[ERROR: {str(e)}]" finally: self._active_tasks.discard(task) await self._release_client(client) self.logger.debug(f"Inference stopped | ClientID:{id(client)}") def _count_tokens(self, text: str) -> int: """Count the number of tokens in a given text.""" return len(self.tokenizer.encode(text)) def _truncate_history(self, request: ChatCompletionRequest) -> List[Message]: messages = request.messages if not messages: return [] last_msg = messages[-1] if isinstance(last_msg.content, list): msg_length = 0 for item in last_msg.content: if item.type == "text": msg_length += self._count_tokens(item.text) else: msg_length = self._count_tokens(last_msg.content) if last_msg.role in ("user", "assistant"): msg_length += 3 if msg_length > self.MAX_CONTEXT_LENGTH: return [] return [last_msg] async def download_image(self, url): try: async with aiohttp.ClientSession() as session: async with session.get(url) as response: if response.status == 200: image_data = await response.read() return base64.b64encode(image_data).decode('utf-8') self.logger.error(f"Image download failed, status code: {response.status}") return None except Exception as e: self.logger.error(f"Image download error: {str(e)}") return None async def generate(self, request: ChatCompletionRequest): try: truncated_messages = self._truncate_history(request) if not truncated_messages: raise HTTPException( status_code=400, detail="The input content exceeds the maximum length supported by the model." ) query_lines = [] base64_images = [] system_prompt = "" for m in truncated_messages: if m.role == "system": system_content = await self._parse_content(m.content, base64_images) system_prompt += f"{system_content}\n" continue message_content = await self._parse_content(m.content, base64_images) if message_content: query_lines.append(message_content) final_query = [] if system_prompt: final_query.append(system_prompt.strip()) if base64_images: pass final_query.append("\n".join(query_lines)) query = "\n\n".join(filter(None, final_query)) self.logger.debug( f"Processed query | System prompt: {len(system_prompt)} chars | " f"Images: {len(base64_images)} | Dialogue lines: {len(query_lines)}" ) if request.stream: async def chunk_generator(): try: async for chunk in self.inference_stream(query, base64_images, request): yield { "id": f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion.chunk", "created": int(time.time()), "model": request.model, "choices": [{ "index": 0, "delta": {"content": chunk}, "finish_reason": None }] } # Add normal completion marker yield { "object": "chat.completion.chunk", "choices": [{ "delta": {}, "finish_reason": "stop" }] } except Exception as e: self.logger.error(f"Stream generation error: {str(e)}") yield { "error": { "message": f"Stream generation failed: {str(e)}", "type": "api_error" } } yield {"choices": [{"delta": {}, "finish_reason": "stop"}]} raise return chunk_generator() else: full_response = "" async for chunk in self.inference_stream(query, base64_images, request): full_response += chunk return { "id": f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", "created": int(time.time()), "model": request.model, "choices": [{ "message": { "role": "assistant", "content": full_response } }] } except RuntimeError as e: self.logger.error(f"Connection error: {str(e)}") raise HTTPException( status_code=400, detail=f"Model service connection failed: {str(e)}" )