185 lines
6.4 KiB
Python
185 lines
6.4 KiB
Python
from venv import logger
|
||
import httpx
|
||
import asyncio
|
||
import os
|
||
import json
|
||
from dotenv import load_dotenv
|
||
load_dotenv()
|
||
|
||
MAIN_HTTP_URL = os.environ.get("MAIN_HTTP_URL")
|
||
MAIN_HTTP_KEY = os.environ.get("MAIN_HTTP_KEY")
|
||
MAIN_MODEL_NAME = os.environ.get("MAIN_MODEL_NAME")
|
||
|
||
EMBEDDING_HTTP_URL = os.environ.get("EMBEDDING_HTTP_URL")
|
||
EMBEDDING_HTTP_KEY = os.environ.get("EMBEDDING_HTTP_KEY")
|
||
EMBEDDING_MODEL_NAME = os.environ.get("EMBEDDING_MODEL_NAME")
|
||
|
||
|
||
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 ""
|
||
|
||
|
||
def fixed_size_chunk(text, chunk_size=2000, overlap=50):
|
||
chunks = []
|
||
start = 0
|
||
while start < len(text):
|
||
end = start + chunk_size
|
||
chunks.append(text[max(0, start - overlap) : min(len(text), end + overlap)])
|
||
start += chunk_size
|
||
return chunks
|
||
|
||
|
||
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 = MAIN_HTTP_URL
|
||
body = {
|
||
"model": MAIN_MODEL_NAME,
|
||
"messages": _merge_messages_and_prompt(messages, prompt),
|
||
"max_tokens": max_tokens,
|
||
"temperature": temperature,
|
||
}
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {MAIN_HTTP_KEY}",
|
||
}
|
||
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 = MAIN_HTTP_URL
|
||
body = {
|
||
"model": MAIN_MODEL_NAME,
|
||
"messages": _merge_messages_and_prompt(messages or [], prompt),
|
||
"max_tokens": max_tokens,
|
||
"temperature": temperature,
|
||
"stream": True,
|
||
}
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {MAIN_HTTP_KEY}",
|
||
}
|
||
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
|
||
|
||
|
||
async def get_embedding(text, timeout: int = 30):
|
||
try:
|
||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||
url = EMBEDDING_HTTP_URL
|
||
body = {
|
||
"model": EMBEDDING_MODEL_NAME,
|
||
"input": [text],
|
||
}
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {EMBEDDING_HTTP_KEY}",
|
||
}
|
||
res = await client.post(url, headers=headers, json=body)
|
||
res.raise_for_status()
|
||
data = res.json()
|
||
return data["data"][0]["embedding"]
|
||
except Exception as e:
|
||
print(f"call_llm[ERROR]: {e}")
|
||
return []
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 设置
|
||
os.environ.setdefault("MAIN_HTTP_URL", "http://localhost:8022/v1/chat/completions")
|
||
os.environ.setdefault("MAIN_HTTP_KEY", "sk-local-827ccb0eea8a706c4c34a16891f84e7b")
|
||
os.environ.setdefault("MAIN_MODEL_NAME", "Qwen2.5")
|
||
|
||
os.environ.setdefault("EMBEDDING_HTTP_URL", "http://localhost:8023/v1/embeddings")
|
||
os.environ.setdefault("EMBEDDING_HTTP_KEY", "sk-local-827ccb0eea8a706c4c34a16891f84e7b")
|
||
os.environ.setdefault("EMBEDDING_MODEL_NAME", "Qwen3-Embedding")
|
||
# 读取
|
||
MAIN_HTTP_URL = os.environ.get("MAIN_HTTP_URL")
|
||
MAIN_HTTP_KEY = os.environ.get("MAIN_HTTP_KEY")
|
||
MAIN_MODEL_NAME = os.environ.get("MAIN_MODEL_NAME")
|
||
EMBEDDING_HTTP_URL = os.environ.get("EMBEDDING_HTTP_URL")
|
||
EMBEDDING_HTTP_KEY = os.environ.get("EMBEDDING_HTTP_KEY")
|
||
EMBEDDING_MODEL_NAME = os.environ.get("EMBEDDING_MODEL_NAME")
|
||
|
||
test_prompt = "你好,你是谁,有什么功能,中文回复"
|
||
|
||
# print("--- 普通对话 ---")
|
||
# text = asyncio.run(call_llm(test_prompt))
|
||
# print(text)
|
||
|
||
# print("--- 流式对话 ---")
|
||
# async def _test_stream():
|
||
# try:
|
||
# async for chunk in call_llm_stream(test_prompt):
|
||
# print(chunk, end="", flush=True)
|
||
# except Exception as e:
|
||
# print(f"_test_stream[ERROR]: {e}")
|
||
# asyncio.run(_test_stream())
|
||
|
||
print("--- embedding ---")
|
||
async def _test_embedding():
|
||
res = await get_embedding(test_prompt)
|
||
print(res)
|
||
|
||
asyncio.run(_test_embedding())
|