tests/integration/vllm_vision_local.py
Rama main · réplica de lectura del repositorio original.
"""
vLLM Vision Inference Test — Local Files
---------------------------------------------------
Tests the base64 local file path end-to-end.
Images are read from disk, encoded as data URIs, and
sent through the full pipeline:
encode → SDK → Samba → file_id → Redis → hydrate → vLLM
No network image fetching — purely local files.
"""
import base64
import json
import os
import time
from config_orc_fc import config
from dotenv import load_dotenv
from projectdavid import (
ContentEvent,
DecisionEvent,
Entity,
ReasoningEvent,
ToolCallRequestEvent,
)
load_dotenv()
# ------------------------------------------------------------------
# ANSI Colors
# ------------------------------------------------------------------
CYAN = "\033[96m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
RED = "\033[91m"
GREY = "\033[90m"
MAGENTA = "\033[95m"
RESET = "\033[0m"
# ------------------------------------------------------------------
# Config
# ------------------------------------------------------------------
BASE_URL = os.getenv("BASE_URL", "http://localhost:9000")
API_KEY = os.getenv("ENTITIES_API_KEY")
ASSISTANT_ID = config.get("assistant_id")
MODEL_ID = "vllm/Qwen/Qwen2.5-VL-3B-Instruct"
VLLM_BASE_URL = os.getenv("VLLM_BASE_URL", "http://vllm_server:8000")
VISION_PROMPT = (
"These are two locally encoded images. "
"Describe what you see in each one — include colours, shapes, and any text visible. "
"Then clearly state what is different between them."
)
# ------------------------------------------------------------------
# Local image encoder
# ------------------------------------------------------------------
def encode_image(path: str) -> str:
"""Read a local image file and return a base64 data URI."""
ext = os.path.splitext(path)[1].lower().lstrip(".")
mime_map = {
"jpg": "image/jpeg",
"jpeg": "image/jpeg",
"png": "image/png",
"webp": "image/webp",
"gif": "image/gif",
}
mime = mime_map.get(ext, "image/jpeg")
with open(path, "rb") as f:
b64 = base64.b64encode(f.read()).decode("utf-8")
return f"data:{mime};base64,{b64}"
# ------------------------------------------------------------------
# Locate test images
# These ship alongside this test script.
# Adjust paths if running from a different working directory.
# ------------------------------------------------------------------
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
IMAGE_1 = os.path.join(SCRIPT_DIR, "test_local_1.jpg")
IMAGE_2 = os.path.join(SCRIPT_DIR, "test_local_2.jpg")
for path in (IMAGE_1, IMAGE_2):
if not os.path.exists(path):
raise FileNotFoundError(
f"Test image not found: {path}\n"
f"Run the image generator script first, or place test_local_1.jpg "
f"and test_local_2.jpg in the same directory as this script."
)
print(f"\n{CYAN}[▶] Encoding local images...{RESET}")
b64_image_1 = encode_image(IMAGE_1)
b64_image_2 = encode_image(IMAGE_2)
print(f"{GREEN}[✓] Image 1: {IMAGE_1} ({len(b64_image_1)} chars){RESET}")
print(f"{GREEN}[✓] Image 2: {IMAGE_2} ({len(b64_image_2)} chars){RESET}")
# ------------------------------------------------------------------
# Multimodal payload — base64 data URIs, no network fetching
# ------------------------------------------------------------------
payload_content = [
{
"type": "text",
"text": VISION_PROMPT,
},
{
"type": "image_url",
"image_url": {"url": b64_image_1},
},
{
"type": "image_url",
"image_url": {"url": b64_image_2},
},
]
# ------------------------------------------------------------------
# SDK Init
# ------------------------------------------------------------------
client = Entity(base_url=BASE_URL, api_key=API_KEY)
# ------------------------------------------------------------------
# Thread + Message Setup
# ------------------------------------------------------------------
print(f"\n{CYAN}[▶] Setting up thread and uploading images to Samba...{RESET}")
thread = client.threads.create_thread()
message = client.messages.create_message(
thread_id=thread.id,
assistant_id=ASSISTANT_ID,
role="user",
content=payload_content,
)
print(f"{GREEN}[✓] Thread: {thread.id}{RESET}")
print(f"{GREEN}[✓] Message: {message.id}{RESET}")
print(f"{GREEN}[✓] Attachments: {message.attachments}{RESET}")
if len(message.attachments) != 2:
print(
f"{RED}[!] Expected 2 attachments, got {len(message.attachments)} — check SDK upload path{RESET}"
)
else:
print(f"{GREEN}[✓] Both local images uploaded and persisted as file_ids{RESET}")
# ------------------------------------------------------------------
# Hydration sanity check
# ------------------------------------------------------------------
formatted = client.messages.get_formatted_messages(thread.id)
user_msgs = [m for m in formatted if m.get("role") == "user"]
last_user = user_msgs[-1] if user_msgs else {}
content = last_user.get("content", "")
if isinstance(content, list):
image_blocks = [b for b in content if b.get("type") == "image"]
print(f"{GREEN}[✓] Hydrated image blocks: {len(image_blocks)}{RESET}\n")
else:
print(
f"{YELLOW}[!] Content is plain string — hydration did not produce image blocks{RESET}"
)
print(f"{YELLOW} content preview: {str(content)[:120]}{RESET}\n")
# ------------------------------------------------------------------
# Run + Stream Setup
# ------------------------------------------------------------------
run = client.runs.create_run(assistant_id=ASSISTANT_ID, thread_id=thread.id)
stream = client.synchronous_inference_stream
stream.setup(
thread_id=thread.id,
assistant_id=ASSISTANT_ID,
message_id=message.id,
run_id=run.id,
)
# ------------------------------------------------------------------
# Stream Loop
# ------------------------------------------------------------------
print(f"{CYAN}[▶] MODEL: {MODEL_ID}{RESET}")
print(f"{CYAN}[▶] VLLM URL: {VLLM_BASE_URL}{RESET}")
print(f"{CYAN}[▶] PROMPT: {VISION_PROMPT}{RESET}\n")
print("-" * 60)
last_tick = time.perf_counter()
global_start = last_tick
content_started = False
try:
for event in stream.stream_events(
model=MODEL_ID,
meta_data={"vllm_base_url": VLLM_BASE_URL},
):
now = time.perf_counter()
delta = now - last_tick
last_tick = now
color = {
ContentEvent: GREEN,
ToolCallRequestEvent: YELLOW,
ReasoningEvent: CYAN,
DecisionEvent: MAGENTA,
}.get(type(event), RESET)
if isinstance(event, ContentEvent):
if not content_started:
print(f"\n{GREEN}[Assistant]{RESET} ", end="", flush=True)
content_started = True
print(event.content, end="", flush=True)
else:
if content_started:
print()
content_started = False
print(
f"{GREY}[{delta:+.4f}s]{RESET} "
f"| {color}{event.__class__.__name__:<25}{RESET} "
f"| {json.dumps(event.to_dict())}"
)
except Exception as e:
if content_started:
print()
print(f"\n{RED}[ERROR] {e}{RESET}")
finally:
if content_started:
print()
total = time.perf_counter() - global_start
print(f"\n{YELLOW}{'=' * 50}")
print(f" TOTAL: {total:.4f}s")
print(f"{'=' * 50}{RESET}\n")