src/cores/client.py

#
# SPDX-FileCopyrightText: Hadad <[email protected]>
# SPDX-License-Identifier: Apache-2.0
#

import asyncio  # Import asyncio for asynchronous programming capabilities
import httpx  # Import httpx to perform asynchronous HTTP requests
import json  # Import json to handle JSON encoding and decoding
import random  # Import random to shuffle lists for load balancing
import uuid  # Import uuid to generate unique session identifiers

from config import *  # Import all configuration constants and variables from config module
from src.cores.server import fetch_response_stream_async  # Import async function to fetch streamed AI responses
from src.cores.session import ensure_stop_event, get_model_key  # Import session helper functions
from datetime import datetime  # Import datetime to get current date and time information

async def chat_with_model_async(history, user_input, model_display, sess, custom_prompt, deep_search):
    """
    Asynchronous function to handle interaction with an AI model and stream its responses.

    Parameters:
    - history: List of tuples containing previous conversation messages (user and assistant)
    - user_input: The current input string from the user
    - model_display: The display name of the AI model to use
    - sess: Session object containing session state, stop event, and cancellation token
    - custom_prompt: Optional custom system instructions to override default instructions
    - deep_search: Boolean flag indicating whether to integrate deep search results into the instructions

    This function prepares the message history and system instructions, optionally enriches the instructions
    with deep search results if enabled, and attempts to fetch streamed responses from multiple backend
    providers with fallback. It yields chunks of the response asynchronously for real-time UI updates.
    """

    # Ensure the session has a stop event initialized to control streaming cancellation
    ensure_stop_event(sess)

    # Clear any previous stop event state to allow new streaming session
    sess.stop_event.clear()

    # Reset the cancellation token to indicate the session is active and not cancelled
    sess.cancel_token["cancelled"] = False

    # Check if provider keys and hosts are configured; if not, yield a predefined error response and exit
    if not LINUX_SERVER_PROVIDER_KEYS or not LINUX_SERVER_HOSTS:
        yield ("content", RESPONSES["RESPONSE_3"])  # Inform user no backend providers are available
        return

    # Assign a unique session ID if not already present to track conversation context
    if not hasattr(sess, "session_id") or not sess.session_id:
        sess.session_id = str(uuid.uuid4())

    # Determine the internal model key based on the display name, falling back to default if not found
    model_key = get_model_key(model_display, MODEL_MAPPING, DEFAULT_MODEL_KEY)

    # Retrieve model-specific configuration parameters or use default configuration
    cfg = MODEL_CONFIG.get(model_key, DEFAULT_CONFIG)

    # Initialize a list to hold the messages that will be sent to the AI model
    msgs = []

    # Obtain the current date and time formatted as a readable string for context in instructions
    current_date = datetime.now().strftime("%A, %B %d, %Y, %I:%M %p %Z")

    # Combine internal AI instructions with the current date to form a comprehensive system instructions
    COMBINED_AI_INSTRUCTIONS = (
        INTERNAL_AI_INSTRUCTIONS
        + "\n\n\n"
        + f"Today is: {current_date}"
        + "\n\n\n"
    )

    # If deep search is enabled and the primary model is selected, prepend deep search instructions and results
    if deep_search and model_display == MODEL_CHOICES[0]:
        # Add deep search instructions as a system message to guide the AI
        msgs.append({"role": "system", "content": DEEP_SEARCH_INSTRUCTIONS})
        try:
            # Create an asynchronous HTTP client session for making the deep search request
            async with httpx.AsyncClient() as client:
                # Define the payload with parameters for the deep search query
                payload = {
                    "query": user_input,
                    "topic": "general",
                    "search_depth": "basic",
                    "chunks_per_source": 5,
                    "max_results": 5,
                    "time_range": None,
                    "days": 7,
                    "include_answer": True,
                    "include_raw_content": False,
                    "include_images": False,
                    "include_image_descriptions": False,
                    "include_domains": [],
                    "exclude_domains": []
                }
                # Send a POST request to the deep search provider with authorization header and JSON payload
                r = await client.post(
                    DEEP_SEARCH_PROVIDER_HOST,
                    headers={"Authorization": f"Bearer {DEEP_SEARCH_PROVIDER_KEY}"},
                    json=payload
                )
                # Parse the JSON response from the deep search provider
                sr_json = r.json()
                # Append the deep search results as a system message in JSON string format
                msgs.append({"role": "system", "content": json.dumps(sr_json)})
        except Exception:
            # If any error occurs during deep search, fail silently without interrupting the chat flow
            pass
        # Append the combined AI instructions after the deep search content to maintain context
        msgs.append({"role": "system", "content": COMBINED_AI_INSTRUCTIONS})

    # If deep search is not enabled but the primary model is selected, use only the combined AI instructions
    elif model_display == MODEL_CHOICES[0]:
        msgs.append({"role": "system", "content": COMBINED_AI_INSTRUCTIONS})

    # For other models, use a custom instructions if provided, otherwise default to the system instructions mapping or default instructions
    else:
        msgs.append({"role": "system", "content": custom_prompt or SYSTEM_PROMPT_MAPPING.get(model_key, SYSTEM_PROMPT_DEFAULT)})

    # Append the conversation history to the message list, alternating user and assistant messages
    # First add all user messages from history
    msgs.extend([{"role": "user", "content": u} for u, _ in history])
    # Then add all assistant messages from history that are not empty
    msgs.extend([{"role": "assistant", "content": a} for _, a in history if a])

    # Append the current user input as the latest user message
    msgs.append({"role": "user", "content": user_input})

    # Create a list of all possible combinations of backend hosts and provider keys for load balancing and fallback
    candidates = [(h, k) for h in LINUX_SERVER_HOSTS for k in LINUX_SERVER_PROVIDER_KEYS]

    # Randomly shuffle the list of host-key pairs to distribute load evenly and avoid bias
    random.shuffle(candidates)

    # Iterate over each host and key pair to attempt fetching a streamed response
    for h, k in candidates:
        # Call the async generator function to fetch streamed response chunks from the backend
        stream_gen = fetch_response_stream_async(
            h, k, model_key, msgs, cfg, sess.session_id, sess.stop_event, sess.cancel_token
        )

        # Flag to track if any response chunks were received from this provider
        got_responses = False

        # Asynchronously iterate over each chunk yielded by the streaming generator
        async for chunk in stream_gen:
            # If the stop event is set or cancellation requested, terminate streaming immediately
            if sess.stop_event.is_set() or sess.cancel_token["cancelled"]:
                return

            # Mark that at least one response chunk has been received
            got_responses = True

            # Yield the current chunk to the caller for incremental UI update or processing
            yield chunk

        # If any responses were received from this host-key pair, stop trying others and return
        if got_responses:
            return

    # If no responses were received from any provider, yield a fallback message indicating failure
    yield ("content", RESPONSES["RESPONSE_2"])