yujunjun/LightRAG
1
0
Fork 0
mirror of https://github.com/HKUDS/LightRAG.git synced 2025-07-25 09:50:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

使用AzureOpenAI实现，支持RPM/TPM限制。修复原先429响应即抛出异常的问题

Browse Source
This commit is contained in:
Magic_yuan 2024-11-21 10:37:09 +08:00 committed by GitHub
parent a50ed0b164
commit 2a3d92b515
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 27 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

57
examples/lightrag_azure_openai_demo.py
View File

@ -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)