| import os |
| import requests |
| from typing import List |
| from tenacity import retry, stop_after_attempt, wait_random_exponential |
| from models.Base import BaseModel |
| from openai import OpenAI |
|
|
|
|
| class VLLMModel(BaseModel): |
| def __init__(self, |
| model_id="local-vllm", |
| base_url="http://localhost:8041/v1", |
| api_key=None): |
| """ |
| model_id: 模型名称或标识 (通常不重要，vLLM 会忽略) |
| base_url: vLLM API 地址，例如 http://localhost:8000/v1 |
| api_key: 可选（vLLM 默认不验证） |
| """ |
| self.model_id = model_id |
|
|
| client = OpenAI( |
| api_key='EMPTY', |
| base_url=base_url, |
| ) |
|
|
| models = client.models.list() |
| self.model_id = models.data[0].id |
|
|
| self.base_url = base_url.rstrip("/") |
| self.api_key = api_key |
|
|
| |
| self.headers = { |
| "Content-Type": "application/json" |
| } |
| if api_key: |
| self.headers["Authorization"] = f"Bearer {api_key}" |
|
|
| |
| self.endpoint = f"{self.base_url}/chat/completions" |
|
|
| @retry(wait=wait_random_exponential(min=1, max=10), stop=stop_after_attempt(3)) |
| def generate(self, |
| messages: List, |
| temperature=0, |
| presence_penalty=0, |
| frequency_penalty=0, |
| max_tokens=5000) -> str: |
| """ |
| 调用本地 vLLM 推理接口 |
| """ |
| payload = { |
| "model": self.model_id, |
| "messages": messages, |
| "temperature": temperature, |
| "max_tokens": max_tokens, |
| "presence_penalty": presence_penalty, |
| "frequency_penalty": frequency_penalty, |
| } |
|
|
| import ipdb; ipdb.set_trace() |
| response = requests.post(self.endpoint, headers=self.headers, json=payload, timeout=3000) |
|
|
| if response.status_code != 200: |
| raise RuntimeError(f"vLLM request failed: {response.status_code}, {response.text}") |
|
|
| data = response.json() |
| if "choices" not in data or len(data["choices"]) == 0: |
| raise ValueError("No response choices returned from vLLM API.") |
|
|
| return data["choices"][0]["message"]["content"] |
|
|
|
|
