fix(llms/vllm): honor VLLM_BASE_URL instead of always using localhost (#6897)
This commit is contained in:
@@ -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
@@ -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):
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user