test24 / api /utils.py
Niansuh's picture
Update api/utils.py
a93ac0e verified
raw
history blame
8.46 kB
from datetime import datetime
import json
import uuid
import asyncio
import random
import string
from typing import Any, Dict, Optional
import httpx
from fastapi import HTTPException
from api.config import (
MODEL_MAPPING,
get_headers_api_chat,
get_headers_chat,
BASE_URL,
AGENT_MODE,
TRENDING_AGENT_MODE,
MODEL_PREFIXES,
MODEL_REFERERS,
)
from api.models import ChatRequest
from api.logger import (
log_generated_chat_id_with_referer,
log_model_delay,
log_http_error,
log_request_error,
log_strip_prefix,
)
# Generate a random alphanumeric chat ID
def generate_chat_id(length: int = 7) -> str:
characters = string.ascii_letters + string.digits
return ''.join(random.choices(characters, k=length))
# Function to create chat completion data
def create_chat_completion_data(
content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
) -> Dict[str, Any]:
return {
"id": f"chatcmpl-{uuid.uuid4()}",
"object": "chat.completion.chunk",
"created": timestamp,
"model": model,
"choices": [
{
"index": 0,
"delta": {"content": content, "role": "assistant"},
"finish_reason": finish_reason,
}
],
"usage": None,
}
# Convert message to dictionary format with optional model prefix
def message_to_dict(message, model_prefix: Optional[str] = None):
content = message.content
if model_prefix:
content = f"{model_prefix} {content}"
return {"role": message.role, "content": content}
# Remove model prefix from content if present
def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
if model_prefix and content.startswith(model_prefix):
log_strip_prefix(model_prefix, content)
return content[len(model_prefix):].strip()
return content
# Streaming response with headers from config.py
async def process_streaming_response(request: ChatRequest):
# Generate chat_id only if model is in MODEL_REFERERS
chat_id = generate_chat_id() if request.model in MODEL_REFERERS else None
referer_path = MODEL_REFERERS.get(request.model, "")
referer_url = f"{BASE_URL}/chat/{chat_id}?model={request.model}" if chat_id else BASE_URL
# Log chat ID, model, and referer URL if applicable
log_generated_chat_id_with_referer(chat_id, request.model, referer_url)
# Fetch model configurations
agent_mode = AGENT_MODE.get(request.model, {})
trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
model_prefix = MODEL_PREFIXES.get(request.model, "")
# Dynamic headers based on referer
headers_api_chat = get_headers_api_chat(referer_url)
# Delay for 'o1-preview' model
if request.model == 'o1-preview':
delay_seconds = random.randint(20, 60)
log_model_delay(delay_seconds, request.model, chat_id)
await asyncio.sleep(delay_seconds)
json_data = {
"agentMode": agent_mode,
"clickedAnswer2": False,
"clickedAnswer3": False,
"clickedForceWebSearch": False,
"codeModelMode": True,
"githubToken": None,
"id": chat_id,
"isChromeExt": False,
"isMicMode": False,
"maxTokens": request.max_tokens,
"messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
"mobileClient": False,
"playgroundTemperature": request.temperature,
"playgroundTopP": request.top_p,
"previewToken": None,
"trendingAgentMode": trending_agent_mode,
"userId": None,
"userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
"userSystemPrompt": None,
"validated": "69783381-2ce4-4dbd-ac78-35e9063feabc",
"visitFromDelta": False,
}
async with httpx.AsyncClient() as client:
try:
async with client.stream(
"POST",
f"{BASE_URL}/api/chat",
headers=headers_api_chat,
json=json_data,
timeout=100,
) as response:
response.raise_for_status()
async for line in response.aiter_lines():
timestamp = int(datetime.now().timestamp())
if line:
content = line.lstrip("$@$v=undefined-rv1$@$")
cleaned_content = strip_model_prefix(content, model_prefix)
yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
yield "data: [DONE]\n\n"
except httpx.HTTPStatusError as e:
log_http_error(e, chat_id)
raise HTTPException(status_code=e.response.status_code, detail=str(e))
except httpx.RequestError as e:
log_request_error(e, chat_id)
raise HTTPException(status_code=500, detail=str(e))
# Non-streaming response with headers from config.py
async def process_non_streaming_response(request: ChatRequest):
chat_id = generate_chat_id() if request.model in MODEL_REFERERS else None
referer_path = MODEL_REFERERS.get(request.model, "")
referer_url = f"{BASE_URL}/chat/{chat_id}?model={request.model}" if chat_id else BASE_URL
# Log with chat ID, model, and referer URL if applicable
log_generated_chat_id_with_referer(chat_id, request.model, referer_url)
# Model configurations
agent_mode = AGENT_MODE.get(request.model, {})
trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
model_prefix = MODEL_PREFIXES.get(request.model, "")
# Headers for API chat and secondary request
headers_api_chat = get_headers_api_chat(referer_url)
headers_chat = get_headers_chat(referer_url, next_action=str(uuid.uuid4()), next_router_state_tree=json.dumps([""]))
# Delay for 'o1-preview' model
if request.model == 'o1-preview':
delay_seconds = random.randint(20, 60)
log_model_delay(delay_seconds, request.model, chat_id)
await asyncio.sleep(delay_seconds)
json_data = {
"agentMode": agent_mode,
"clickedAnswer2": False,
"clickedAnswer3": False,
"clickedForceWebSearch": False,
"codeModelMode": True,
"githubToken": None,
"id": chat_id,
"isChromeExt": False,
"isMicMode": False,
"maxTokens": request.max_tokens,
"messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
"mobileClient": False,
"playgroundTemperature": request.temperature,
"playgroundTopP": request.top_p,
"previewToken": None,
"trendingAgentMode": trending_agent_mode,
"userId": None,
"userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
"userSystemPrompt": None,
"validated": "69783381-2ce4-4dbd-ac78-35e9063feabc",
"visitFromDelta": False,
}
full_response = ""
async with httpx.AsyncClient() as client:
try:
async with client.stream(
method="POST", url=f"{BASE_URL}/api/chat", headers=headers_api_chat, json=json_data
) as response:
response.raise_for_status()
async for chunk in response.aiter_text():
full_response += chunk
except httpx.HTTPStatusError as e:
log_http_error(e, chat_id)
raise HTTPException(status_code=e.response.status_code, detail=str(e))
except httpx.RequestError as e:
log_request_error(e, chat_id)
raise HTTPException(status_code=500, detail=str(e))
if full_response.startswith("$@$v=undefined-rv1$@$"):
full_response = full_response[21:]
# Strip the model prefix
cleaned_full_response = strip_model_prefix(full_response, model_prefix)
return {
"id": f"chatcmpl-{uuid.uuid4()}",
"object": "chat.completion",
"created": int(datetime.now().timestamp()),
"model": request.model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": cleaned_full_response},
"finish_reason": "stop",
}
],
"usage": None,
}