test24 / api /utils.py
Niansuh's picture
Update api/utils.py
c9b2576 verified
raw
history blame
8.26 kB
from datetime import datetime
import json
import uuid
import asyncio
import random
import string
from typing import Any, Dict, Optional, AsyncGenerator
import httpx
from fastapi import HTTPException
from api.config import (
models,
model_aliases,
ALLOWED_MODELS,
MODEL_MAPPING,
get_headers_api_chat,
BASE_URL,
)
from api.models import ChatRequest, Message
from api.logger import setup_logger
logger = setup_logger(__name__)
# Editee API endpoint
EDITE_API_ENDPOINT = "https://editee.com/submit/chatgptfree"
# Helper function to create 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))
# Helper 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,
}
# Function to convert message to dictionary format, ensuring base64 data
def message_to_dict(message: Message):
if isinstance(message.content, str):
content = message.content
elif isinstance(message.content, list) and isinstance(message.content[0], dict) and "text" in message.content[0]:
content = message.content[0]["text"]
else:
content = ""
if isinstance(message.content, list) and len(message.content) == 2 and "image_url" in message.content[1]:
# Ensure base64 images are always included for all models
return {
"role": message.role,
"content": content,
"data": {
"imageBase64": message.content[1]["image_url"]["url"],
"fileText": "",
"title": "snapshot",
},
}
return {"role": message.role, "content": content}
# Function to strip model prefix from content if present (Removed as MODEL_PREFIXES is removed)
# def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
# """Remove the model prefix from the response content if present."""
# if model_prefix and content.startswith(model_prefix):
# logger.debug(f"Stripping prefix '{model_prefix}' from content.")
# return content[len(model_prefix):].strip()
# return content
# Function to get the correct referer URL for logging (Removed as MODEL_REFERERS is removed)
# def get_referer_url(chat_id: str, model: str) -> str:
# """Generate the referer URL based on specific models listed in MODEL_REFERERS."""
# if model in MODEL_REFERERS:
# return f"{BASE_URL}/chat/{chat_id}?model={model}"
# return BASE_URL
# Function to resolve model aliases
def resolve_model(model: str) -> str:
if model in MODEL_MAPPING:
return model
elif model in model_aliases:
return model_aliases[model]
else:
logger.warning(f"Model '{model}' not recognized. Using default model '{default_model}'.")
return "claude" # default_model
# Process streaming response with headers from config.py
async def process_streaming_response(request: ChatRequest) -> AsyncGenerator[str, None]:
chat_id = generate_chat_id()
resolved_model = resolve_model(request.model)
# referer_url = get_referer_url(chat_id, resolved_model) # Removed
logger.info(f"Generated Chat ID: {chat_id} - Model: {resolved_model}")
# model_prefix = MODEL_PREFIXES.get(resolved_model, "") # Removed
headers_api_chat = get_headers_api_chat(BASE_URL) # Using BASE_URL as referer
# Removed agent mode and delay logic
prompt = format_prompt(request.messages)
data = {
"user_input": prompt,
"context": " ",
"template_id": "",
"selected_model": resolved_model
}
async with httpx.AsyncClient() as client:
try:
response = await client.post(
EDITE_API_ENDPOINT,
headers=headers_api_chat,
json=data,
timeout=100,
)
response.raise_for_status()
response_data = response.json()
# Assuming response_data contains 'text' field
text = response_data.get('text', '')
timestamp = int(datetime.now().timestamp())
if text:
# cleaned_content = strip_model_prefix(text, model_prefix) # Removed
yield f"data: {json.dumps(create_chat_completion_data(text, resolved_model, timestamp))}\n\n"
# Indicate completion
yield f"data: {json.dumps(create_chat_completion_data('', resolved_model, timestamp, 'stop'))}\n\n"
yield "data: [DONE]\n\n"
except httpx.HTTPStatusError as e:
logger.error(f"HTTP error occurred for Chat ID {chat_id}: {e.response.status_code} - {e.response.text}")
raise HTTPException(status_code=e.response.status_code, detail=str(e))
except httpx.RequestError as e:
logger.error(f"Request error occurred for Chat ID {chat_id}: {e}")
raise HTTPException(status_code=500, detail=str(e))
# Process non-streaming response with headers from config.py
async def process_non_streaming_response(request: ChatRequest) -> Dict[str, Any]:
chat_id = generate_chat_id()
resolved_model = resolve_model(request.model)
# referer_url = get_referer_url(chat_id, resolved_model) # Removed
logger.info(f"Generated Chat ID: {chat_id} - Model: {resolved_model}")
# model_prefix = MODEL_PREFIXES.get(resolved_model, "") # Removed
headers_api_chat = get_headers_api_chat(BASE_URL) # Using BASE_URL as referer
# Removed agent mode and delay logic
prompt = format_prompt(request.messages)
data = {
"user_input": prompt,
"context": " ",
"template_id": "",
"selected_model": resolved_model
}
async with httpx.AsyncClient() as client:
try:
response = await client.post(
EDITE_API_ENDPOINT,
headers=headers_api_chat,
json=data,
timeout=100,
)
response.raise_for_status()
response_data = response.json()
text = response_data.get('text', '')
# if text.startswith("$@$v=undefined-rv1$@$"):
# text = text[21:] # Removed
# cleaned_full_response = strip_model_prefix(text, model_prefix) # Removed
return {
"id": f"chatcmpl-{uuid.uuid4()}",
"object": "chat.completion",
"created": int(datetime.now().timestamp()),
"model": resolved_model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": text},
"finish_reason": "stop",
}
],
"usage": None,
}
except httpx.HTTPStatusError as e:
logger.error(f"HTTP error occurred for Chat ID {chat_id}: {e.response.status_code} - {e.response.text}")
raise HTTPException(status_code=e.response.status_code, detail=str(e))
except httpx.RequestError as e:
logger.error(f"Request error occurred for Chat ID {chat_id}: {e}")
raise HTTPException(status_code=500, detail=str(e))
# Helper function to format prompt from messages
def format_prompt(messages: list[Message]) -> str:
# Implement the prompt formatting as per Editee's requirements
# Placeholder implementation
formatted_messages = []
for msg in messages:
if isinstance(msg.content, str):
formatted_messages.append(msg.content)
elif isinstance(msg.content, list):
text = msg.content[0].get("text", "")
formatted_messages.append(text)
return "\n".join(formatted_messages)