Update app.py

app.py

@@ -1,44 +1,123 @@
+from llama_index.core.agent import ReActAgent
+from llama_index.llms.openai import OpenAI
+from llama_index.core.tools import FunctionTool
+from opensearchpy import OpenSearch
+from gradio_client import Client
 import streamlit as st
-import streamlit as st
-from openai import OpenAI
 import json
+import openai
+import warnings
+warnings.filterwarnings('ignore')
+from UI_Config import *
 import os
 
+openai_api_key = os.getenv("OPENAI_API_KEY")
+job_id = os.getenv('JOB_ID')
+user = os.getenv('USERNAME')
+password = os.getenv('PASSWORD')
+host = os.getenv('HOST')
+port = int(os.getenv('PORT'))
+
+auth = (user,password)
+
+client = openai.OpenAI(api_key=openai_api_key)
+
+os_client = OpenSearch(
+    hosts = [{'host': host, 'port': port}],
+    http_auth = auth,
+    use_ssl = True,
+    verify_certs = False
+)
+indices = os_client.cat.indices(format="json")
+list_of_indeces = []
+for index in indices:
+    list_of_indeces.append(index['index'])
+
+def rag_app(user_input: str) -> str:
+    gr_client = Client("anasmkh/QdrantVectorStore_Llamaindex")
+    result = gr_client.predict(
+    		user_input=user_input,
+    		api_name="/chat_with_ai"
+    )
+    return result
 
-openai_api_key = os.getenv('OPENAI_API_KEY')
+rag_tool = FunctionTool.from_defaults(fn=rag_app)
 
-client = OpenAI(api_key=openai_api_key)
+def query_generator(user_input:str) -> str:
+    job = job_id
+    response = client.fine_tuning.jobs.retrieve(job)
+    completion = client.chat.completions.create(
+    model=response.fine_tuned_model,
+    messages=[
+    {"role": "system", "content": f"""You are a highly skilled assistant trained to translate natural language requests into accurate and efficient OpenSearch JSON queries. Follow a clear, step-by-step process to:
 
-  
+Understand the user's request by breaking it down into components such as filters, aggregations, sort criteria, and specific fields.
+Pay special attention to fields with unique names, such as Date (instead of timestamp) and Stream (instead of type), and ensure they are used correctly in the query.
+Recognize that the user operates within two main opcos: Zambia and Eswatini, each containing ptm_counters, ptm_events, and multiple streams like ers-daily.
+Generate a valid JSON query strictly based on the provided indices, ensuring it aligns with the user's prompt. The available indices are: {list_of_indeces}.
+When generating the query:
+Be precise and include only necessary fields and components relevant to the request.
+Assume any unspecified context or detail needs clarification and provide a clear explanation of your assumptions if needed.
+Optimize the query for OpenSearch performance and readability.
+Your goal is to provide a query that directly addresses the user's needs while being efficient and valid within the OpenSearch framework."""
 
-job_id = os.getenv('job_id')
-response = client.fine_tuning.jobs.retrieve(job_id)
+},
+    {"role": "user", "content": user_input}
+  ])
+    return completion.choices[0].message
 
 
+query_tool = FunctionTool.from_defaults(fn=query_generator)
 
-st.title("AI Assistant - Fine-Tuned Model")
+llm = OpenAI(model="gpt-3.5-turbo", temperature=0)
+agent = ReActAgent.from_tools([query_tool,rag_tool], llm=llm, verbose=True)
 
+def implement_query(generated_query):
+    try:
+        st.write("Raw Query:", generated_query)
 
+        if isinstance(generated_query, str):
+            generated_query = generated_query.replace("'", '"')
+            query = json.loads(generated_query)
+        else:
+            query = generated_query
 
-user_input = st.text_input("Enter your question:")
+        st.write("Validated Query:", query)
 
 
-if st.button("Ask"):
+        response = os_client.search(body=query)
+        return response
+    except json.JSONDecodeError as e:
+        st.error("Error: The generated query is not valid JSON.")
+        st.write(f"JSONDecodeError Details: {e}")
+    except Exception as e:
+        st.error(f"Error executing OpenSearch query: {e}")
+        st.write(f"Exception Details: {e}")
+
+
+st.subheader('OpenSearch Assistant')
+user_input = st.text_input("Enter your query:", "")
+
+
+if st.button("Submit"):
     if user_input:
-     
-        try: 
-                    completion = client.chat.completions.create(
-                          model="ft:gpt-3.5-turbo-0125:personal::AVdWfIMJ",
-                          messages=[
-                            {"role": "system", "content": "You are an assistant trained to translate natural language requests into OpenSearch queries. Use a step, generate a valid OpenSearch JSON query only."},
-                            {"role": "user", "content": user_input}
-                                  ]
-                    )
-     
-                    st.write("Assistant's response:")
-                    st.write(completion.choices[0].message.content)
-
-        except Exception as e:
-            st.write("Error:", e)
+        with st.spinner("Processing..."):
+            try:
+                response = agent.chat(user_input)
+                st.success("Query Processed Successfully!")
+                st.subheader("Agent Response:")
+                sources = response.sources
+                for source in sources:
+                    st.write('Used Tool: ',source.tool_name)
+                    if source.tool_name =='query_generator':
+                        st.write(source.raw_output.content)
+                        os_response = implement_query(source.raw_output.content)
+                        st.subheader('OS Response')
+                        st.write(os_response)
+                    else:
+                        st.write(source.raw_output[0][0][1])
+            except Exception as e:
+                st.error(f"Error: {e}")
     else:
-        st.write("Please enter a question.")
+        st.warning("Please enter a query to process.")
+