mirror of
https://github.com/HKUDS/LightRAG.git
synced 2025-07-23 08:52:21 +00:00
使用AzureOpenAI实现,支持RPM/TPM限制。修复原先429响应即抛出异常的问题
This commit is contained in:
parent
a50ed0b164
commit
2a3d92b515
@ -6,6 +6,7 @@ import numpy as np
|
|||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
import aiohttp
|
import aiohttp
|
||||||
import logging
|
import logging
|
||||||
|
from openai import AzureOpenAI
|
||||||
|
|
||||||
logging.basicConfig(level=logging.INFO)
|
logging.basicConfig(level=logging.INFO)
|
||||||
|
|
||||||
@ -32,11 +33,12 @@ os.mkdir(WORKING_DIR)
|
|||||||
async def llm_model_func(
|
async def llm_model_func(
|
||||||
prompt, system_prompt=None, history_messages=[], **kwargs
|
prompt, system_prompt=None, history_messages=[], **kwargs
|
||||||
) -> str:
|
) -> str:
|
||||||
headers = {
|
|
||||||
"Content-Type": "application/json",
|
client = AzureOpenAI(
|
||||||
"api-key": AZURE_OPENAI_API_KEY,
|
api_key=AZURE_OPENAI_API_KEY,
|
||||||
}
|
api_version=AZURE_OPENAI_API_VERSION,
|
||||||
endpoint = f"{AZURE_OPENAI_ENDPOINT}openai/deployments/{AZURE_OPENAI_DEPLOYMENT}/chat/completions?api-version={AZURE_OPENAI_API_VERSION}"
|
azure_endpoint=AZURE_OPENAI_ENDPOINT
|
||||||
|
)
|
||||||
|
|
||||||
messages = []
|
messages = []
|
||||||
if system_prompt:
|
if system_prompt:
|
||||||
@ -45,40 +47,29 @@ async def llm_model_func(
|
|||||||
messages.extend(history_messages)
|
messages.extend(history_messages)
|
||||||
messages.append({"role": "user", "content": prompt})
|
messages.append({"role": "user", "content": prompt})
|
||||||
|
|
||||||
payload = {
|
chat_completion = client.chat.completions.create(
|
||||||
"messages": messages,
|
model=AZURE_OPENAI_DEPLOYMENT, # model = "deployment_name".
|
||||||
"temperature": kwargs.get("temperature", 0),
|
messages=messages,
|
||||||
"top_p": kwargs.get("top_p", 1),
|
temperature=kwargs.get("temperature", 0),
|
||||||
"n": kwargs.get("n", 1),
|
top_p=kwargs.get("top_p", 1),
|
||||||
}
|
n=kwargs.get("n", 1),
|
||||||
|
|
||||||
async with aiohttp.ClientSession() as session:
|
|
||||||
async with session.post(endpoint, headers=headers, json=payload) as response:
|
|
||||||
if response.status != 200:
|
|
||||||
raise ValueError(
|
|
||||||
f"Request failed with status {response.status}: {await response.text()}"
|
|
||||||
)
|
)
|
||||||
result = await response.json()
|
return chat_completion.choices[0].message.content
|
||||||
return result["choices"][0]["message"]["content"]
|
|
||||||
|
|
||||||
|
|
||||||
async def embedding_func(texts: list[str]) -> np.ndarray:
|
async def embedding_func(texts: list[str]) -> np.ndarray:
|
||||||
headers = {
|
|
||||||
"Content-Type": "application/json",
|
|
||||||
"api-key": AZURE_OPENAI_API_KEY,
|
|
||||||
}
|
|
||||||
endpoint = f"{AZURE_OPENAI_ENDPOINT}openai/deployments/{AZURE_EMBEDDING_DEPLOYMENT}/embeddings?api-version={AZURE_EMBEDDING_API_VERSION}"
|
|
||||||
|
|
||||||
payload = {"input": texts}
|
client = AzureOpenAI(
|
||||||
|
api_key=AZURE_OPENAI_API_KEY,
|
||||||
async with aiohttp.ClientSession() as session:
|
api_version=AZURE_EMBEDDING_API_VERSION,
|
||||||
async with session.post(endpoint, headers=headers, json=payload) as response:
|
azure_endpoint=AZURE_OPENAI_ENDPOINT
|
||||||
if response.status != 200:
|
|
||||||
raise ValueError(
|
|
||||||
f"Request failed with status {response.status}: {await response.text()}"
|
|
||||||
)
|
)
|
||||||
result = await response.json()
|
embedding = client.embeddings.create(
|
||||||
embeddings = [item["embedding"] for item in result["data"]]
|
model=AZURE_EMBEDDING_DEPLOYMENT,
|
||||||
|
input=texts
|
||||||
|
)
|
||||||
|
|
||||||
|
embeddings = [item.embedding for item in embedding.data]
|
||||||
return np.array(embeddings)
|
return np.array(embeddings)
|
||||||
|
|
||||||
|
|
||||||
|
Loading…
x
Reference in New Issue
Block a user