File size: 16,961 Bytes
e448d98
32f5b77
 
 
 
 
 
95c52e2
e448d98
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c76addc
e448d98
c76addc
 
e448d98
c76addc
 
 
e448d98
c76addc
 
 
 
 
e448d98
4c67f45
 
e448d98
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4c67f45
 
76271c7
3d87c18
32f5b77
4c67f45
e448d98
76271c7
 
 
 
ac253c3
e448d98
 
0614630
 
e448d98
0614630
 
e448d98
c76addc
e448d98
c76addc
 
 
 
 
 
e448d98
c76addc
 
 
 
 
 
e448d98
c76addc
 
e448d98
 
c76addc
 
 
e448d98
c76addc
e448d98
3d87c18
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
from utils import message_dict

from haystack.dataclasses import ChatMessage
from haystack.components.generators.chat import OpenAIChatGenerator

chat_generator = OpenAIChatGenerator(model="gpt-4o")
response = None

def example_question_message(data_source, name, titles, schema):

    example_message_dict = {
        'file_upload' : ["You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source'.",
                         f"""We have a SQLite database with the following {titles}. 

                        We also have an AI agent with access to the same database that will be performing data analysis.

                        Please return an array of seven strings, each one being a question for our data analysis agent

                        that we can suggest that you believe will be insightful or helpful to a data analysis looking for

                        data insights. Return nothing more than the array of questions because I need that specific data structure

                        to process your response. No other response type or data structure will work."""],

        'sql' : [f"You are a helpful and knowledgeable agent who has access to an MongoDB NoSQL document database called {name}.",
                 f"""We have a PostgreSQL database with the following tables: {titles}. 

                        We also have an AI agent with access to the same database that will be performing data analysis.

                        Please return an array of seven strings, each one being a question for our data analysis agent

                        that we can suggest that you believe will be insightful or helpful to a data analysis looking for

                        data insights. Return nothing more than the array of questions because I need that specific data structure

                        to process your response. No other response type or data structure will work."""],

        'doc_db' : [f"You are a helpful and knowledgeable agent who has access to an MongoDB NoSQL document database called {name}.",
                    f"""We have a MongoDB NoSQL document database with the following collections: {titles}.

                        The schema of these collections is: {schema}. 

                        We also have an AI agent with access to the same database that will be performing data analysis.

                        Please return an array of seven strings, each one being a question for our data analysis agent

                        that we can suggest that you believe will be insightful or helpful to a data analysis looking for

                        data insights. Return nothing more than the array of questions because I need that specific data structure

                        to process your response. No other response type or data structure will work."""],

        'graphql' : [f"You are a helpful and knowledgeable agent who has access to an GraphQL API endpoint called {name}.",
                     f"""We have a GraphQL API endpoint with the following types: {titles}.

                        We also have an AI agent with access to the same GraphQL API endpoint that will be performing data analysis.

                        Please return an array of seven strings, each one being a question for our data analysis agent

                        that we can suggest that you believe will be insightful or helpful to a data analysis looking for

                        data insights. Return nothing more than the array of questions because I need that specific data structure

                        to process your response. No other response type or data structure will work."""]         

    }

    return example_message_dict[data_source]

def example_question_generator(session_hash, data_source, name, titles, schema):
    example_response = None
    example_message_list = example_question_message(data_source, name, titles, schema)
    example_messages = [
        ChatMessage.from_system(
            example_message_list[0]
        )
    ]

    example_messages.append(ChatMessage.from_user(text=example_message_list[1]))

    example_response = chat_generator.run(messages=example_messages)

    return example_response["replies"][0].text

def system_message(data_source, titles, schema=""):
    print("TITLES")
    print(titles)
    system_message_dict = {
        'file_upload' : f"""You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source' that contains the following columns: {titles}. 

                    You also have access to a function, called table_generation_func, that can take a query.csv file generated from our sql query and returns an iframe that we should display in our chat window.

                    You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.

                    You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.

                    You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.

                    You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.

                    You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.

                    You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our sql query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.

                    Could you please always display the generated charts, tables, and visualizations as part of your output?""",

        'sql' : f"""You are a helpful and knowledgeable agent who has access to an PostgreSQL database which has a series of tables called {titles}. 

                    You also have access to a function, called table_generation_func, that can take a query.csv file generated from our sql query and returns an iframe that we should display in our chat window.

                    You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.

                    You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.

                    You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.

                    You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.

                    You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.

                    You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our sql query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.

                    Could you please always display the generated charts, tables, and visualizations as part of your output?""",

        'doc_db' : f"""You are a helpful and knowledgeable agent who has access to a NoSQL MongoDB Document database which has a series of collections called {titles}.

                    The schema of these collections is: {schema}.

                    You also have access to a function, called table_generation_func, that can take a query.csv file generated from our MongoDB query and returns an iframe that we should display in our chat window.

                    You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.

                    You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.

                    You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.

                    You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.

                    You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.

                    You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our MongoDB query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.

                    Could you please always display the generated charts, tables, and visualizations as part of your output?""",

        'graphql' : f"""You are a helpful and knowledgeable agent who has access to a GraphQL API which has the following types: {titles}.

                    We have also saved a schema.json file that contains the entire introspection query that we can use to find out more about each type before making a query.

                    You also have access to a function, called table_generation_func, that can take a query.csv file generated from our GraphQL API query and returns an iframe that we should display in our chat window.

                    You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.

                    You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.

                    You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.

                    You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.

                    You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.

                    You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our GraphQL API query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.

                    Could you please always display the generated charts, tables, and visualizations as part of your output?"""            

    }

    return system_message_dict[data_source]

def chatbot_func(message, history, session_hash, data_source, titles, schema, *args):
    from functions import table_generation_func, regression_func, scatter_chart_generation_func, \
        query_func, graphql_schema_query, graphql_csv_query, \
        line_chart_generation_func,bar_chart_generation_func,pie_chart_generation_func,histogram_generation_func
    import tools.tools as tools

    available_functions = {"query_func":query_func,"graphql_schema_query": graphql_schema_query,"graphql_csv_query": graphql_csv_query,
                           "table_generation_func":table_generation_func,
                           "line_chart_generation_func":line_chart_generation_func,"bar_chart_generation_func":bar_chart_generation_func,
                           "scatter_chart_generation_func":scatter_chart_generation_func, "pie_chart_generation_func":pie_chart_generation_func,
                           "histogram_generation_func":histogram_generation_func,
                           "regression_func":regression_func }
    
    if message_dict[session_hash][data_source] != None:
        message_dict[session_hash][data_source].append(ChatMessage.from_user(message))
    else:
        messages = [
            ChatMessage.from_system(system_message(data_source, titles, schema))
        ]
        messages.append(ChatMessage.from_user(message))
        message_dict[session_hash][data_source] = messages
    
    response = chat_generator.run(messages=message_dict[session_hash][data_source], generation_kwargs={"tools": tools.tools_call(session_hash, data_source, titles)})

    while True:
        # if OpenAI response is a tool call
        if response and response["replies"][0].meta["finish_reason"] == "tool_calls" or response["replies"][0].tool_calls:
            function_calls = response["replies"][0].tool_calls
            for function_call in function_calls:
                message_dict[session_hash][data_source].append(ChatMessage.from_assistant(tool_calls=[function_call]))
                ## Parse function calling information
                function_name = function_call.tool_name
                function_args = function_call.arguments

                ## Find the corresponding function and call it with the given arguments
                function_to_call = available_functions[function_name]
                function_response = function_to_call(**function_args, session_hash=session_hash, session_folder=data_source, args=args)
                print(function_name)
                ## Append function response to the messages list using `ChatMessage.from_tool`
                message_dict[session_hash][data_source].append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
                response = chat_generator.run(messages=message_dict[session_hash][data_source], generation_kwargs={"tools": tools.tools_call(session_hash, data_source, titles)})

        # Regular Conversation
        else:
            message_dict[session_hash][data_source].append(response["replies"][0])
            break

    return response["replies"][0].text