Pular para conteúdo

Query database

call_rag(query, message_history_id)

Sends a query to a RAG (Retrieval-Augmented Generation) system.

Parameters:

Name Type Description Default
query str

The user's input query.

required
message_history_id int

The ID of the message history session.

required

Returns:

Name Type Description
dict dict

The response from the RAG system.

Source code in app/ui/pages/query_database.py
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def call_rag(query: str, message_history_id: int) -> dict:
    """
    Sends a query to a RAG (Retrieval-Augmented Generation) system.

    Args:
        query (str): 
            The user's input query.
        message_history_id (int): 
            The ID of the message history session.

    Returns:
        dict:
            The response from the RAG system.
    """
    response = api_request(
        api_url=f"http://0.0.0.0:8200/api/text-to-sql/",
        json={
             "message_history_id": message_history_id,
             "query": query
        }
    )
    return response

get_historic_message(message_history_id)

Retrieves the historic messages associated with a given message history ID.

Parameters:

Name Type Description Default
message_history_id int

The ID of the message history session.

required

Returns:

Name Type Description
list list

A list of previous messages in the conversation.

Source code in app/ui/pages/query_database.py
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def get_historic_message(message_history_id: int) -> list:
    """
    Retrieves the historic messages associated with a given message history ID.

    Args:
        message_history_id (int): 
            The ID of the message history session.

    Returns:
        list:
            A list of previous messages in the conversation.
    """
    response = api_request(
        api_url=f"http://0.0.0.0:8200/api/historic-message/",
        json={
                "message_history_id": message_history_id,
                "question": ""
            })

    return response

get_new_message_history_id()

Generates a new random message history ID.

Returns:

Name Type Description
int int

A randomly generated integer ID.

Source code in app/ui/pages/query_database.py
111
112
113
114
115
116
117
118
119
def get_new_message_history_id() -> int:
    """
    Generates a new random message history ID.

    Returns:
        int:
            A randomly generated integer ID.
    """
    return random.randint(0, 2**16 - 1)

return_historic()

Retrieves the chat history from Streamlit's session state.

Returns:

Name Type Description
list list

The stored message history, or an empty list if no history is found.

Source code in app/ui/pages/query_database.py
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def return_historic() -> list:
    """
    Retrieves the chat history from Streamlit's session state.

    Returns:
        list:
            The stored message history, or an empty list if no history is found.
    """
    if "historic" not in st.session_state:
        if "message_history_id_site" not in st.session_state:
            st.session_state.message_history_id_site = get_new_message_history_id()
            return []
        else:
            response = get_historic_message(
                st.session_state.message_history_id_site)
            return response
    return st.session_state.historic

return_last_message()

Retrieves the last message stored in Streamlit's session state.

Returns:

Type Description
list | dict

list | dict: The last message, or an empty list if no message is stored.

Source code in app/ui/pages/query_database.py
155
156
157
158
159
160
161
162
163
def return_last_message() -> list | dict:
    """
    Retrieves the last message stored in Streamlit's session state.

    Returns:
        list | dict:
            The last message, or an empty list if no message is stored.
    """
    return st.session_state.get("last_message", [])

update_historic(q_and_a)

Updates the chat history stored in Streamlit's session state.

Parameters:

Name Type Description Default
q_and_a list[dict]

A list of messages to be added to the session history.

required

Returns:

Type Description
None

None

Source code in app/ui/pages/query_database.py
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def update_historic(q_and_a: list[dict]) -> None:
    """
    Updates the chat history stored in Streamlit's session state.

    Args:
        q_and_a (list[dict]): 
            A list of messages to be added to the session history.

    Returns:
        None
    """
    if "historic" not in st.session_state:
        st.session_state.historic = q_and_a
    else:
        for msg in q_and_a:
            st.session_state.historic.append(msg)

update_last_message(msg)

Updates the last message stored in Streamlit's session state.

Parameters:

Name Type Description Default
msg dict

The message to be stored.

required

Returns:

Type Description
None

None

Source code in app/ui/pages/query_database.py
141
142
143
144
145
146
147
148
149
150
151
152
def update_last_message(msg: dict) -> None:
    """
    Updates the last message stored in Streamlit's session state.

    Args:
        msg (dict): 
            The message to be stored.

    Returns:
        None
    """
    st.session_state.last_message = msg

verify_sql_injection(query)

Checks if a given SQL query is vulnerable to SQL injection.

Parameters:

Name Type Description Default
query str

The SQL query to be checked.

required

Returns:

Name Type Description
str str

The security status of the query ("Secure" or "Insecure").

Source code in app/ui/pages/query_database.py
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
def verify_sql_injection(query: str) -> str:
    """
    Checks if a given SQL query is vulnerable to SQL injection.

    Args:
        query (str): 
            The SQL query to be checked.

    Returns:
        str:
            The security status of the query ("Secure" or "Insecure").
    """
    encoded_query = quote(query)
    response = api_request(
        api_url=f"http://0.0.0.0:8200/api/verify-sql-injection/{encoded_query}",
        json=None
    )['status']

    return response

write_user_and_assistant_messages(q_and_a)

Displays chat messages for both the user and assistant in a Streamlit app.

Parameters:

Name Type Description Default
q_and_a list[dict]

A list of messages with 'role' and 'content' keys.

required

Returns:

Type Description
None

None

Source code in app/ui/pages/query_database.py
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def write_user_and_assistant_messages(q_and_a: list[dict]) -> None:
    """
    Displays chat messages for both the user and assistant in a Streamlit app.

    Args:
        q_and_a (list[dict]): 
            A list of messages with 'role' and 'content' keys.

    Returns:
        None
    """
    for msg in q_and_a:
        owner = "assistant" if msg["role"] == "assistant" else "user"
        with st.chat_message(name=owner):
            st.write(msg["content"])