Spaces:
Sleeping
Sleeping
File size: 2,295 Bytes
da21448 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | from collections.abc import Sequence
from itertools import batched
from typing import Any
from openai import APIConnectionError, APITimeoutError, InternalServerError, OpenAI, RateLimitError
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_random_exponential
from techdocs_graphrag.config import Settings
EMBEDDING_BATCH_SIZE = 96
TRANSIENT_ERRORS = (APIConnectionError, APITimeoutError, InternalServerError, RateLimitError)
with_retry = retry(
retry=retry_if_exception_type(TRANSIENT_ERRORS),
wait=wait_random_exponential(min=1, max=30),
stop=stop_after_attempt(5),
reraise=True,
)
class OpenAIGateway:
def __init__(self, settings: Settings) -> None:
self._client = OpenAI(api_key=settings.openai_api_key.get_secret_value())
self._chat_model = settings.chat_model
self._chat_temperature = settings.chat_temperature
self._embedding_model = settings.embedding_model
self._embedding_dimensions = settings.embedding_dimensions
@with_retry
def complete(self, system_prompt: str, user_prompt: str, *, json_mode: bool = False) -> str:
options: dict[str, Any] = {}
if json_mode:
options["response_format"] = {"type": "json_object"}
if self._chat_temperature is not None:
options["temperature"] = self._chat_temperature
response = self._client.chat.completions.create(
model=self._chat_model,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
**options,
)
return response.choices[0].message.content or ""
def embed(self, texts: Sequence[str]) -> list[list[float]]:
return [
vector
for batch in batched(texts, EMBEDDING_BATCH_SIZE)
for vector in self._embed_batch(batch)
]
@with_retry
def _embed_batch(self, texts: Sequence[str]) -> list[list[float]]:
response = self._client.embeddings.create(
model=self._embedding_model,
input=list(texts),
dimensions=self._embedding_dimensions,
)
return [item.embedding for item in sorted(response.data, key=lambda item: item.index)]
|