app.py

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 json
import openai
import warnings
warnings.filterwarnings('ignore')
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

rag_tool = FunctionTool.from_defaults(fn=rag_app)

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": """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.
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."""

},
    {"role": "user", "content": user_input}
  ])
    return completion.choices[0].message


query_tool = FunctionTool.from_defaults(fn=query_generator)

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

        st.write("Validated Query:", query)


    #     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:
        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.warning("Please enter a query to process.")