app.py

"""gradio app for Chat Interface for DataStax Langflow calls"""

from requests.exceptions import ConnectTimeout, ReadTimeout
import re
from datetime import datetime
from typing import Optional, Sequence, Tuple
from uuid import uuid4
import os
from retry import retry
import requests
import gradio as gr
from dotenv import load_dotenv

load_dotenv()


BASE_API_URL = "https://api.langflow.astra.datastax.com"
LANGFLOW_ID = "e99b525a-84dd-4cf2-bac6-8e5ebc6a7d17"
FLOW_ID = "bfbf7237-50dd-40dd-baba-cc680288d788"
ASTRA_DB_CLIENT_ID = os.getenv('ASTRA_DB_CLIENT_ID')
ASTRA_DB_ID = os.getenv('ASTRA_DB_ID')
ASTRA_DB_KEYSPACE = os.getenv('ASTRA_DB_KEYSPACE')
ASTRA_DB_REGION = os.getenv('ASTRA_DB_REGION')
ASTRA_DB_SECRET = os.getenv('ASTRA_DB_SECRET')
ASTRA_DB_TOKEN = os.getenv('ASTRA_DB_TOKEN')
DATASTAX_TOKEN = os.getenv('DATASTAX_TOKEN')

ENDPOINT = 'ai_traveller'

TWEAKS = {
    "Memory-fCuCs": {
        "n_messages": 100,
        "order": "Ascending",
        "sender": "Machine and User",
        "sender_name": "",
        "template": "{sender_name}: {text}"
    },
    "ChatInput-PN5E4": {},
    "Prompt-vGvVG": {},
    "GoogleGenerativeAIModel-6n0Ft": {},
    "ChatOutput-HjfF1": {}
}


def email_validator(_email: str) -> bool:
    """RegEx based Validate the email address."""
    return bool(re.match(r"[^@]+@[^@]+\.[^@]+", _email))


def contact_validator(_contact: str) -> bool:
    """RegEx based Validate the contact number."""
    return bool(re.match(r'^\d{10}$', _contact))


def required_validator(_name: str, _email: str, _contact: str) -> bool:
    """Validate the required fields."""
    return bool(_name.strip() and _email.strip() and _contact.strip())


def submit_signup(_email: str, _name: str, _contact: str):
    """Add signups to the database."""

    if not email_validator(_email):
        return 'Invalid Email'

    if not contact_validator(_contact):
        return 'Invalid Contact Number'

    if not required_validator(_name, _email, _contact):
        return 'All Fields Required for Sign Up'

    response = requests.post(
        f'https://{ASTRA_DB_ID}-{ASTRA_DB_REGION}.apps.astra.datastax.com/api/rest/v2/keyspaces/{ASTRA_DB_KEYSPACE}/signups',
        headers={
            'Content-Type': 'application/json',
            'X-Cassandra-Token': ASTRA_DB_TOKEN
        },
        json={
            'email': _email,
            'name': _name,
            'contact_no': _contact,
            'created_at': datetime.now().strftime('%Y-%m-%dT%H:%M:%SZ')
        },
        timeout=360
    )
    if response.status_code != 201:
        return 'Error'
    return 'Success'


@retry((ConnectTimeout, ReadTimeout), tries=3, delay=2)
def call_travel_ai(
    message: str,
    history: Sequence[Tuple[str, str]],
    local_storage: str,
    endpoint: str = ENDPOINT,
    output_type: str = "chat",
    input_type: str = "chat",
    tweaks: Optional[dict] = None,
    application_token: Optional[str] = DATASTAX_TOKEN
) -> dict:
    """
    Run a flow with a given message and optional tweaks.

    :param message: The message to send to the flow
    :param endpoint: The ID or the endpoint name of the flow
    :param tweaks: Optional tweaks to customize the flow
    :return: The JSON response from the flow
    """
    del history
    api_url = f"{BASE_API_URL}/lf/{LANGFLOW_ID}/api/v1/run/{endpoint}"

    payload = {
        'session_id': local_storage[0],
        "input_value": message,
        "output_type": output_type,
        "input_type": input_type,
    }
    headers = None
    if tweaks:
        payload["tweaks"] = tweaks
    if application_token:
        headers = {
            "Authorization": "Bearer " + application_token,
            "Content-Type": "application/json"
        }
    response = requests.post(
        api_url,
        json=payload,
        headers=headers,
        timeout=120,
    )
    response = response.json()
    resp_message = ''
    try:
        for resp in response['outputs']:
            for _resp in resp['outputs']:
                for message in _resp['messages']:
                    resp_message = message['message']
    except KeyError as e:
        print(e)
        print(response)

    return resp_message


def generate_session_id():
    """Generate a unique session ID."""
    return str(uuid4())


# def render(session_id: str):
#     """Render the session ID."""
#     return session_id


def on_load(local_storage):
    """Update the visibility of the sign up and chat interfaces."""
    signed_up = local_storage[1]
    return gr.update(visible=not signed_up), gr.update(visible=signed_up)


with gr.Blocks() as demo:
    with gr.Column():
        gr.Image(
            "./aem_logo.jpeg",
            width=300, label="AEM", show_download_button=False,
            show_fullscreen_button=False,
            show_label=False,
            interactive=False,
            show_share_button=False,
        )

    session = gr.BrowserState(["", False])
    # gr.Textbox(
    #     value=render,
    #     visible=False,
    #     interactive=False,
    #     label="Session ID",
    #     inputs=[session],
    # )

    with gr.Accordion(
        'Sign Up',
        visible=lambda session: not session[1],
    ) as signup:

        with gr.Column():
            name = gr.Textbox(label="Name")
            email = gr.Textbox(label="Email")
            phone = gr.Textbox(label="Phone")
            submit_button = gr.Button("Submit")

    with gr.Accordion(
        'Chat With Agent',
        visible=False,
    ) as chat:
        chat_interface = gr.ChatInterface(
            call_travel_ai,
            type='messages',
            title="AI Travel Partner",
            additional_inputs=[session],
            autofocus=True,
            fill_height=True,
        )

    def update_visibility(
        _name: str,
        _email: str,
        _phone: str,
        local_storage: str
    ):
        """Update the visibility of the sign up and chat interfaces."""
        # session_id, signed_up = local_storage
        local_storage[0] = local_storage[0] or generate_session_id()
        _status = submit_signup(_email, _name, _phone)
        gr.Info(_status, duration=5, title='Sign Up Status')
        if _status == 'Success':
            local_storage[1] = True

        return (
            gr.update(visible=not local_storage[1]),
            gr.update(visible=local_storage[1]),
            local_storage
        )

    submit_button.click(  # pylint: disable=no-member
        update_visibility,
        [name, email, phone, session],
        outputs=[signup, chat, session],
    )
    demo.load(  # pylint: disable=no-member
        on_load,
        [session],
        outputs=[signup, chat],
    )

demo.launch(share=False, debug=True,)