semantic_search/query_rewrite.py

import json
import os
import sys
import boto3
import amazon_rekognition
from botocore.config import Config
import getpass
#from nltk.stem import PorterStemmer
#from nltk.tokenize import word_tokenize
import os
import streamlit as st
from langchain.schema import Document
#from langchain_community.vectorstores import OpenSearchVectorSearch,ElasticsearchStore
from requests_aws4auth import AWS4Auth
from requests.auth import HTTPBasicAuth
from langchain.chains.query_constructor.base import (
    StructuredQueryOutputParser,
    get_query_constructor_prompt,
)
from langchain.retrievers.self_query.opensearch import OpenSearchTranslator
from langchain.chains import ConversationChain
from langchain.llms.bedrock import Bedrock
from langchain.memory import ConversationBufferMemory
from langchain.chains.query_constructor.base import AttributeInfo
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain_core.prompts.few_shot import FewShotPromptTemplate
from langchain_core.prompts.prompt import PromptTemplate
from langchain.embeddings import BedrockEmbeddings
#from langchain.vectorstores import OpenSearchVectorSearch
from opensearchpy import OpenSearch, RequestsHttpConnection
import utilities.invoke_models as invoke_models



bedrock_params = {
    "max_tokens_to_sample":2048,
    "temperature":0.0001,
    "top_k":250,
    "top_p":1,
    "stop_sequences":["\\n\\nHuman:"]
}
bedrock_region="us-east-1"

#boto3_bedrock = boto3.client(service_name="bedrock-runtime", endpoint_url=f"https://bedrock-runtime.{bedrock_region}.amazonaws.com")
boto3_bedrock = boto3.client(service_name="bedrock-runtime", config=Config(region_name=bedrock_region))

bedrock_titan_llm = Bedrock(model_id="anthropic.claude-instant-v1", client=boto3_bedrock)
bedrock_titan_llm.model_kwargs = bedrock_params
bedrock_embeddings = BedrockEmbeddings(model_id='amazon.titan-embed-text-v1',client=boto3_bedrock)


schema = """{{
    "content": "Brief summary of a retail product",
    "attributes": {{
    "category": {{
        "description": "The category of the product, the available categories are apparel, footwear, outdoors, electronics, beauty, jewelry, accessories, housewares, homedecor, furniture, seasonal, floral, books, groceries, instruments, tools, hot dispensed, cold dispensed, food service and salty snacks",
        "type": "string"
    }},
    "gender_affinity": {{
        "description": "The gender that the product relates to, the choices are Male and Female", 
        "type": "string"
    }},
    "price": {{
        "description": "Cost of the product",
        "type": "double"
    }},
    "description": {{
        "description": "The detailed description of the product",
        "type": "string"
    }},
    "color": {{
        "description": "The color of the product",
        "type": "string"
    }},
     "caption": {{
        "description": "The short description of the product",
        "type": "string"
    }},
     "current_stock": {{
        "description": "The available quantity of the product in stock for sale",
        "type": "integer"
    }},
     "style": {{
        "description": "The style of the product",
        "type": "string"
    }}
}}
}}"""
metadata_field_info_ = [
    AttributeInfo(
        name="price",
        description="Cost of the product",
        type="string",
    ),
    AttributeInfo(
        name="style",
        description="The style of the product",
        type="string",
    ),
    AttributeInfo(
        name="category",
        description="The category of the product, the available categories are apparel, footwear, outdoors, electronics, beauty, jewelry, accessories, housewares, homedecor, furniture, seasonal, floral, books, groceries, instruments, tools, hot dispensed, cold dispensed, food service and salty snacks",
        type="string",
    ),
    AttributeInfo(
        name="current_stock",
        description="The available quantity of the product",
        type="string",
    ),
    AttributeInfo(
        name="gender_affinity", 
        description="The gender that the product relates to, the choices are Male and Female", 
        type="string"
    ),
    AttributeInfo(
        name="caption", 
        description="The short description of the product", 
        type="string"
    ),
    AttributeInfo(
        name="description", 
        description="The detailed description of the product", 
        type="string"
    ),
    AttributeInfo(
        name="color", 
        description="The color of the product", 
        type="string"
    )
]
document_content_description_ = "Brief summary of a retail product"

# open_search_vector_store = OpenSearchVectorSearch(
#                                     index_name="retail-ml-search-index",#"self-query-rewrite-retail",
#                                     embedding_function=bedrock_embeddings,
#                                     opensearch_url=os_domain_ep,
#                                     http_auth=auth
#                                     )  

examples = [
    {  "i":1,
        "data_source": schema,
        "user_query": "black shoes for men",
        "structured_request": """{{
    "query": "shoes",
    "filter": "and(eq(\"color\", \"black\"),  eq(\"category\", \"footwear\")), eq(\"gender_affinity\", \"Male\")"
            }}""",
    },
    
        {  "i":2,
        "data_source": schema,
        "user_query": "black or brown jackets for men under 50 dollars",
        "structured_request": """{{
    "query": "jackets",
    "filter": "and(eq(\"style\", \"jacket\"), or(eq(\"color\", \"brown\"),eq(\"color\", \"black\")),eq(\"category\", \"apparel\"),eq(\"gender_affinity\", \"male\"),lt(\"price\", \"50\"))"
            }}""",
    },
          {  "i":2,
        "data_source": schema,
        "user_query": "trendy handbags for women",
        "structured_request": """{{
    "query": "handbag",
    "filter": "and(eq(\"style\", \"bag\") ,eq(\"category\", \"accessories\"),eq(\"gender_affinity\", \"female\"))"
            }}""",
    }
]


example_prompt = PromptTemplate(
    input_variables=["question", "answer"], template="Question: {question}\n{answer}"
)
example_prompt=PromptTemplate(input_variables=['data_source', 'i', 'structured_request', 'user_query'],
template='<< Example {i}. >>\nData Source:\n{data_source}\n\nUser Query:\n{user_query}\n\nStructured Request:\n{structured_request}\n')
#print(example_prompt.format(**examples[0]))

prefix_ = """
Your goal is to structure the user's query to match the request schema provided below.

<< Structured Request Schema >>
When responding use a markdown code snippet with a JSON object formatted in the following schema:

```json
{{
    "query": string \ text string to compare to document contents
    "filter": string \ logical condition statement for filtering documents
}}
```

The query string should contain only text that is expected to match the contents of documents. Any conditions in the filter should not be mentioned in the query as well.

A logical condition statement is composed of one or more comparison and logical operation statements.

A comparison statement takes the form: `comp(attr, val)`:
- `comp` (eq | ne | gt | gte | lt | lte | contain | like | in | nin): comparator
- `attr` (string):  name of attribute to apply the comparison to
- `val` (string): is the comparison value

A logical operation statement takes the form `op(statement1, statement2, ...)`:
- `op` (and | or | not): logical operator
- `statement1`, `statement2`, ... (comparison statements or logical operation statements): one or more statements to apply the operation to

Make sure that you only use the comparators and logical operators listed above and no others.
Make sure that filters only refer to attributes that exist in the data source.
Make sure that filters only use the attributed names with its function names if there are functions applied on them.
Make sure that filters only use format `YYYY-MM-DD` when handling date data typed values.
Make sure that filters take into account the descriptions of attributes and only make comparisons that are feasible given the type of data being stored.
Make sure that filters are only used as needed. If there are no filters that should be applied return "NO_FILTER" for the filter value.
"""

suffix_ = """<< Example 3. >>
Data Source:
{schema}

User Query:
{query}

Structured Request:
"""

prompt_ = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    suffix=suffix_,
    prefix=prefix_,
    input_variables=["query","schema"],
)





# retriever = SelfQueryRetriever.from_llm(
#     bedrock_titan_llm, open_search_vector_store, document_content_description_, metadata_field_info_, verbose=True
# )

# res = retriever.get_relevant_documents("bagpack for men")

# st.write(res)

######### use this for self query retriever ########
# prompt = get_query_constructor_prompt(
#     document_content_description_,
#     metadata_field_info_,
# )
# output_parser = StructuredQueryOutputParser.from_components()
# query_constructor = prompt | bedrock_titan_llm | output_parser

def get_new_query_res(query):
    field_map = {'Price':'price','Gender':'gender_affinity','Category':'category','Style':'style','Color':'color'}
    field_map_filter = {key: field_map[key] for key in st.session_state.input_must}
    if(query == ""):
        query = st.session_state.input_rekog_label
    if(st.session_state.input_is_rewrite_query == 'enabled'):

        res = invoke_models.invoke_llm_model( prompt_.format(query=query,schema = schema)  ,False)
        inter_query = res[7:-3].replace('\\"',"'").replace("\n","")
        print("inter_query")
        print(inter_query)
        query_struct = StructuredQueryOutputParser.from_components().parse(inter_query)  
        print("query_struct")
        print(query_struct)
        opts = OpenSearchTranslator()
        result_query_llm = opts.visit_structured_query(query_struct)[1]['filter']
        print(result_query_llm)
        draft_new_query = {'bool':{'should':[],'must':[]}}
        if('bool' in result_query_llm and ('must' in result_query_llm['bool'] or 'should' in result_query_llm['bool'])):
            #draft_new_query['bool']['should'] = []
            if('must' in result_query_llm['bool']):
                for q in result_query_llm['bool']['must']:
                    old_clause = list(q.keys())[0]
                    if(old_clause == 'term'):
                        new_clause = 'match'
                    else:
                        new_clause = old_clause
                    q_dash = {}
                    q_dash[new_clause] = {}
                    long_field = list(q[old_clause].keys())[0]
                    #print(long_field)
                    get_attr = long_field.split(".")[1]
                    #print(get_attr)
                    q_dash[new_clause][get_attr] = q[old_clause][long_field]
                    #print(q_dash)
                    if(get_attr in list(field_map_filter.values())):
                        draft_new_query['bool']['must'].append(q_dash)
                    else:
                        draft_new_query['bool']['should'].append(q_dash)

        query_ = draft_new_query
        
        
        ###### find the main subject of the query
        #imp_item = ""
        # if("bool" in query_ and  'should' in query_['bool']):
        #     for i in query_['bool']['should']:
        #         if("term" in i.keys()):
        #             if("metadata.category.keyword" in i["term"]):
        #                 imp_item = imp_item + i["term"]["metadata.category.keyword"]+ " "
        #             if("metadata.style.keyword" in i["term"]):
        #                 imp_item = imp_item + i["term"]["metadata.style.keyword"]+ " "
        #         if("match" in i.keys()):
        #             if("metadata.category.keyword" in i["match"]):
        #                 imp_item = imp_item + i["match"]["metadata.category.keyword"]+ " "
        #             if("metadata.style.keyword" in i["match"]):
        #                 imp_item = imp_item + i["match"]["metadata.style.keyword"]+ " "
        # else:
        #     if("term" in query_):
        #             if("metadata.category.keyword" in query_):
        #                 imp_item = imp_item + query_["metadata.category.keyword"] + " "
        #             if("metadata.style.keyword" in query_):
        #                 imp_item = imp_item + query_["metadata.style.keyword"]+ " "
        #     if("match" in query_):
        #             if("metadata.category.keyword" in query_):
        #                 imp_item = imp_item + query_["metadata.category.keyword"]+ " "
        #             if("metadata.style.keyword" in query_):
        #                 imp_item = imp_item + query_["metadata.style.keyword"]+ " "
        ###### find the main subject of the query
        imp_item = (opts.visit_structured_query(query_struct)[0]).replace(",","")
            
                    
        if(imp_item == ""):
            imp_item = query
            
        #ps = PorterStemmer()        
        # def stem_(sentence):
        #     words = word_tokenize(sentence)
            
        #     words_stem = ""

        #     for w in words:
        #         words_stem = words_stem +" "+ps.stem(w)
        #     return words_stem.strip()
        
        #imp_item = stem_(imp_item)
        print("imp_item---------------")
        print(imp_item)
        print(query_)
        if('must' in query_['bool']):
            query_['bool']['must'].append({
                    "simple_query_string": {
                    
                        "query": imp_item.strip(),
                      "fields":['description',"style","caption"]#'rekog_all^3'
                    
                    }
                    #"match":{"description":imp_item.strip()}
                })
        else:
            query_['bool']['must']={
                    "multi_match": {
                    
                        "query": imp_item.strip(),
                      "fields":['description',"style"]#'rekog_all^3'
                    
                    }
                    #"match":{"description":imp_item.strip()}
                }
       
            
        #query_['bool']["minimum_should_match"] = 1
            
        st.session_state.input_rewritten_query = {"query":query_}
        print(st.session_state.input_rewritten_query)