235 lines
7.5 KiB
Python
235 lines
7.5 KiB
Python
from collections.abc import Iterable
|
|
import os
|
|
import pipmaster as pm # Pipmaster for dynamic library install
|
|
|
|
# install specific modules
|
|
if not pm.is_installed("openai"):
|
|
pm.install("openai")
|
|
|
|
from openai import (
|
|
AsyncAzureOpenAI,
|
|
APIConnectionError,
|
|
RateLimitError,
|
|
APITimeoutError,
|
|
)
|
|
from openai.types.chat import ChatCompletionMessageParam
|
|
|
|
from tenacity import (
|
|
retry,
|
|
stop_after_attempt,
|
|
wait_exponential,
|
|
retry_if_exception_type,
|
|
)
|
|
|
|
from lightrag.utils import (
|
|
wrap_embedding_func_with_attrs,
|
|
safe_unicode_decode,
|
|
logger,
|
|
)
|
|
|
|
import numpy as np
|
|
|
|
|
|
@retry(
|
|
stop=stop_after_attempt(3),
|
|
wait=wait_exponential(multiplier=1, min=4, max=10),
|
|
retry=retry_if_exception_type(
|
|
(RateLimitError, APIConnectionError, APIConnectionError)
|
|
),
|
|
)
|
|
async def azure_openai_complete_if_cache(
|
|
model,
|
|
prompt,
|
|
system_prompt: str | None = None,
|
|
history_messages: Iterable[ChatCompletionMessageParam] | None = None,
|
|
enable_cot: bool = False,
|
|
base_url: str | None = None,
|
|
api_key: str | None = None,
|
|
api_version: str | None = None,
|
|
token_tracker: Any | None = None,
|
|
**kwargs,
|
|
):
|
|
if enable_cot:
|
|
logger.debug(
|
|
"enable_cot=True is not supported for the Azure OpenAI API and will be ignored."
|
|
)
|
|
deployment = os.getenv("AZURE_OPENAI_DEPLOYMENT") or model or os.getenv("LLM_MODEL")
|
|
base_url = (
|
|
base_url or os.getenv("AZURE_OPENAI_ENDPOINT") or os.getenv("LLM_BINDING_HOST")
|
|
)
|
|
api_key = (
|
|
api_key or os.getenv("AZURE_OPENAI_API_KEY") or os.getenv("LLM_BINDING_API_KEY")
|
|
)
|
|
api_version = (
|
|
api_version
|
|
or os.getenv("AZURE_OPENAI_API_VERSION")
|
|
or os.getenv("OPENAI_API_VERSION")
|
|
)
|
|
|
|
kwargs.pop("hashing_kv", None)
|
|
kwargs.pop("keyword_extraction", None)
|
|
timeout = kwargs.pop("timeout", None)
|
|
|
|
openai_async_client = AsyncAzureOpenAI(
|
|
azure_endpoint=base_url,
|
|
azure_deployment=deployment,
|
|
api_key=api_key,
|
|
api_version=api_version,
|
|
timeout=timeout,
|
|
)
|
|
messages = []
|
|
if system_prompt:
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
if history_messages:
|
|
messages.extend(history_messages)
|
|
if prompt is not None:
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
if "response_format" in kwargs:
|
|
response = await openai_async_client.beta.chat.completions.parse(
|
|
model=model, messages=messages, **kwargs
|
|
)
|
|
else:
|
|
response = await openai_async_client.chat.completions.create(
|
|
model=model, messages=messages, **kwargs
|
|
)
|
|
|
|
if hasattr(response, "__aiter__"):
|
|
final_chunk_usage = None
|
|
accumulated_response = ""
|
|
|
|
async def inner():
|
|
nonlocal final_chunk_usage, accumulated_response
|
|
try:
|
|
async for chunk in response:
|
|
if len(chunk.choices) == 0:
|
|
continue
|
|
content = chunk.choices[0].delta.content
|
|
if content is None:
|
|
continue
|
|
accumulated_response += content
|
|
if r"\u" in content:
|
|
content = safe_unicode_decode(content.encode("utf-8"))
|
|
yield content
|
|
|
|
# Check for usage in the last chunk
|
|
if hasattr(chunk, "usage") and chunk.usage is not None:
|
|
final_chunk_usage = chunk.usage
|
|
except Exception as e:
|
|
logger.error(f"Error in Azure OpenAI stream response: {str(e)}")
|
|
raise
|
|
finally:
|
|
# After streaming is complete, track token usage
|
|
if token_tracker and final_chunk_usage:
|
|
# Use actual usage from the API
|
|
token_counts = {
|
|
"prompt_tokens": getattr(final_chunk_usage, "prompt_tokens", 0),
|
|
"completion_tokens": getattr(
|
|
final_chunk_usage, "completion_tokens", 0
|
|
),
|
|
"total_tokens": getattr(final_chunk_usage, "total_tokens", 0),
|
|
}
|
|
token_tracker.add_usage(token_counts)
|
|
logger.debug(f"Azure OpenAI streaming token usage: {token_counts}")
|
|
elif token_tracker:
|
|
logger.debug(
|
|
"No usage information available in Azure OpenAI streaming response"
|
|
)
|
|
|
|
return inner()
|
|
else:
|
|
content = response.choices[0].message.content
|
|
if r"\u" in content:
|
|
content = safe_unicode_decode(content.encode("utf-8"))
|
|
|
|
# Track token usage for non-streaming response
|
|
if token_tracker and hasattr(response, "usage"):
|
|
token_counts = {
|
|
"prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
|
|
"completion_tokens": getattr(response.usage, "completion_tokens", 0),
|
|
"total_tokens": getattr(response.usage, "total_tokens", 0),
|
|
}
|
|
token_tracker.add_usage(token_counts)
|
|
logger.debug(f"Azure OpenAI non-streaming token usage: {token_counts}")
|
|
|
|
return content
|
|
|
|
|
|
async def azure_openai_complete(
|
|
prompt,
|
|
system_prompt=None,
|
|
history_messages=[],
|
|
keyword_extraction=False,
|
|
token_tracker=None,
|
|
**kwargs,
|
|
) -> str:
|
|
kwargs.pop("keyword_extraction", None)
|
|
result = await azure_openai_complete_if_cache(
|
|
os.getenv("LLM_MODEL", "gpt-4o-mini"),
|
|
prompt,
|
|
system_prompt=system_prompt,
|
|
history_messages=history_messages,
|
|
token_tracker=token_tracker,
|
|
**kwargs,
|
|
)
|
|
return result
|
|
|
|
|
|
@wrap_embedding_func_with_attrs(embedding_dim=1536)
|
|
@retry(
|
|
stop=stop_after_attempt(3),
|
|
wait=wait_exponential(multiplier=1, min=4, max=10),
|
|
retry=retry_if_exception_type(
|
|
(RateLimitError, APIConnectionError, APITimeoutError)
|
|
),
|
|
)
|
|
async def azure_openai_embed(
|
|
texts: list[str],
|
|
model: str | None = None,
|
|
base_url: str | None = None,
|
|
api_key: str | None = None,
|
|
api_version: str | None = None,
|
|
token_tracker: Any | None = None,
|
|
) -> np.ndarray:
|
|
deployment = (
|
|
os.getenv("AZURE_EMBEDDING_DEPLOYMENT")
|
|
or model
|
|
or os.getenv("EMBEDDING_MODEL", "text-embedding-3-small")
|
|
)
|
|
base_url = (
|
|
base_url
|
|
or os.getenv("AZURE_EMBEDDING_ENDPOINT")
|
|
or os.getenv("EMBEDDING_BINDING_HOST")
|
|
)
|
|
api_key = (
|
|
api_key
|
|
or os.getenv("AZURE_EMBEDDING_API_KEY")
|
|
or os.getenv("EMBEDDING_BINDING_API_KEY")
|
|
)
|
|
api_version = (
|
|
api_version
|
|
or os.getenv("AZURE_EMBEDDING_API_VERSION")
|
|
or os.getenv("OPENAI_API_VERSION")
|
|
)
|
|
|
|
openai_async_client = AsyncAzureOpenAI(
|
|
azure_endpoint=base_url,
|
|
azure_deployment=deployment,
|
|
api_key=api_key,
|
|
api_version=api_version,
|
|
)
|
|
|
|
response = await openai_async_client.embeddings.create(
|
|
model=model, input=texts, encoding_format="float"
|
|
)
|
|
|
|
# Track token usage for embeddings if token tracker is provided
|
|
if token_tracker and hasattr(response, "usage"):
|
|
token_counts = {
|
|
"prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
|
|
"total_tokens": getattr(response.usage, "total_tokens", 0),
|
|
}
|
|
token_tracker.add_usage(token_counts)
|
|
logger.debug(f"Azure OpenAI embedding token usage: {token_counts}")
|
|
|
|
return np.array([dp.embedding for dp in response.data])
|