@@ -383,24 +383,31 @@ LOGGING = { 'version': 1, 'disable_existing_loggers': False, 'formatters': { + 'pretty': { + 'format': '[{asctime}] [{levelname}] [{name}] {message}', + 'style': '{', + 'datefmt': '%d.%m.%Y %H:%M:%S', + }, 'verbose': { - 'format': '[{asctime}] {levelname} [pid={process}] [thread={thread}] [{pathname}:{lineno}] {message}', + 'format': '[{asctime}] {levelname} [{name}] [{process}:{threadName}] [{pathname}:{lineno}] {message}', 'style': '{', - 'datefmt': '%Y-%m-%d %H:%M:%S', + 'datefmt': '%Y-%m-%dT%H:%M:%SZ', }, }, 'handlers': { 'console': { 'class': 'logging.StreamHandler', - 'formatter': 'verbose', + 'formatter': 'pretty' if DEBUG else 'verbose', }, }, 'loggers': { 'UnleashClient': { + 'handlers': ['console'], 'level': 'CRITICAL', 'propagate': False, }, 'apscheduler': { + 'handlers': ['console'], 'level': 'CRITICAL', 'propagate': False, }, @@ -471,13 +478,13 @@ if CACHEOPS_REDIS: 'token_blacklist.outstandingtoken': {'ops': 'get', 'timeout': 60 * 60 * 24}, } -# UNLEASH settings -FEATURE_FLAG_API_URL = env.str('FEATURE_FLAG_API_URL') -FEATURE_FLAG_APP_NAME = env.str('FEATURE_FLAG_APP_NAME', 'staging') -FEATURE_FLAG_INSTANCE_ID = env.str('FEATURE_FLAG_INSTANCE_ID') -FEATURE_FLAG_WEBHOOK_SECRET_KEY = env.str( - 'FEATURE_FLAG_WEBHOOK_SECRET_KEY', 'FEATURE_FLAG_WEBHOOK_SECRET_KEY' -) +# Unleash settings +UNLEASH_API_URL = env.str('UNLEASH_API_URL', 'https://example.com') +UNLEASH_APP_NAME = env.str('UNLEASH_APP_NAME', 'Development') +UNLEASH_REQUEST_TIMEOUT = env.int('UNLEASH_REQUEST_TIMEOUT', 3) +UNLEASH_REQUEST_RETRIES = env.int('UNLEASH_REQUEST_RETRIES', 1) +UNLEASH_INSTANCE_ID = env.str('UNLEASH_INSTANCE_ID', '') +UNLEASH_WEBHOOK_SECRET_KEY = env.str('UNLEASH_WEBHOOK_SECRET_KEY', 'defaultsecretkey') # RECURRING SETTINGS MAX_RECURRING_ATTEMPTS = env.int('MAX_RECURRING_ATTEMPTS', 1) @@ -23,6 +23,9 @@ logger = logging.getLogger(__name__) api = NinjaAPI(title='AIR API', version='1.0.0', parser=MultiContentTypeParser(), docs_url=None) compatibility_api = NinjaAPI(title='AIR API DEBUG', version='0.0.1', docs_url=None) compatibility_api_v2 = NinjaAPI(title='AIR API DEBUG v2', version='2.0.0', docs_url=None) +public_api = NinjaAPI( + title='PUBLIC AIR API', urls_namespace='public-api-1.0.0', parser=MultiContentTypeParser(), docs_url=None +) api.add_router('users/', 'users.routes.v1.router') api.add_router('chats/', 'tools.chats.routes.v1.router') @@ -35,6 +38,8 @@ compatibility_api.add_router('ml_model/', 'ml_model.routes.v1.router') compatibility_api_v2.add_router('auth/', 'authentication.routes.v2.router') +public_api.add_router('', 'tools.public_api.routes.v1.router') + class Status(Schema): status: Literal['ok', 'dead'] @@ -97,6 +102,7 @@ urlpatterns = ( path('api/v1/api/', api.urls), path('api/v1/', compatibility_api.urls), path('api/v1/v2/', compatibility_api_v2.urls), + path('public/', public_api.urls), ] + static(settings.STATIC_URL, document_root=settings.STATIC_ROOT) + public_urlpatterns @@ -122,3 +128,4 @@ if settings.DEBUG: api.docs_url = '/docs' compatibility_api.docs_url = '/docs' compatibility_api_v2.docs_url = '/docs' + public_api.docs_url = '/docs' @@ -13,11 +13,15 @@ from lib.unleash.cache import UnleashRedisCache class UnleashFeatureFlagService(FeatureFlagService): def __init__(self) -> None: self.client = UnleashClient( - url=settings.FEATURE_FLAG_API_URL, - app_name=settings.FEATURE_FLAG_APP_NAME, - instance_id=settings.FEATURE_FLAG_INSTANCE_ID, + url=settings.UNLEASH_API_URL, + app_name=settings.UNLEASH_APP_NAME, + environment=settings.UNLEASH_APP_NAME, + instance_id=settings.UNLEASH_INSTANCE_ID, + request_timeout=settings.UNLEASH_REQUEST_TIMEOUT, + request_retries=settings.UNLEASH_REQUEST_RETRIES, cache=UnleashRedisCache(), - environment=settings.FEATURE_FLAG_APP_NAME, + disable_registration=True, + disable_metrics=True, ) def get_flag_state_by_emails(self, name: str, emails: List[Email]) -> Mapping[Email, State]: @@ -39,7 +39,7 @@ def stream_message_reconnect(request, chat_uid: UUID, offset: int = 0): return StreamingHttpResponse( sse_chat_stream.event_stream(request=request, offset=offset), content_type='text/event-stream', - headers={'Cache-Control': 'no-cache', 'Connection': 'keep-alive', 'X-Accel-Buffering': 'no'}, + headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'}, ) @@ -69,5 +69,5 @@ def stream_message(request, chat_uid: UUID, body: MessageInSchema): return StreamingHttpResponse( sse_chat_stream.event_stream(request=request), content_type='text/event-stream', - headers={'Cache-Control': 'no-cache', 'Connection': 'keep-alive', 'X-Accel-Buffering': 'no'}, + headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'}, ) @@ -66,3 +66,13 @@ class SSEStoreService: def cleanup(self) -> None: self.delete_stream() + + +class PublicSSEStoreService(SSEStoreService): + def __init__(self, idempotency_key: UUID, user_uuid: UUID): + self.user_uuid = user_uuid + self.idempotency_key = idempotency_key + self.redis_client = self._get_redis_client() + + def _get_cache_key(self) -> str: + return f'sse:tokens:{self.user_uuid}:{self.idempotency_key}' \ No newline at end of file @@ -1,25 +1,28 @@ +from decimal import Decimal + from celery import shared_task from django.conf import settings +from django.db.models import F, Value +from django.db.models.functions import Greatest + from messages.models import Message +from ml_model.models import NeuronModel +from ml_model.services.base import StreamSimpleService from tools.chats.models import Chat from tools.chats.services.sse_chunk_service import SSEChunkService -from tools.chats.services.sse_store import SSEStoreService +from tools.chats.services.sse_store import PublicSSEStoreService, SSEStoreService +from tools.public_api.models import APIKey, APIStore -@shared_task(soft_time_limit=570, time_limit=600) -def event_stream_task(chat_uuid: str, message_uuid: str, user_uuid: str) -> None: - store = SSEStoreService(user_uuid=user_uuid, chat_uuid=chat_uuid) +def _run_stream(store: SSEStoreService, message_uuid: str, service: StreamSimpleService, message): stream = None event_id = 0 try: - chat = Chat.objects.select_related('model').get(pk=chat_uuid) - message = Message.objects.get(pk=message_uuid) - event_id += 1 store.push(SSEChunkService.start(event_id, message_uuid)) - stream = chat.model.service(chat).make_stream(message) + stream = service.make_stream(message) while True: token = next(stream) if not token: @@ -29,8 +32,6 @@ def event_stream_task(chat_uuid: str, message_uuid: str, user_uuid: str) -> None except StopIteration as exc: event_id += 1 store.push(SSEChunkService.done(event_id, exc.value or ''), ttl=settings.SSE_DONE_STREAM_TTL) - except (Chat.DoesNotExist, Message.DoesNotExist) as exc: - store.push(SSEChunkService.error(event_id + 1, str(exc)), ttl=settings.SSE_DONE_STREAM_TTL) except Exception as exc: if event_id < 2: message.is_sent = False @@ -40,3 +41,56 @@ def event_stream_task(chat_uuid: str, message_uuid: str, user_uuid: str) -> None finally: if stream is not None: stream.close() + + +@shared_task(soft_time_limit=570, time_limit=600) +def event_stream_task(chat_uuid: str, message_uuid: str, user_uuid: str) -> None: + store = SSEStoreService(user_uuid=user_uuid, chat_uuid=chat_uuid) + + chat = Chat.objects.select_related('model').get(pk=chat_uuid) + message = Message.objects.get(pk=message_uuid) + + service = chat.model.service(chat) + + _run_stream(store, message_uuid, service, message) + + +@shared_task(soft_time_limit=570, time_limit=600) +def public_event_stream_task( + start_user_balance: Decimal, + message_uuid: str, + user_uuid: str, + model_slug: str, + idempotency_key: str, + api_key_uuid: str, + debit_api_key_limit: bool, +): + store = PublicSSEStoreService(user_uuid=user_uuid, idempotency_key=idempotency_key) + + message = Message.objects.get(pk=message_uuid) + api_store = APIStore.objects.select_related( + 'user', + 'user__payment_plan', + 'user__payment_plan__plan', + 'user__business_account', + 'user__business_account__group', + 'user__business_account__parent_company', + 'user__business_account__parent_company__user__payment_plan', + 'user__business_account__parent_company__user__payment_plan__plan', + ).get(pk=message.object_id) + model = NeuronModel.objects.get(slug=model_slug) + + message.content_object = api_store + message.content_object.model = model + + service = model.service(api_store) + + try: + _run_stream(store, message_uuid, service, message) + finally: + if debit_api_key_limit: + spent = start_user_balance - api_store.user.balance + if spent > 0: + APIKey.objects.filter(pk=api_key_uuid).update( + token_limit=Greatest(F('token_limit') - spent, Value(Decimal('0'))), + ) @@ -0,0 +1,122 @@ +from django.db.models import Q +from django.http import StreamingHttpResponse +from django.utils.translation import gettext as _ +from ninja.errors import HttpError + +from ml_model.models import NeuronModel +from tools.chats.services.sse_store import PublicSSEStoreService +from tools.public_api.services.openai_errors import OpenAIErrorService +from tools.public_api.services.openai_stream import OpenAIStreamService + + +def _resolve_chat_model(model_ref: str) -> NeuronModel: + try: + return ( + NeuronModel.objects.filter( + Q(model_modelversions__slug=model_ref) | Q(slug=model_ref), + category__slug='chat-bots', + ) + .distinct() + .get() + ) + except NeuronModel.DoesNotExist: + raise HttpError(404, _('Model not found')) + + +def _require_idempotency_key(request) -> str: + if not (key := request.headers.get('Idempotency-Key', '').strip()): + raise HttpError(400, _('Idempotency-Key header is not provided')) + return key + + +def _api_user(request): + from tools.public_api.routes.v1 import _get_api_key + + api_key = _get_api_key( + request, + select_related=['user__host_account', 'user__business_account'], + check_usage_limit=False, + ) + user = api_key.user + if user.account_type not in ('business_host', 'regular', 'business_admin'): + raise HttpError(403, _('API key is not available for this account type')) + return user + + +@OpenAIErrorService.view +def openai_responses_stream(request, body: dict): + from tools.public_api.routes.v1 import public_stream_message + + if not body.get('stream'): + raise HttpError(400, _('Only streaming is supported')) + if not (model_ref := body.get('model')): + raise HttpError(400, _('You must provide a model parameter')) + + model = _resolve_chat_model(model_ref) + store = PublicSSEStoreService( + user_uuid=_api_user(request).pk, + idempotency_key=_require_idempotency_key(request), + ) + stream = OpenAIStreamService(store) + stream.cleanup_if_finished() + + response = public_stream_message(request, model.slug, OpenAIStreamService.user_message_body(body)) + meta = stream.init_meta(model_ref) + response.streaming_content = stream.event_stream(request, meta=meta) + return response + + +@OpenAIErrorService.view +def openai_responses_stream_reconnect( + request, + response_id: str | None = None, + *, + stream: bool = True, + starting_after: int = 0, +): + if not stream: + raise HttpError(400, _('Only streaming reconnect is supported')) + + user = _api_user(request) + store = PublicSSEStoreService( + user_uuid=user.pk, idempotency_key=response_id or _require_idempotency_key(request) + ) + if not store.exists(): + raise HttpError(404, _('Stream not found')) + + openai_stream = OpenAIStreamService(store) + meta = openai_stream.init_meta('', reconnect=True) + if str(meta.get('user_uuid')) != str(user.pk): + raise HttpError(404, _('Stream not found')) + + return StreamingHttpResponse( + openai_stream.event_stream(request, reconnect=True, meta=meta), + content_type='text/event-stream', + headers=OpenAIStreamService.SSE_HEADERS, + ) + + +from tools.public_api.routes import v1 as public_v1_routes + + +def _register_reconnect_get(path: str, *, with_response_id: bool): + if with_response_id: + + @public_v1_routes.router.get(path, tags=['openai/responses']) + def handler(request, response_id: str, stream: bool = True, starting_after: int = 0): + return openai_responses_stream_reconnect( + request, response_id, stream=stream, starting_after=starting_after + ) + else: + + @public_v1_routes.router.get(path, tags=['openai/responses']) + def handler(request, stream: bool = True, starting_after: int = 0): + return openai_responses_stream_reconnect(request, stream=stream, starting_after=starting_after) + + +for with_rid, paths in ( + (False, ('openai/v1/responses', 'openai/responses')), + (True, ('openai/v1/responses/{response_id}', 'openai/responses/{response_id}')), +): + for path in paths: + _register_reconnect_get(path, with_response_id=with_rid) @@ -0,0 +1,153 @@ +from datetime import date + +from django.http import StreamingHttpResponse +from ninja import Body, Router +from ninja.errors import HttpError + +from django.utils.translation import gettext as _ + +from messages.models import Message +from ml_model.selectors.ml_models_selector import NeuronModelSelector + +from tools.chats.schemas import MessageInSchema +from tools.chats.services.sse_chat_stream import SSEChatStreamService +from tools.chats.services.sse_store import PublicSSEStoreService +from tools.chats.tasks import public_event_stream_task + +from tools.public_api.models import APIKey, APIStore + +router = Router(auth=None, tags=['public']) + + +def _get_api_key( + request, + select_related: list[str] = None, + prefetch_related: list[str] = None, + *, + check_usage_limit: bool = True, +): + raw_api_key = request.headers.get('Authorization', '') + if not raw_api_key: + raise HttpError(401, _('No API Key in Authorization header')) + + if (split_api_key := raw_api_key.split())[0] == 'Bearer': + raw_api_key = split_api_key[-1] + + api_key = ( + APIKey.objects.select_related('user', *(select_related or [])) + .prefetch_related(*(prefetch_related or [])) + .filter(key=raw_api_key, is_deleted=False) + ).first() + + if not api_key: + raise HttpError(404, _('API key not found')) + if check_usage_limit: + if api_key.expires_at and api_key.expires_at < date.today(): + raise HttpError(401, _('API key expired')) + if api_key.token_limit is not None and api_key.token_limit < 1: + raise HttpError(403, _('API key limit exceeded')) + + return api_key + + +@router.get('text/{model_slug}/stream', tags=['public/text']) +def public_stream_message_reconnect(request, model_slug: str, offset: int = 0): + api_key = _get_api_key( + request, + select_related=['user__host_account', 'user__business_account'], + check_usage_limit=False, + ) + user = api_key.user + if user.account_type not in ('business_host', 'regular', 'business_admin'): + raise HttpError(403, _('API key is not available for this account type')) + + idempotency_key = request.headers.get('Idempotency-Key', '') + if not idempotency_key: + raise HttpError(400, _('Idempotency-Key header is not provided')) + + store = PublicSSEStoreService(user_uuid=user.pk, idempotency_key=idempotency_key) + if not store.exists(): + raise HttpError(404, _('Stream not found')) + + sse_chat_stream = SSEChatStreamService(store) + return StreamingHttpResponse( + sse_chat_stream.event_stream(request=request, offset=offset), + content_type='text/event-stream', + headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'}, + ) + + +@router.post('text/{model_slug}/stream', tags=['public/text']) +def public_stream_message(request, model_slug: str, body: MessageInSchema): + api_key = _get_api_key( + request, + select_related=[ + 'user__host_account', + 'user__business_account', + 'user__payment_plan', + 'user__payment_plan__plan', + 'user__business_account__parent_company', + 'user__business_account__parent_company__user__payment_plan', + 'user__business_account__parent_company__user__payment_plan__plan', + ], + prefetch_related=[ + 'user__payment_plan__plan__features', + 'user__business_account__parent_company__user__payment_plan__plan__features', + ], + ) + user = api_key.user + if user.account_type not in ('business_host', 'regular', 'business_admin'): + raise HttpError(403, _('API key is not available for this account type')) + balance = user.balance + + idempotency_key = request.headers.get('Idempotency-Key', '') + if not idempotency_key: + raise HttpError(400, _('Idempotency-Key header is not provided')) + + api_store, created = APIStore.objects.get_or_create(user=user) + + selector = NeuronModelSelector(user) + model = selector.get_model_by_slug(slug=model_slug) + if model.blocked: + raise HttpError(403, _('Model is blocked by outdating or temporary block, please retry later')) + if not model.streaming: + raise HttpError(501, _('Stream not supported for this model')) + + if not body.content: + raise HttpError(400, _('The request must not be empty')) + + store = PublicSSEStoreService(user_uuid=user.pk, idempotency_key=idempotency_key) + if store.exists(): + raise HttpError(409, _('Stream already in progress')) + + data = body.model_dump(include={'content', 'file', 'info'}, exclude_unset=True) + input_message = Message.objects.create( + content_object=api_store, from_model=False, from_public_api=True, **data + ) + store.start() + + public_event_stream_task.delay( + start_user_balance=balance, + message_uuid=str(input_message.pk), + user_uuid=str(user.pk), + model_slug=model_slug, + idempotency_key=idempotency_key, + api_key_uuid=str(api_key.pk), + debit_api_key_limit=api_key.token_limit is not None, + ) + + sse_chat_stream = SSEChatStreamService(store) + return StreamingHttpResponse( + sse_chat_stream.event_stream(request=request), + content_type='text/event-stream', + headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'}, + ) + + +from tools.public_api.routes.providers.openai import openai_responses_stream + + +@router.post('openai/v1/responses', tags=['openai/responses']) +@router.post('openai/responses', tags=['openai/responses']) +def openai_responses(request, body: dict = Body(...)): + return openai_responses_stream(request, body) @@ -1 +1,3 @@ from .api_key import APIKeyService +from .openai_errors import OpenAIErrorService +from .openai_stream import OpenAIStreamService @@ -0,0 +1,108 @@ +from functools import wraps + +from django.http import JsonResponse +from ninja.errors import HttpError + + +class OpenAIErrorService: + ERR_TYPES = { + 400: 'invalid_request_error', + 401: 'authentication_error', + 403: 'permission_error', + 404: 'invalid_request_error', + 409: 'invalid_request_error', + 501: 'api_error', + } + ERR_CODES = { + 'only streaming is supported': 'streaming_not_supported', + 'only streaming reconnect is supported': 'streaming_not_supported', + 'you must provide a model parameter': 'missing_required_parameter', + 'missing required parameter: input': 'missing_required_parameter', + 'model not found': 'model_not_found', + 'response not found': 'response_not_found', + 'stream not found': 'stream_not_found', + 'stream already in progress': 'stream_in_progress', + 'idempotency-key header is not provided': 'missing_idempotency_key', + 'no api key in authorization header': 'missing_api_key', + 'api key not found': 'invalid_api_key', + 'api key expired': 'expired_api_key', + 'api key limit exceeded': 'insufficient_quota', + 'api key is not available for this account type': 'account_not_allowed', + 'stream not supported for this model': 'streaming_not_supported', + } + + @classmethod + def response( + cls, + status: int, + message, + *, + code: str | None = None, + param: str | None = None, + ) -> JsonResponse: + msg = str(message).strip() + if not code: + code = cls._infer_code(status, msg) + return JsonResponse( + { + 'error': { + 'message': msg, + 'type': cls.ERR_TYPES.get(status, 'api_error'), + 'param': param, + 'code': code, + } + }, + status=status, + ) + + @classmethod + def from_http_error(cls, exc: HttpError) -> JsonResponse: + return cls.response(exc.status_code, exc.message) + + @classmethod + def view(cls, fn): + @wraps(fn) + def wrapper(*args, **kwargs): + try: + return fn(*args, **kwargs) + except HttpError as exc: + return cls.from_http_error(exc) + + return wrapper + + @classmethod + def _infer_code(cls, status: int, message: str) -> str | None: + lower = message.lower() + if code := cls.ERR_CODES.get(lower): + return code + if status == 409: + return 'stream_in_progress' + if 'idempotency' in lower: + return 'missing_idempotency_key' + if status == 404: + return next( + ( + c + for k, c in ( + ('api key', 'invalid_api_key'), + ('model', 'model_not_found'), + ('response', 'response_not_found'), + ('stream', 'stream_not_found'), + ) + if k in lower + ), + None, + ) + if status == 403: + if 'limit' in lower: + return 'insufficient_quota' + if 'account' in lower: + return 'account_not_allowed' + if status == 501: + return 'streaming_not_supported' + if status == 401: + if 'expired' in lower: + return 'expired_api_key' + if 'api key' in lower: + return 'missing_api_key' + return None @@ -0,0 +1,320 @@ +import secrets +import time +from typing import Iterator + +import orjson +from django.http import HttpRequest +from django.utils.translation import gettext as _ +from ninja.errors import HttpError + +from tools.chats.schemas import MessageInSchema +from tools.chats.services.sse_chat_stream import SSEChatStreamService +from tools.chats.services.sse_store import PublicSSEStoreService + + +class OpenAIStreamService: + SSE_HEADERS = {'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'} + TERMINAL_EVENTS = frozenset({'done', 'error'}) + REASONING = {'effort': None, 'summary': None} + ITEM_IDX = {'output_index': 0, 'content_index': 0} + RESPONSE_DEFAULTS = { + 'parallel_tool_calls': True, + 'previous_response_id': None, + 'temperature': 1.0, + 'text': {'format': {'type': 'text'}}, + 'tool_choice': 'auto', + 'tools': [], + 'top_p': 1.0, + 'truncation': 'disabled', + 'user': None, + 'metadata': {}, + 'incomplete_details': None, + 'instructions': None, + 'max_output_tokens': None, + } + + def __init__(self, store: PublicSSEStoreService) -> None: + self.store = store + + @staticmethod + def parse_input(raw_input) -> str: + if raw_input is None: + raise HttpError(400, _('Missing required parameter: input')) + if isinstance(raw_input, str): + content = raw_input.strip() + elif not isinstance(raw_input, list): + raise HttpError(400, _('Invalid input payload')) + else: + lines = [] + for item in raw_input: + if not isinstance(item, dict): + continue + role = item.get('role', 'user').capitalize() + item_content = item.get('content') + if isinstance(item_content, str): + if text := item_content.strip(): + lines.append(f'[{role}] {text}') + continue + for part in item_content or []: + if ( + isinstance(part, dict) + and part.get('type') in ('input_text', 'text') + and (t := part.get('text')) + and (text := str(t).strip()) + ): + lines.append(f'[{role}] {text}') + content = '\n'.join(lines).strip() + if not content: + raise HttpError(400, _('The request must not be empty')) + return content + + @classmethod + def user_message_body(cls, body: dict) -> MessageInSchema: + kwargs = {'content': cls.parse_input(body.get('input'))} + if isinstance(body.get('metadata'), dict) and body['metadata']: + info = {} + for k, v in body['metadata'].items(): + if k in ('ttft', 'tbt') and isinstance(v, str): + try: + v = float(v) + except ValueError: + pass + info[k] = v + kwargs['info'] = info + return MessageInSchema(**kwargs) + + def cleanup_if_finished(self) -> None: + if self.store.exists() and any( + orjson.loads(raw).get('event') in self.TERMINAL_EVENTS + for raw in self.store.redis_client.lrange(self.store._get_cache_key(), 0, -1) + ): + self._cleanup() + + def init_meta(self, model: str, *, reconnect: bool = False) -> dict: + if meta := self._load_meta(): + return meta + if reconnect and not self.store.exists(): + raise HttpError(404, _('Stream not found')) + meta = { + 'rid': str(self.store.idempotency_key), + 'mid': f'msg_{secrets.token_hex(24)}', + 'ts': int(time.time()), + 'model': model, + 'user_uuid': str(self.store.user_uuid), + 'last_seq': -1, + 'internal_offset': 0, + 'completed': False, + 'setup_emitted': False, + } + self._save_meta(meta) + return meta + + def event_stream( + self, + request: HttpRequest | None = None, + *, + reconnect: bool = False, + meta: dict | None = None, + model: str = '', + ) -> Iterator[str]: + meta = meta or self.init_meta(model, reconnect=reconnect) + replayed: list[str] = [] + if reconnect: + if meta.get('completed'): + try: + yield from self._replay_events() + finally: + self._cleanup() + return + replayed = self._replay_events() + yield from replayed + if (self._load_meta() or {}).get('completed'): + self._cleanup() + return + if not reconnect or not replayed: + yield from self._emit_setup(meta) + try: + yield from self._map_internal_stream(meta, request) + except Exception as exc: + yield from self._finish_error(meta, exc) + self._cleanup() + + def _access(self) -> tuple[int, dict, dict]: + key = self.store._get_cache_key() + raws = self.store.redis_client.lrange(key, 0, -1) + if not raws: + chunk = {'event_id': 0, 'event': 'pending', 'data': {}} + self.store.redis_client.rpush(key, orjson.dumps(chunk)) + index = 0 + else: + index, chunk = 0, None + for i, raw in enumerate(raws): + chunk = orjson.loads(raw) + if chunk.get('event') == 'pending' or (chunk.get('data') or {}).get('openai') is not None: + index = i + break + else: + chunk = {'event_id': 0, 'event': 'pending', 'data': {}} + index = len(raws) + self.store.redis_client.rpush(key, orjson.dumps(chunk)) + bucket = chunk.setdefault('data', {}).setdefault('openai', {'meta': None, 'events': []}) + return index, chunk, bucket + + def _commit(self, index: int, chunk: dict) -> None: + self.store.redis_client.lset(self.store._get_cache_key(), index, orjson.dumps(chunk)) + + def _load_meta(self) -> dict | None: + if not self.store.exists(): + return None + return self._access()[2].get('meta') + + def _save_meta(self, meta: dict) -> None: + if self.store.exists(): + index, chunk, bucket = self._access() + bucket['meta'] = meta + self._commit(index, chunk) + + def _replay_events(self) -> list[str]: + if not self.store.exists(): + return [] + events = self._access()[2].get('events') or [] + return [e if isinstance(e, str) else e['sse'] for e in events] + + def _cleanup(self) -> None: + if self.store.exists(): + self.store.cleanup() + + @staticmethod + def _omit_none(value): + if isinstance(value, dict): + return {k: OpenAIStreamService._omit_none(v) for k, v in value.items() if v is not None} + if isinstance(value, list): + return [OpenAIStreamService._omit_none(v) for v in value] + return value + + @classmethod + def _sse(cls, event_type: str, seq: int, **fields) -> str: + payload = cls._omit_none({'type': event_type, 'sequence_number': seq, **fields}) + return f'event: {event_type}\ndata: {orjson.dumps(payload).decode()}\n\n' + + @staticmethod + def _parse_sse(chunk: str) -> tuple[str, dict, int | None]: + event, data, event_id = '', {}, None + for line in chunk.split('\n'): + if line.startswith('id:'): + event_id = int(line[3:].strip()) + elif line.startswith('event:'): + event = line[6:].strip() + elif line.startswith('data:'): + data = orjson.loads(line[5:].strip()) + return event, data, event_id + + def _response_obj(self, meta: dict, status: str, **fields): + return { + 'id': meta['rid'], + 'object': 'response', + 'created_at': meta['ts'], + 'model': meta['model'], + 'status': status, + 'error': None, + **self.RESPONSE_DEFAULTS, + **fields, + } + + def _emit(self, meta: dict, event_type: str, **fields) -> str: + meta['last_seq'] += 1 + sse = self._sse(event_type, meta['last_seq'], **fields) + index, chunk, bucket = self._access() + bucket['events'].append(sse) + self._commit(index, chunk) + return sse + + def _mark_done(self, meta: dict) -> None: + meta['completed'] = True + self._save_meta(meta) + + def _emit_setup(self, meta: dict) -> Iterator[str]: + if meta.get('setup_emitted'): + return + progress = self._response_obj(meta, 'in_progress', output=[], reasoning=self.REASONING, store=True, usage=None) + yield self._emit(meta, 'response.created', response=progress) + yield self._emit(meta, 'response.in_progress', response=progress) + yield self._emit( + meta, + 'response.output_item.added', + item={'id': meta['mid'], 'type': 'message', 'status': 'in_progress', 'role': 'assistant', 'content': []}, + **self.ITEM_IDX, + ) + yield self._emit( + meta, + 'response.content_part.added', + item_id=meta['mid'], + part={'type': 'output_text', 'text': '', 'annotations': []}, + **self.ITEM_IDX, + ) + meta['setup_emitted'] = True + self._save_meta(meta) + + def _finish_success(self, meta: dict, text: str) -> Iterator[str]: + part = {'type': 'output_text', 'text': text, 'annotations': []} + item = {'id': meta['mid'], 'type': 'message', 'status': 'completed', 'role': 'assistant', 'content': [part]} + ctx = {'item_id': meta['mid'], **self.ITEM_IDX} + yield self._emit(meta, 'response.output_text.done', text=text, **ctx) + yield self._emit(meta, 'response.content_part.done', part=part, **ctx) + yield self._emit(meta, 'response.output_item.done', item=item, output_index=0) + yield self._emit( + meta, + 'response.completed', + response=self._response_obj( + meta, + 'completed', + output=[item], + reasoning=self.REASONING, + store=True, + usage={ + 'input_tokens': 0, + 'output_tokens': 0, + 'output_tokens_details': {'reasoning_tokens': 0}, + 'total_tokens': 0, + }, + ), + ) + self._mark_done(meta) + + def _finish_error(self, meta: dict, message) -> Iterator[str]: + yield self._emit( + meta, + 'response.failed', + response=self._response_obj( + meta, 'failed', store=False, error={'code': 'server_error', 'message': str(message)} + ), + ) + self._mark_done(meta) + + def _map_internal_stream(self, meta: dict, request: HttpRequest | None) -> Iterator[str]: + text_parts = [] + ctx = {'item_id': meta['mid'], **self.ITEM_IDX} + try: + for chunk in SSEChatStreamService(self.store).event_stream( + request=request, offset=int(meta.get('internal_offset', 0)) + ): + if chunk == SSEChatStreamService.HEARTBEAT: + continue + event, data, event_id = self._parse_sse(chunk) + if event == 'token' and (token := data.get('content', '')): + text_parts.append(token) + yield self._emit(meta, 'response.output_text.delta', delta=token, **ctx) + elif event == 'done': + yield from self._finish_success(meta, data.get('content') or ''.join(text_parts)) + return + elif event == 'error': + yield from self._finish_error( + meta, data.get('detail', 'The model failed to generate a response.') + ) + return + if event_id is not None: + meta['internal_offset'] = int(event_id) + self._save_meta(meta) + finally: + if meta.get('completed'): + self._cleanup() @@ -95,9 +95,10 @@ DOMAIN=localhost PROVIDER=docker # UNLEASH -FEATURE_FLAG_API_URL=https://gitlab.kisulkens.ru/api/v4/feature_flags/unleash/243 -FEATURE_FLAG_INSTANCE_ID=glffct-ic8xsVF5eR9BaUySR-_w -FEATURE_FLAG_APP_NAME=Production +UNLEASH_API_URL=https://example.com +UNLEASH_APP_NAME=Development +UNLEASH_REQUEST_TIMEOUT=3 +UNLEASH_REQUEST_RETRIES=1 # ZROK ZROK2_API_ENDPOINT=https://zrok2.example.com @@ -108,5 +109,6 @@ COMPOSE_FILE=docker-compose.yml:docker-compose.local.yml COMPOSE_PROFILES="" # use-tunnel available BUILDKIT_PROGRESS=plain -# HTTP FRAMEWORK -GUNICORN_CMD_ARGS="-b 0.0.0.0:8000 -k gthread -w 3 -t 600 --reload" \ No newline at end of file +# SERVER +DJANGO_RUNSERVER_HIDE_WARNING=true +PYTHONWARNINGS=ignore::UserWarning:polymorphic # temporarily \ No newline at end of file @@ -7,6 +7,12 @@ x-dev-app-config: &dev-app-config services: app: <<: *dev-app-config + command: + - /bin/sh + - -c + - | + python manage.py compilemessages --locale ru_RU + python manage.py runserver 0.0.0.0:8000 ports: - "8000:8000" @@ -96,7 +96,7 @@ services: - /bin/sh - -c - | - python manage.py create_indexes --skip-system-checks + python manage.py create_indexes --skip-checks python manage.py initialize_buckets --skip-checks python manage.py collectstatic --no-input --skip-checks python manage.py migrate