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
- headers = {
36
- "Content-Type": "application/json",
37
- "api-key": AZURE_OPENAI_API_KEY,
38
- }
39
- endpoint = f"{AZURE_OPENAI_ENDPOINT}openai/deployments/{AZURE_OPENAI_DEPLOYMENT}/chat/completions?api-version={AZURE_OPENAI_API_VERSION}"
 
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
- payload = {
49
- "messages": messages,
50
- "temperature": kwargs.get("temperature", 0),
51
- "top_p": kwargs.get("top_p", 1),
52
- "n": kwargs.get("n", 1),
53
- }
54
-
55
- async with aiohttp.ClientSession() as session:
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
- headers = {
67
- "Content-Type": "application/json",
68
- "api-key": AZURE_OPENAI_API_KEY,
69
- }
70
- endpoint = f"{AZURE_OPENAI_ENDPOINT}openai/deployments/{AZURE_EMBEDDING_DEPLOYMENT}/embeddings?api-version={AZURE_EMBEDDING_API_VERSION}"
71
-
72
- payload = {"input": texts}
73
-
74
- async with aiohttp.ClientSession() as session:
75
- async with session.post(endpoint, headers=headers, json=payload) as response:
76
- if response.status != 200:
77
- raise ValueError(
78
- f"Request failed with status {response.status}: {await response.text()}"
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():