Spaces:
Sleeping
Sleeping
| 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 | |
| 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) | |
| ] | |
| 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)] | |