api/utils.py

from datetime import datetime
import json
from typing import Any, Dict, Optional

import httpx
from api.config import (
    MODEL_MAPPING,
    headers,
    AGENT_MODE,
    TRENDING_AGENT_MODE,
    BASE_URL,
    MODEL_PREFIXES,
    MODEL_REFERERS
)
from fastapi import HTTPException
from api.models import ChatRequest

from api.logger import setup_logger

import uuid  # Ensure this is imported if needed elsewhere

logger = setup_logger(__name__)

async def process_streaming_response(request: ChatRequest):
    agent_mode = AGENT_MODE.get(request.model, {})
    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
    referer_path = MODEL_REFERERS.get(request.model, f"/?model={request.model}")
    referer_url = f"{BASE_URL}{referer_path}"

    # Update headers with dynamic Referer
    dynamic_headers = headers.copy()
    dynamic_headers['Referer'] = referer_url

    json_data = {
        "messages": request.messages,  # Directly pass messages without processing
        "previewToken": None,
        "userId": None,
        "codeModelMode": True,
        "agentMode": agent_mode,
        "trendingAgentMode": trending_agent_mode,
        "isMicMode": False,
        "userSystemPrompt": None,
        "maxTokens": request.max_tokens,
        "playgroundTopP": request.top_p,
        "playgroundTemperature": request.temperature,
        "isChromeExt": False,
        "githubToken": None,
        "clickedAnswer2": False,
        "clickedAnswer3": False,
        "clickedForceWebSearch": False,
        "visitFromDelta": False,
        "mobileClient": False,
        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
    }

    async with httpx.AsyncClient() as client:
        try:
            async with client.stream(
                "POST",
                f"{BASE_URL}/api/chat",
                headers=dynamic_headers,
                json=json_data,
                timeout=100,
            ) as response:
                response.raise_for_status()
                async for chunk in response.aiter_text():
                    yield chunk  # Directly yield the chunk without modification
        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP error occurred: {e}")
            raise HTTPException(status_code=e.response.status_code, detail=str(e))
        except httpx.RequestError as e:
            logger.error(f"Error occurred during request: {e}")
            raise HTTPException(status_code=500, detail=str(e))

async def process_non_streaming_response(request: ChatRequest):
    agent_mode = AGENT_MODE.get(request.model, {})
    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
    referer_path = MODEL_REFERERS.get(request.model, f"/?model={request.model}")
    referer_url = f"{BASE_URL}{referer_path}"

    # Update headers with dynamic Referer
    dynamic_headers = headers.copy()
    dynamic_headers['Referer'] = referer_url

    json_data = {
        "messages": request.messages,  # Directly pass messages without processing
        "previewToken": None,
        "userId": None,
        "codeModelMode": True,
        "agentMode": agent_mode,
        "trendingAgentMode": trending_agent_mode,
        "isMicMode": False,
        "userSystemPrompt": None,
        "maxTokens": request.max_tokens,
        "playgroundTopP": request.top_p,
        "playgroundTemperature": request.temperature,
        "isChromeExt": False,
        "githubToken": None,
        "clickedAnswer2": False,
        "clickedAnswer3": False,
        "clickedForceWebSearch": False,
        "visitFromDelta": False,
        "mobileClient": False,
        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(
                f"{BASE_URL}/api/chat",
                headers=dynamic_headers,
                json=json_data,
                timeout=100,
            )
            response.raise_for_status()
            return response.json()  # Directly return the JSON response
        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP error occurred: {e}")
            raise HTTPException(status_code=e.response.status_code, detail=str(e))
        except httpx.RequestError as e:
            logger.error(f"Error occurred during request: {e}")
            raise HTTPException(status_code=500, detail=str(e))