HighRoller/Ai_Agent

1from imports import *
2from pipelines import *
3from hal_check import *
4from prompts import *
5from prompt_templates import * 
6from retrieval import *
7from all_apis import *
8from prompts import *
9from update_api_toolset import *
10from mem_check import *
11warnings.filterwarnings('ignore')
12
13# retrieval examples
14retrieval_loader = CSVLoader(file_path=r'Seed_Dataset.csv', source_column = 'QUERY')
15retrieval_data = retrieval_loader.load()
16retrieval_embeddings = HuggingFaceEmbeddings()
17
18if "available_tools" not in st.session_state:
19    st.session_state.available_tools = available_tools.copy()
20    
21if "available_arguments" not in st.session_state:
22    st.session_state.available_arguments = available_arguments.copy()
23    
24if "arg_allowed_values_dict" not in st.session_state:
25    st.session_state.arg_allowed_values_dict = arg_allowed_values_dict.copy()
26    
27if "args_in_list_dict" not in st.session_state:
28    st.session_state.args_in_list_dict = args_in_list_dict.copy()
29    
30if "api_list_updated" not in st.session_state:
31    st.session_state.api_list_updated = API_LIST.copy()
32
33if "retrieval_vector_db" not in st.session_state:
34    st.session_state.retrieval_vector_db = FAISS.from_documents(
35        documents=retrieval_data,
36        embedding=retrieval_embeddings,
37        )
38
39# Initialize session state for messages if not already present
40if "messages" not in st.session_state:
41    st.session_state.messages = []
42
43if "PAST_QUERY" not in st.session_state:
44    st.session_state.PAST_QUERY = "NO PAST QUERIES"
45    st.session_state.PAST_RESPONSE = "NO PAST RESPONSES"
46    st.session_state.PREV_QUERY = ""
47    st.session_state.PREV_RESPONSE = ""
48
49# Function to clear the session state variable
50def clear_api_list_updated():
51    st.session_state.api_list_updated = API_LIST.copy()
52    st.session_state.args_in_list_dict = args_in_list_dict.copy()
53    st.session_state.arg_allowed_values_dict = arg_allowed_values_dict.copy()
54    st.session_state.available_arguments = available_arguments.copy()
55    st.session_state.available_tools = available_tools.copy()
56    st.session_state.retrieval_vector_db = FAISS.from_documents(
57        documents=retrieval_data,
58        embedding=retrieval_embeddings,
59        )
60
61# file_path = r'Updated_API_list.json'
62
63# Page navigation
64st.sidebar.title("Navigation")
65page = st.sidebar.radio("Go to", ["Chatbot", "Tool Management", "API_LIST"])
66
67if page == "Chatbot":
68    # Chatbot UI Code
69    st.title("DevRev AI Agent")
70
71    # Display chat messages from history
72    for message in st.session_state.messages:
73        with st.chat_message(message["role"]):
74            st.markdown(message["content"])
75
76    # Accept user input
77    if prompt := st.chat_input("What is up?"):
78        print("okay")
79        st.session_state.messages.append({"role": "user", "content": prompt})
80        with st.chat_message("user"):
81            st.markdown(prompt)
82        print(f"User query: {prompt}")
83        query = prompt
84        with st.chat_message("assistant"):
85            message_placeholder = st.empty()
86            json_answer = pipeline(query, st.session_state.api_list_updated, st.session_state.available_tools, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict, st.session_state.retrieval_vector_db) # allowed args dict ka placeholder modify karna bacha
87            full_response = json_answer
88            try:
89                message_placeholder.json(full_response)
90            except:
91                message_placeholder.markdown(full_response)
92        st.session_state.messages.append({"role": "assistant", "content": full_response})
93
94elif page == "Tool Management":
95    print("TOOLS")
96    # Tool Management Code
97    st.title("API Tool Management")
98
99    # Load or initialize the API list in session state
100    if "api_list_updated" not in st.session_state:
101        st.session_state.api_list_updated = API_LIST
102    
103    # Button to trigger the clearing action
104    if st.button("Reset API list"):
105        clear_api_list_updated()
106
107    # Add a new tool
108    st.header("Add a New Tool")
109    with st.form("new_tool_form"):
110        new_tool_name = st.text_input("Tool Name")
111        new_tool_desc = st.text_area("Tool Description")
112        submitted = st.form_submit_button("Add Tool")
113        if submitted:
114            st.session_state.api_list_updated, st.session_state.available_tools, st.session_state.available_arguments = add_tool(st.session_state.api_list_updated, new_tool_name, new_tool_desc, st.session_state.available_tools, st.session_state.available_arguments, st.session_state.retrieval_vector_db)
115            st.success("Tool Added Successfully!")
116            
117    if st.session_state.api_list_updated:
118        selected_tool_name = st.selectbox("Select a tool", options=[tool["name"] for tool in st.session_state.api_list_updated])
119        selected_tool = next((tool for tool in st.session_state.api_list_updated if tool['name'] == selected_tool_name), None)
120
121    # Update or Delete Tools
122    st.header("Update or Delete Tools")
123    if selected_tool:
124        with st.form("update_delete_tool_form"):
125            st.write(f"Selected Tool: {selected_tool_name}")
126            new_name = st.text_input("New Name", value=selected_tool['name'])
127            new_description = st.text_area("New Description", value=selected_tool['description'])
128            update_button = st.form_submit_button("Update Tool")
129            delete_button = st.form_submit_button("Delete Tool")
130
131            if update_button:
132                st.session_state.api_list_updated, st.session_state.available_tools, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict = update_tool(st.session_state.api_list_updated, selected_tool_name, new_name, new_description, st.session_state.available_tools, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict, st.session_state.retrieval_vector_db)
133                st.success("Tool Updated Successfully!")
134
135            if delete_button:
136                st.session_state.api_list_updated, st.session_state.available_tools, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict, st.session_state.retrieval_vector_db = delete_tool(st.session_state.api_list_updated, selected_tool_name, st.session_state.available_tools, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict, st.session_state.retrieval_vector_db)
137                st.success("Tool Deleted Successfully!")
138
139    # Add a New Argument to a Tool
140    st.header("Add a New Argument to a Tool")
141    # if st.session_state.api_list_updated:  # Check if there are tools available
142    #     selected_tool_name = st.selectbox("Select a tool to add an argument", options=[tool["name"] for tool in st.session_state.api_list_updated])
143    if selected_tool:
144        with st.form("new_argument_form"):
145            new_arg_name = st.text_input("Argument Name")
146            new_arg_desc = st.text_area("Argument Description")
147            new_arg_allowed_values = st.text_area("Argument Allowed Values")
148            new_arg_type = st.text_input("Argument Type")
149            submitted_arg = st.form_submit_button("Add Argument")
150            if submitted_arg:
151                st.session_state.api_list_updated, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict= add_argument(st.session_state.api_list_updated, selected_tool_name, new_arg_name, new_arg_desc, new_arg_type, new_arg_allowed_values, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict, st.session_state.retrieval_vector_db)
152                arg_allowed_values_dict[f"{selected_tool_name}/{new_arg_name}"] = new_arg_allowed_values
153                st.success("Argument Added Successfully!")
154    else:
155        st.write("No tools available. Add a tool first.")
156
157
158    # Update or Delete Arguments
159    st.header("Update or Delete Arguments")
160    if selected_tool:
161        st.write(f"Selected Tool: {selected_tool_name}")
162        selected_arg_name = st.selectbox("Select an argument", options=[arg["argument_name"] for arg in selected_tool["arguments"]])
163        selected_arg = next((arg for arg in selected_tool['arguments'] if arg['argument_name'] == selected_arg_name), None)
164
165        if selected_arg:
166            with st.form("update_delete_arg_form"):
167                new_arg_name = st.text_input("New Argument Name", value=selected_arg['argument_name'])
168                new_arg_desc = st.text_area("New Argument Description", value=selected_arg['argument_description'])
169                if f"{selected_tool_name}/{selected_arg_name}" in arg_allowed_values_dict:    
170                    new_arg_allowed_values = st.text_area("New Argument Allowed Values", value=arg_allowed_values_dict[f"{selected_tool_name}/{selected_arg_name}"])
171                else:
172                    new_arg_allowed_values = st.text_area("New Argument Allowed Values")
173                new_arg_type = st.text_input("New Argument Type", value=selected_arg['argument_type'])
174                update_arg_button = st.form_submit_button("Update Argument")
175                delete_arg_button = st.form_submit_button("Delete Argument")
176
177                if update_arg_button:
178                    st.session_state.api_list_updated, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict = update_argument(st.session_state.api_list_updated, selected_tool_name, selected_arg_name, new_arg_name, new_arg_desc, new_arg_type, new_arg_allowed_values, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict, st.session_state.retrieval_vector_db)
179                    arg_allowed_values_dict[f"{selected_tool_name}/{new_arg_name}"] = new_arg_allowed_values
180                    st.success("Argument Updated Successfully!")
181
182                if delete_arg_button:
183                    st.session_state.api_list_updated, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict = delete_argument(st.session_state.api_list_updated, selected_tool_name, selected_arg_name, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict, st.session_state.retrieval_vector_db)
184                    st.success("Argument Deleted Successfully!")
185
186    # Delete Multiple Arguments
187    if selected_tool:
188        st.header("Delete Multiple Arguments from Tool")
189        all_arg_names = [arg["argument_name"] for arg in selected_tool["arguments"]]
190        selected_args_to_delete = st.multiselect("Select arguments to delete from the tool", options=all_arg_names)
191        if st.button("Delete Selected Arguments"):
192            st.session_state.api_list_updated, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict = delete_multiple_arguments(st.session_state.api_list_updated, selected_tool_name, selected_args_to_delete, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict, st.session_state.retrieval_vector_db)
193            st.success("Selected Arguments Deleted Successfully!")
194    
195    # Display the current API list
196    st.header("Current API List")
197    st.text(json.dumps(st.session_state.api_list_updated, indent=2))
198
199    with open(file_path, 'w') as file:
200        json.dump(st.session_state.api_list_updated, file)
201elif page == "API_LIST":
202    st.session_state.api_list_updated, st.session_state.available_tools, st.session_state.available_arguments, st.session_state.arg_allowed_values_dict, st.session_state.args_in_list_dict
203    st.write(len(st.session_state.retrieval_vector_db.docstore._dict))
204