fix(llms/vllm): honor VLLM_BASE_URL instead of always using localhost (#6897)

This commit is contained in:
Kartik
2026-08-14 16:49:18 +05:30
committed by GitHub
parent 96d45b78c7
commit c883f52130
3 changed files with 57 additions and 3 deletions
+2 -1
View File
@@ -1,3 +1,4 @@
import os
from typing import Optional from typing import Optional
from mem0.configs.llms.base import BaseLlmConfig from mem0.configs.llms.base import BaseLlmConfig
@@ -53,4 +54,4 @@ class VllmConfig(BaseLlmConfig):
) )
# vLLM-specific parameters # vLLM-specific parameters
self.vllm_base_url = vllm_base_url or "http://localhost:8000/v1" self.vllm_base_url = vllm_base_url or os.getenv("VLLM_BASE_URL") or "http://localhost:8000/v1"
+1 -2
View File
@@ -37,8 +37,7 @@ class VllmLLM(LLMBase):
self.config.model = "Qwen/Qwen2.5-32B-Instruct" self.config.model = "Qwen/Qwen2.5-32B-Instruct"
self.config.api_key = self.config.api_key or os.getenv("VLLM_API_KEY") or "vllm-api-key" self.config.api_key = self.config.api_key or os.getenv("VLLM_API_KEY") or "vllm-api-key"
base_url = self.config.vllm_base_url or os.getenv("VLLM_BASE_URL") self.client = OpenAI(api_key=self.config.api_key, base_url=self.config.vllm_base_url)
self.client = OpenAI(api_key=self.config.api_key, base_url=base_url)
def _parse_response(self, response, tools): def _parse_response(self, response, tools):
""" """
+54
View File
@@ -4,6 +4,7 @@ import pytest
from mem0 import AsyncMemory, Memory from mem0 import AsyncMemory, Memory
from mem0.configs.llms.base import BaseLlmConfig from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.vllm import VllmConfig
from mem0.llms.vllm import VllmLLM from mem0.llms.vllm import VllmLLM
@@ -15,6 +16,59 @@ def mock_vllm_client():
yield mock_client yield mock_client
def test_base_url_explicit_arg_wins_over_env(mock_vllm_client, monkeypatch):
monkeypatch.setenv("VLLM_BASE_URL", "http://from-env:8000/v1")
config = VllmConfig(vllm_base_url="http://from-arg:8000/v1")
VllmLLM(config)
assert config.vllm_base_url == "http://from-arg:8000/v1"
def test_base_url_env_wins_over_default(mock_vllm_client, monkeypatch):
monkeypatch.setenv("VLLM_BASE_URL", "http://from-env:8000/v1")
config = VllmConfig()
VllmLLM(config)
assert config.vllm_base_url == "http://from-env:8000/v1"
def test_base_url_default_when_neither_set(mock_vllm_client, monkeypatch):
monkeypatch.delenv("VLLM_BASE_URL", raising=False)
config = VllmConfig()
VllmLLM(config)
assert config.vllm_base_url == "http://localhost:8000/v1"
def test_openai_client_receives_explicit_base_url(monkeypatch):
monkeypatch.delenv("VLLM_BASE_URL", raising=False)
with patch("mem0.llms.vllm.OpenAI") as mock_openai:
VllmLLM(VllmConfig(vllm_base_url="http://from-arg:8000/v1"))
assert mock_openai.call_args.kwargs["base_url"] == "http://from-arg:8000/v1"
def test_openai_client_receives_base_url_from_env(monkeypatch):
monkeypatch.setenv("VLLM_BASE_URL", "http://from-env:8000/v1")
with patch("mem0.llms.vllm.OpenAI") as mock_openai:
VllmLLM(VllmConfig())
base_url = mock_openai.call_args.kwargs["base_url"]
assert base_url == "http://from-env:8000/v1"
assert base_url != "http://localhost:8000/v1"
def test_openai_client_receives_base_url_from_env_via_base_llm_config(monkeypatch):
"""Original #6256 repro: BaseLlmConfig (not VllmConfig) plus VLLM_BASE_URL only."""
monkeypatch.setenv("VLLM_BASE_URL", "http://remote-vllm:8000/v1")
with patch("mem0.llms.vllm.OpenAI") as mock_openai:
VllmLLM(BaseLlmConfig(model="Qwen/Qwen2.5-32B-Instruct"))
base_url = mock_openai.call_args.kwargs["base_url"]
assert base_url == "http://remote-vllm:8000/v1"
assert base_url != "http://localhost:8000/v1"
def test_generate_response_without_tools(mock_vllm_client): def test_generate_response_without_tools(mock_vllm_client):
config = BaseLlmConfig(model="Qwen/Qwen2.5-32B-Instruct", temperature=0.7, max_tokens=100, top_p=1.0) config = BaseLlmConfig(model="Qwen/Qwen2.5-32B-Instruct", temperature=0.7, max_tokens=100, top_p=1.0)
llm = VllmLLM(config) llm = VllmLLM(config)