Magic_yuan
commited on
Commit
·
5570390
1
Parent(s):
c417ca3
使用AzureOpenAI实现,支持RPM/TPM限制。修复原先429响应即抛出异常的问题
Browse files
examples/lightrag_azure_openai_demo.py
CHANGED
@@ -6,6 +6,7 @@ import numpy as np
|
|
6 |
from dotenv import load_dotenv
|
7 |
import aiohttp
|
8 |
import logging
|
|
|
9 |
|
10 |
logging.basicConfig(level=logging.INFO)
|
11 |
|
@@ -32,11 +33,12 @@ os.mkdir(WORKING_DIR)
|
|
32 |
async def llm_model_func(
|
33 |
prompt, system_prompt=None, history_messages=[], **kwargs
|
34 |
) -> str:
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
|
|
40 |
|
41 |
messages = []
|
42 |
if system_prompt:
|
@@ -45,41 +47,30 @@ async def llm_model_func(
|
|
45 |
messages.extend(history_messages)
|
46 |
messages.append({"role": "user", "content": prompt})
|
47 |
|
48 |
-
|
49 |
-
"
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
async with session.post(endpoint, headers=headers, json=payload) as response:
|
57 |
-
if response.status != 200:
|
58 |
-
raise ValueError(
|
59 |
-
f"Request failed with status {response.status}: {await response.text()}"
|
60 |
-
)
|
61 |
-
result = await response.json()
|
62 |
-
return result["choices"][0]["message"]["content"]
|
63 |
|
64 |
|
65 |
async def embedding_func(texts: list[str]) -> np.ndarray:
|
66 |
-
|
67 |
-
|
68 |
-
|
69 |
-
|
70 |
-
|
71 |
-
|
72 |
-
|
73 |
-
|
74 |
-
|
75 |
-
|
76 |
-
|
77 |
-
|
78 |
-
|
79 |
-
)
|
80 |
-
result = await response.json()
|
81 |
-
embeddings = [item["embedding"] for item in result["data"]]
|
82 |
-
return np.array(embeddings)
|
83 |
|
84 |
|
85 |
async def test_funcs():
|
|
|
6 |
from dotenv import load_dotenv
|
7 |
import aiohttp
|
8 |
import logging
|
9 |
+
from openai import AzureOpenAI
|
10 |
|
11 |
logging.basicConfig(level=logging.INFO)
|
12 |
|
|
|
33 |
async def llm_model_func(
|
34 |
prompt, system_prompt=None, history_messages=[], **kwargs
|
35 |
) -> str:
|
36 |
+
|
37 |
+
client = AzureOpenAI(
|
38 |
+
api_key=AZURE_OPENAI_API_KEY,
|
39 |
+
api_version=AZURE_OPENAI_API_VERSION,
|
40 |
+
azure_endpoint=AZURE_OPENAI_ENDPOINT
|
41 |
+
)
|
42 |
|
43 |
messages = []
|
44 |
if system_prompt:
|
|
|
47 |
messages.extend(history_messages)
|
48 |
messages.append({"role": "user", "content": prompt})
|
49 |
|
50 |
+
chat_completion = client.chat.completions.create(
|
51 |
+
model=AZURE_OPENAI_DEPLOYMENT, # model = "deployment_name".
|
52 |
+
messages=messages,
|
53 |
+
temperature=kwargs.get("temperature", 0),
|
54 |
+
top_p=kwargs.get("top_p", 1),
|
55 |
+
n=kwargs.get("n", 1),
|
56 |
+
)
|
57 |
+
return chat_completion.choices[0].message.content
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
58 |
|
59 |
|
60 |
async def embedding_func(texts: list[str]) -> np.ndarray:
|
61 |
+
|
62 |
+
client = AzureOpenAI(
|
63 |
+
api_key=AZURE_OPENAI_API_KEY,
|
64 |
+
api_version=AZURE_EMBEDDING_API_VERSION,
|
65 |
+
azure_endpoint=AZURE_OPENAI_ENDPOINT
|
66 |
+
)
|
67 |
+
embedding = client.embeddings.create(
|
68 |
+
model=AZURE_EMBEDDING_DEPLOYMENT,
|
69 |
+
input=texts
|
70 |
+
)
|
71 |
+
|
72 |
+
embeddings = [item.embedding for item in embedding.data]
|
73 |
+
return np.array(embeddings)
|
|
|
|
|
|
|
|
|
74 |
|
75 |
|
76 |
async def test_funcs():
|