feat: call 方法

This commit is contained in:
李如威 2025-11-05 17:52:42 +08:00
parent 4372d7c52d
commit 370301cd15
6 changed files with 136 additions and 0 deletions

View File

@ -0,0 +1 @@
httpx==0.28.1

View File

0
src/sample_ai/flow.py Normal file
View File

0
src/sample_ai/main.py Normal file
View File

0
src/sample_ai/nodes.py Normal file
View File

135
src/sample_ai/utils.py Normal file
View File

@ -0,0 +1,135 @@
import httpx
import asyncio
import os
import json
VLLM_HTTP_KEY = os.environ.get("VLLM_HTTP_KEY")
VLLM_HTTP_URL = os.environ.get("VLLM_HTTP_URL")
VLLM_CHAT_MODEL_NAME = os.environ.get("VLLM_CHAT_MODEL_NAME")
def _headers():
h = {"Content-Type": "application/json"}
if VLLM_HTTP_KEY:
h["Authorization"] = f"Bearer {VLLM_HTTP_KEY}"
return h
def _merge_messages_and_prompt(messages: list[dict] = [], prompt: str = ""):
"""
messages历史对话和单条 prompt 合并成新的 OpenAI-style messages 列表
- 如果 messages None 或空则创建一个只包含 prompt user 消息 prompt 非空
- prompt 非空总是以一条 {"role":"user","content": prompt} 追加到 messages 末尾
"""
merged = []
if messages:
merged = [dict(x) for x in messages]
if prompt:
merged.append({"role": "user", "content": prompt})
return merged
def _get_content(resp: dict) -> str:
choices = resp.get("choices")
if not isinstance(choices, list) or not choices:
return ""
first = choices[0]
if not isinstance(first, dict):
return ""
content_obj = first.get("message") or first.get("delta")
if isinstance(content_obj, dict):
return content_obj.get("content", "")
return ""
async def call_llm(prompt: str, messages = [], max_tokens: int = 512, temperature: float = 0.0, timeout: int = 30):
"""简单对话"""
try:
async with httpx.AsyncClient(timeout=timeout) as client:
url = VLLM_HTTP_URL.rstrip("/") + "/v1/chat/completions"
body = {
"model": VLLM_CHAT_MODEL_NAME,
"messages": _merge_messages_and_prompt(messages, prompt),
"max_tokens": max_tokens,
"temperature": temperature,
}
res = await client.post(url, headers=_headers(), json=body)
res.raise_for_status()
data = res.json()
return _get_content(data)
except Exception as e:
print(f"call_llm[ERROR]: {e}")
return ""
async def call_llm_stream(prompt: str, messages = [], max_tokens: int = 512, temperature: float = 0.0, timeout: int = 60):
"""
流式对话
- 使用: async for chunk call_llm_stream("prompt"):
"""
try:
async with httpx.AsyncClient(timeout=timeout) as client:
url = (VLLM_HTTP_URL or "").rstrip("/") + "/v1/chat/completions"
body = {
"model": VLLM_CHAT_MODEL_NAME,
"messages": _merge_messages_and_prompt(messages or [], prompt),
"max_tokens": max_tokens,
"temperature": temperature,
"stream": True,
}
async with client.stream("POST", url, headers=_headers(), json=body) as resp:
resp.raise_for_status()
async for raw_line in resp.aiter_lines():
if not raw_line:
continue
line = raw_line.strip()
if line.startswith("data:"):
payload = line[len("data:"):].strip()
data = {}
if payload in ("[DONE]", ""):
break
try:
data = json.loads(payload)
except Exception:
yield payload
continue
# print(data)
yield _get_content(data)
continue
yield line
except Exception as e:
print(f"call_llm_stream[ERROR]: {e}")
return
if __name__ == "__main__":
# 设置
os.environ.setdefault("VLLM_HTTP_URL", "http://localhost:8022")
os.environ.setdefault("VLLM_HTTP_KEY", "sk-local-827ccb0eea8a706c4c34a16891f84e7b")
os.environ.setdefault("VLLM_CHAT_MODEL_NAME", "Qwen2.5-0.5B-Instruct")
# 读取
VLLM_HTTP_KEY = os.environ.get("VLLM_HTTP_KEY")
VLLM_HTTP_URL = os.environ.get("VLLM_HTTP_URL")
VLLM_CHAT_MODEL_NAME = os.environ.get("VLLM_CHAT_MODEL_NAME")
test_prompt = "你好,你是谁,有什么功能,中文回复"
print("--- 普通对话 ---")
text = asyncio.run(call_llm(test_prompt))
print(text)
print("--- 对话结束 ---\n")
print("--- 流式对话 ---")
async def _test_stream():
try:
async for chunk in call_llm_stream(test_prompt):
print(chunk, end="", flush=True)
print("\n--- 流式结束 ---")
except Exception as e:
print(f"_test_stream[ERROR]: {e}")
asyncio.run(_test_stream())