|
| 1 | +import asyncio |
| 2 | +import json |
| 3 | +import random |
| 4 | +import re |
| 5 | +import string |
| 6 | +import time |
| 7 | +import uuid |
| 8 | + |
| 9 | +import pybase64 |
| 10 | +import websockets |
| 11 | +from fastapi import HTTPException |
| 12 | + |
| 13 | +from api.files import get_file_content |
| 14 | +from api.models import model_system_fingerprint |
| 15 | +from api.tokens import split_tokens_from_content, calculate_image_tokens, num_tokens_from_messages |
| 16 | +from utils.Logger import logger |
| 17 | + |
| 18 | +moderation_message = "I'm sorry, I cannot provide or engage in any content related to pornography, violence, or any unethical material. If you have any other questions or need assistance, please feel free to let me know. I'll do my best to provide support and assistance." |
| 19 | + |
| 20 | + |
| 21 | +async def format_not_stream_response(response, prompt_tokens, max_tokens, model): |
| 22 | + chat_id = f"chatcmpl-{''.join(random.choice(string.ascii_letters + string.digits) for _ in range(29))}" |
| 23 | + system_fingerprint_list = model_system_fingerprint.get(model, None) |
| 24 | + system_fingerprint = random.choice(system_fingerprint_list) if system_fingerprint_list else None |
| 25 | + created_time = int(time.time()) |
| 26 | + all_text = "" |
| 27 | + async for chunk in response: |
| 28 | + try: |
| 29 | + if chunk.startswith("data: [DONE]"): |
| 30 | + break |
| 31 | + elif not chunk.startswith("data: "): |
| 32 | + continue |
| 33 | + else: |
| 34 | + chunk = json.loads(chunk[6:]) |
| 35 | + if not chunk["choices"][0].get("delta"): |
| 36 | + continue |
| 37 | + all_text += chunk["choices"][0]["delta"]["content"] |
| 38 | + except Exception as e: |
| 39 | + logger.error(f"Error: {chunk}, error: {str(e)}") |
| 40 | + continue |
| 41 | + content, completion_tokens, finish_reason = await split_tokens_from_content(all_text, max_tokens, model) |
| 42 | + message = { |
| 43 | + "role": "assistant", |
| 44 | + "content": content, |
| 45 | + } |
| 46 | + usage = { |
| 47 | + "prompt_tokens": prompt_tokens, |
| 48 | + "completion_tokens": completion_tokens, |
| 49 | + "total_tokens": prompt_tokens + completion_tokens |
| 50 | + } |
| 51 | + if not message.get("content"): |
| 52 | + raise HTTPException(status_code=403, detail="No content in the message.") |
| 53 | + |
| 54 | + data = { |
| 55 | + "id": chat_id, |
| 56 | + "object": "chat.completion", |
| 57 | + "created": created_time, |
| 58 | + "model": model, |
| 59 | + "choices": [ |
| 60 | + { |
| 61 | + "index": 0, |
| 62 | + "message": message, |
| 63 | + "logprobs": None, |
| 64 | + "finish_reason": finish_reason |
| 65 | + } |
| 66 | + ], |
| 67 | + "usage": usage |
| 68 | + } |
| 69 | + if system_fingerprint: |
| 70 | + data["system_fingerprint"] = system_fingerprint |
| 71 | + return data |
| 72 | + |
| 73 | + |
| 74 | +async def head_process_response(response): |
| 75 | + async for chunk in response: |
| 76 | + chunk = chunk.decode("utf-8") |
| 77 | + if chunk.startswith("data: {"): |
| 78 | + chunk_old_data = json.loads(chunk[6:]) |
| 79 | + message = chunk_old_data.get("message", {}) |
| 80 | + if not message and "error" in chunk_old_data: |
| 81 | + return response, False |
| 82 | + role = message.get('author', {}).get('role') |
| 83 | + if role == 'user' or role == 'system': |
| 84 | + continue |
| 85 | + |
| 86 | + status = message.get("status") |
| 87 | + if status == "in_progress": |
| 88 | + return response, True |
| 89 | + return response, False |
| 90 | + |
| 91 | + |
| 92 | +async def stream_response(service, response, model, max_tokens): |
| 93 | + chat_id = f"chatcmpl-{''.join(random.choice(string.ascii_letters + string.digits) for _ in range(29))}" |
| 94 | + created_time = int(time.time()) |
| 95 | + |
| 96 | + chunk_new_data = { |
| 97 | + "id": chat_id, |
| 98 | + "object": "chat.completion.chunk", |
| 99 | + "created": created_time, |
| 100 | + "model": model, |
| 101 | + "choices": [ |
| 102 | + { |
| 103 | + "index": 0, |
| 104 | + "delta": {"role": "assistant", "content": ""}, |
| 105 | + "logprobs": None, |
| 106 | + "finish_reason": None |
| 107 | + } |
| 108 | + ] |
| 109 | + } |
| 110 | + yield f"data: {json.dumps(chunk_new_data)}\n\n" |
| 111 | + |
| 112 | + async for chunk in response: |
| 113 | + chunk = chunk.decode("utf-8") |
| 114 | + try: |
| 115 | + if chunk.startswith("data: {"): |
| 116 | + chunk_old_data = json.loads(chunk[6:].strip()) |
| 117 | + except Exception as e: |
| 118 | + logger.error(f"Error: {chunk}, error: {str(e)}") |
| 119 | + continue |
| 120 | + |
| 121 | + |
| 122 | +def get_url_from_content(content): |
| 123 | + if isinstance(content, str) and content.startswith('http'): |
| 124 | + try: |
| 125 | + url = re.match( |
| 126 | + r'(?i)\b((?:[a-z][\w-]+:(?:/{1,3}|[a-z0-9%])|www\d{0,3}[.]|[a-z0-9.\-]+[.][a-z]{2,4}/)(?:[^\s()<>]+|\(([^\s()<>]+|(\([^\s()<>]+\)))*\))+(?:\(([^\s()<>]+|(\([^\s()<>]+\)))*\)|[^\s`!()\[\]{};:\'".,<>?«»“”‘’]))', |
| 127 | + content.split(' ')[0])[0] |
| 128 | + content = content.replace(url, '').strip() |
| 129 | + return url, content |
| 130 | + except Exception: |
| 131 | + return None, content |
| 132 | + return None, content |
| 133 | + |
| 134 | + |
| 135 | +def format_messages_with_url(content): |
| 136 | + url_list = [] |
| 137 | + while True: |
| 138 | + url, content = get_url_from_content(content) |
| 139 | + if url: |
| 140 | + url_list.append(url) |
| 141 | + logger.info(f"Found a file_url from messages: {url}") |
| 142 | + else: |
| 143 | + break |
| 144 | + if not url_list: |
| 145 | + return content |
| 146 | + new_content = [ |
| 147 | + { |
| 148 | + "type": "text", |
| 149 | + "text": content |
| 150 | + } |
| 151 | + ] |
| 152 | + for url in url_list: |
| 153 | + new_content.append({ |
| 154 | + "type": "image_url", |
| 155 | + "image_url": { |
| 156 | + "url": url |
| 157 | + } |
| 158 | + }) |
| 159 | + return new_content |
| 160 | + |
| 161 | + |
| 162 | +async def api_messages_to_chat(service, api_messages, upload_by_url=False): |
| 163 | + file_tokens = 0 |
| 164 | + chat_messages = [] |
| 165 | + for api_message in api_messages: |
| 166 | + role = api_message.get('role') |
| 167 | + content = api_message.get('content') |
| 168 | + if upload_by_url: |
| 169 | + if isinstance(content, str): |
| 170 | + content = format_messages_with_url(content) |
| 171 | + if isinstance(content, list): |
| 172 | + parts = [] |
| 173 | + attachments = [] |
| 174 | + content_type = "multimodal_text" |
| 175 | + for i in content: |
| 176 | + if i.get("type") == "text": |
| 177 | + parts.append(i.get("text")) |
| 178 | + elif i.get("type") == "image_url": |
| 179 | + image_url = i.get("image_url") |
| 180 | + url = image_url.get("url") |
| 181 | + detail = image_url.get("detail", "auto") |
| 182 | + file_content, mime_type = await get_file_content(url) |
| 183 | + file_meta = await service.upload_file(file_content, mime_type) |
| 184 | + if file_meta: |
| 185 | + file_id = file_meta["file_id"] |
| 186 | + file_size = file_meta["size_bytes"] |
| 187 | + file_name = file_meta["file_name"] |
| 188 | + mime_type = file_meta["mime_type"] |
| 189 | + use_case = file_meta["use_case"] |
| 190 | + if mime_type.startswith("image/"): |
| 191 | + width, height = file_meta["width"], file_meta["height"] |
| 192 | + file_tokens += await calculate_image_tokens(width, height, detail) |
| 193 | + parts.append({ |
| 194 | + "content_type": "image_asset_pointer", |
| 195 | + "asset_pointer": f"file-service://{file_id}", |
| 196 | + "size_bytes": file_size, |
| 197 | + "width": width, |
| 198 | + "height": height |
| 199 | + }) |
| 200 | + attachments.append({ |
| 201 | + "id": file_id, |
| 202 | + "size": file_size, |
| 203 | + "name": file_name, |
| 204 | + "mime_type": mime_type, |
| 205 | + "width": width, |
| 206 | + "height": height |
| 207 | + }) |
| 208 | + else: |
| 209 | + if not use_case == "ace_upload": |
| 210 | + await service.check_upload(file_id) |
| 211 | + file_tokens += file_size // 1000 |
| 212 | + attachments.append({ |
| 213 | + "id": file_id, |
| 214 | + "size": file_size, |
| 215 | + "name": file_name, |
| 216 | + "mime_type": mime_type, |
| 217 | + }) |
| 218 | + metadata = { |
| 219 | + "attachments": attachments |
| 220 | + } |
| 221 | + else: |
| 222 | + content_type = "text" |
| 223 | + parts = [content] |
| 224 | + metadata = {} |
| 225 | + chat_message = { |
| 226 | + "id": f"{uuid.uuid4()}", |
| 227 | + "author": {"role": role}, |
| 228 | + "content": {"content_type": content_type, "parts": parts}, |
| 229 | + "metadata": metadata |
| 230 | + } |
| 231 | + chat_messages.append(chat_message) |
| 232 | + text_tokens = await num_tokens_from_messages(api_messages, service.resp_model) |
| 233 | + prompt_tokens = text_tokens + file_tokens |
| 234 | + return chat_messages, prompt_tokens |
0 commit comments