Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

from __future__ import annotations

import asyncio
import os
from dataclasses import asdict, dataclass
from typing import Any, Literal
Expand All @@ -22,7 +23,7 @@
import httpx

import openai
from livekit.agents import llm
from livekit.agents import llm, utils
from livekit.agents.inference.llm import LLMStream as _LLMStream
from livekit.agents.llm import ToolChoice, utils as llm_utils
from livekit.agents.llm.chat_context import ChatContext
Expand Down Expand Up @@ -157,6 +158,7 @@ def __init__(
),
),
)
self._prewarm_task: asyncio.Task[None] | None = None

@property
def model(self) -> str:
Expand Down Expand Up @@ -912,6 +914,26 @@ def chat(
extra_kwargs=extra,
)

def prewarm(self) -> None:
"""Pre-warm the HTTP connection pool to reduce first-request latency."""

async def _prewarm_impl() -> None:
try:
await self._client.get("/", cast_to=str)
except Exception:
pass

# Cancel any existing prewarm task before creating a new one
if self._prewarm_task is not None and not self._prewarm_task.done():
self._prewarm_task.cancel()

self._prewarm_task = asyncio.create_task(_prewarm_impl())

async def aclose(self) -> None:
"""Clean up resources including any pending prewarm tasks."""
if self._prewarm_task is not None:
await utils.aio.gracefully_cancel(self._prewarm_task)


class LLMStream(_LLMStream):
def __init__(
Expand Down
163 changes: 163 additions & 0 deletions tests/test_llm_prewarm.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,163 @@
"""
Test LLM prewarming functionality (Issue #3240).

This test suite verifies that the prewarm() method reduces first-request latency
by pre-establishing HTTP connections to the LLM service.
"""

from __future__ import annotations

import asyncio
import os
import time

import pytest

from livekit.agents import llm
from livekit.plugins import openai

pytestmark = pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set"
)

llm_model = "gpt-4o-mini"


@pytest.mark.asyncio
async def test_llm_prewarm_reduces_latency():
"""Test that prewarming reduces time to first token (TTFT).
This test verifies that calling prewarm() before making an LLM request
reduces the latency of the first request by pre-establishing the HTTP connection.
"""
# Test 1: WITHOUT prewarming
llm_no_prewarm = openai.LLM(model=llm_model)

chat_ctx = llm.ChatContext()
chat_ctx.add_message(role="user", content="Say 'test' in one word only")

start = time.perf_counter()
stream = llm_no_prewarm.chat(chat_ctx=chat_ctx)

# Measure time to first chunk
ttft_no_prewarm = 0
async for chunk in stream:
if chunk.delta and chunk.delta.content:
ttft_no_prewarm = time.perf_counter() - start
break

# Fully consume the stream to avoid leaks
async for _ in stream:
pass

await llm_no_prewarm.aclose()

# Test 2: WITH prewarming
llm_with_prewarm = openai.LLM(model=llm_model)
llm_with_prewarm.prewarm()

# Give the prewarm task a moment to establish the connection
await asyncio.sleep(0.3)

chat_ctx = llm.ChatContext()
chat_ctx.add_message(role="user", content="Say 'test' in one word only")

start = time.perf_counter()
stream = llm_with_prewarm.chat(chat_ctx=chat_ctx)

# Measure time to first chunk
ttft_with_prewarm = 0
async for chunk in stream:
if chunk.delta and chunk.delta.content:
ttft_with_prewarm = time.perf_counter() - start
break

# Fully consume the stream to avoid leaks
async for _ in stream:
pass

await llm_with_prewarm.aclose()

# Verify prewarming helped (should be at least slightly faster)
# We don't assert a specific improvement because network conditions vary,
# but we print the results for visibility
print("Prewarm Test Results:")
print(f" Without prewarm: {ttft_no_prewarm:.3f}s")
print(f" With prewarm: {ttft_with_prewarm:.3f}s")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for your contribution! I hope the livekit team reviews this and gets it in. Looks like a must. To facilitate review, I would recommend:

  • run ruff
  • move into a method everything from chat_ctx = llm.ChatContext() to await llm_no_prewarm.aclose(). It is likely that this can be reused for test_llm_prewarm of other providers (gemini, etc)
  • move into a constant the LLM you're choosing. In this case: gpt-4o-mini


if ttft_with_prewarm < ttft_no_prewarm:
improvement = ttft_no_prewarm - ttft_with_prewarm
improvement_pct = (improvement / ttft_no_prewarm) * 100
print(f"Improvement: {improvement:.3f}s ({improvement_pct:.1f}% faster)")
else:
print(" No improvement detected (network conditions may vary)")

# The test passes if both requests succeeded
# We don't strictly assert latency improvements due to network variability
assert ttft_no_prewarm > 0
assert ttft_with_prewarm > 0


@pytest.mark.asyncio
async def test_llm_prewarm_task_cleanup():
"""Test that prewarm task is properly cleaned up on aclose()."""
llm_instance = openai.LLM(model=llm_model)

# Start prewarming
llm_instance.prewarm()

# Verify task was created
assert llm_instance._prewarm_task is not None

# Close immediately (should cancel the prewarm task gracefully)
await llm_instance.aclose()

# Task should be completed or cancelled
assert llm_instance._prewarm_task.done() or llm_instance._prewarm_task.cancelled()


@pytest.mark.asyncio
async def test_llm_prewarm_idempotent():
"""Test that calling prewarm() multiple times doesn't cause issues."""
llm_instance = openai.LLM(model=llm_model)

# Call prewarm multiple times
llm_instance.prewarm()
first_task = llm_instance._prewarm_task

# Calling prewarm again should create a new task
llm_instance.prewarm()
second_task = llm_instance._prewarm_task

# Both tasks should exist
assert first_task is not None
assert second_task is not None

# Clean up - must wait for tasks to complete or aclose will leak
await llm_instance.aclose()


@pytest.mark.asyncio
async def test_llm_works_without_prewarm():
"""Test that LLM works normally even without calling prewarm()."""
llm_instance = openai.LLM(model=llm_model)

# Don't call prewarm() at all
chat_ctx = llm.ChatContext()
chat_ctx.add_message(role="user", content="Say 'hello' in one word")

stream = llm_instance.chat(chat_ctx=chat_ctx)

# Should still work fine
response_received = False
async for chunk in stream:
if chunk.delta and chunk.delta.content:
response_received = True
break

# Fully consume the stream to avoid leaks
async for _ in stream:
pass

await llm_instance.aclose()

assert response_received, "Should receive response even without prewarm"