import torch
from transformers import AutoTokenizer, AutoModel
from transformers import DistilBertTokenizer, DistilBertModel

#model_name = "distilbert-base-uncased"
#model_name = "sentence-transformers/msmarco-distilbert-base-dot-prod-v3"
model_name = "sentence-transformers/distilbert-base-nli-stsb-mean-tokens"


#Mean Pooling - Take attention mask into account for correct averaging
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
    return sum_embeddings / sum_mask


def sentence_to_vector(raw_inputs):
    tokenizer = DistilBertTokenizer.from_pretrained(model_name)
    model = DistilBertModel.from_pretrained(model_name)
    inputs_tokens = tokenizer(raw_inputs, padding=True, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs_tokens)

    sentence_embeddings = mean_pooling(outputs, inputs_tokens['attention_mask'])
    return sentence_embeddings

knn_index = {
    "settings": {
        "index.knn": True,
        "index.knn.space_type": "cosinesimil",
        "analysis": {
          "analyzer": {
            "default": {
              "type": "standard",
              "stopwords": "_english_"
            }
          }
        }
    },
    "mappings": {
        "properties": {
            "question_vector": {
                "type": "knn_vector",
                "dimension": 768,
                "store": True
            },
            "question": {
                "type": "text",
                "store": True
            },
            "answer": {
                "type": "text",
                "store": True
            }
        }
    }
}

query_raw_sentences = ['does this work with xbox?']
search_vector = sentence_to_vector(query_raw_sentences)[0].tolist()

query={
    "size": 30,
    "query": {
        "knn": {
            "question_vector":{
                "vector":search_vector,
                "k":30
            }
        }
    }
}

res = aos_client.search(index="nlp_pqa", 
                       body=query,
                       stored_fields=["question","answer"])
#print("Got %d Hits:" % res['hits']['total']['value'])
query_result=[]
for hit in res['hits']['hits']:
    row=[hit['_id'],hit['_score'],hit['fields']['question'][0],hit['fields']['answer'][0]]
    query_result.append(row)

query_result_df = pd.DataFrame(data=query_result,columns=["_id","_score","question","answer"])
display(query_result_df)

from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth
import boto3
import json

kms = boto3.client('secretsmanager')
aos_credentials = json.loads(kms.get_secret_value(SecretId=outputs['OpenSearchSecret'])['SecretString'])

region = 'us-east-1' 

#credentials = boto3.Session().get_credentials()
#auth = AWSV4SignerAuth(credentials, region)
auth = (aos_credentials['username'], aos_credentials['password'])

index_name = 'nlp_pqa'

aos_client = OpenSearch(
    hosts = [{'host': aos_host, 'port': 443}],
    http_auth = auth,
    use_ssl = True,
    verify_certs = True,
    connection_class = RequestsHttpConnection
)

from opensearch_py_ml.ml_models import SentenceTransformerModel
from opensearch_py_ml.ml_commons import MLCommonClient

ml_client = MLCommonClient(aos_client)
model_id = ml_client.register_pretrained_model(model_name = "huggingface/sentence-transformers/all-MiniLM-L12-v2", model_version = "1.0.1", model_format = "TORCH_SCRIPT", deploy_model=False, wait_until_deployed=False)
print(model_id)

load_model_output = ml_client.deploy_model(model_id)

print(load_model_output)

pipeline={
  "description": "An example neural search pipeline",
  "processors" : [
    {
      "text_embedding": {
        "model_id": model_id,
        "field_map": {
           "question": "question_vector"
        }
      }
    }
  ]
}
pipeline_id = 'nlp_pipeline'
aos_client.ingest.put_pipeline(id=pipeline_id,body=pipeline)

knn_index = {
    "settings": {
        "index.knn": True,
        "index.knn.space_type": "cosinesimil",
        "default_pipeline": pipeline_id,
        "analysis": {
          "analyzer": {
            "default": {
              "type": "standard",
              "stopwords": "_english_"
            }
          }
        }
    },
    "mappings": {
        "properties": {
            "question_vector": {
                "type": "knn_vector",
                "dimension": 384,
                "method": {
                    "name": "hnsw",
                    "space_type": "l2",
                    "engine": "faiss"
                },
                "store": True
            },
            "question": {
                "type": "text",
                "store": True
            },
            "answer": {
                "type": "text",
                "store": True
            }
        }
    }
}

query={
  "_source": {
        "exclude": [ "question_vector" ]
    },
  "size": 30,
  "query": {
    "neural": {
      "question_vector": {
        "query_text": "does this work with xbox?",
        "model_id": model_id,
        "k": 30
      }
    }
  }
}

query={
    "size": 30,
    "query": {
        "match": {
            "question":"does this work with xbox?"
        }
    }
}

from langchain.document_loaders import UnstructuredURLLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(chunk_size = 1000, chunk_overlap = 100)

urls = ["https://docs.aws.amazon.com/opensearch-service/latest/developerguide/bp.html",
        "https://docs.aws.amazon.com/opensearch-service/latest/developerguide/sizing-domains.html", 
        "https://docs.aws.amazon.com/opensearch-service/latest/developerguide/petabyte-scale.html",
        "https://docs.aws.amazon.com/opensearch-service/latest/developerguide/managedomains-dedicatedmasternodes.html",
        "https://docs.aws.amazon.com/opensearch-service/latest/developerguide/cloudwatch-alarms.html"]
url_loader = UnstructuredURLLoader(urls=urls)
url_texts = url_loader.load_and_split(text_splitter=text_splitter)
all_splits = url_texts
print(all_splits[0])

page_content='AWS\n\nDocumentation\n\nAmazon OpenSearch Service\n\nDeveloper Guide\n\nMonitoring and alerting\n\nShard strategy\n\nStability\n\nPerformance\n\nSecurity\n\nCost optimization\n\nOperational best practices for Amazon OpenSearch Service\n\nThis chapter provides best practices for operating Amazon OpenSearch Service domains and includes general\n        guidelines that apply to many use cases. Each workload is unique, with unique\n        characteristics, so no generic recommendation is exactly right for every use case. The most\n        important best practice is to deploy, test, and tune your domains in a continuous cycle to\n        find the optimal configuration, stability, and cost for your workload.\n\nTopics\n\nMonitoring and alerting\n\nShard strategy\n\nStability\n\nPerformance\n\nSecurity\n\nCost optimization\n\nSizing Amazon OpenSearch Service domains\n\nPetabyte scale in Amazon OpenSearch Service\n\nDedicated master nodes in\n                Amazon OpenSearch Service' metadata={'source': 'https://docs.aws.amazon.com/opensearch-service/latest/developerguide/bp.html'}

from typing import Any, Dict, Iterable, List, Optional, Tuple, Callable
import json
from langchain.embeddings import SagemakerEndpointEmbeddings
from langchain.embeddings.sagemaker_endpoint import EmbeddingsContentHandler
from langchain.schema import Document

class BulkSagemakerEndpointEmbeddings(SagemakerEndpointEmbeddings):
        def embed_documents(
            self, texts: List[str], chunk_size: int = 5
        ) -> List[List[float]]:
            results = []
            _chunk_size = len(texts) if chunk_size > len(texts) else chunk_size

            for i in range(0, len(texts), _chunk_size):
                response = self._embedding_func(texts[i:i + _chunk_size])
                results.extend(response)
            return results
        
class EmbeddingContentHandler(EmbeddingsContentHandler):
        content_type = "application/json"
        accepts = "application/json"

        def transform_input(self, prompt: str, model_kwargs={}) -> bytes:

            input_str = json.dumps({"text_inputs": prompt, **model_kwargs})
            return input_str.encode('utf-8') 

        def transform_output(self, output: bytes) -> str:

            response_json = json.loads(output.read().decode("utf-8"))
            embeddings = response_json["embedding"]
            if len(embeddings) == 1:
                return [embeddings[0]]
            return embeddings

print(embedding_endpoint_name)
embeddings = BulkSagemakerEndpointEmbeddings( 
            endpoint_name=embedding_endpoint_name,
            region_name=aws_region, 
            content_handler=EmbeddingContentHandler())

from langchain.vectorstores import OpenSearchVectorSearch

os_domain_ep = 'https://'+aos_host

embedding_index_name = 'opensearch_kb_vector'

if len(all_splits) > 500:
    for i in range(0, len(all_splits), 500):
        start = i
        end = i+500
        if end > len(all_splits):
            end = len(all_splits)-1
        docs = all_splits[start:end]
        OpenSearchVectorSearch.from_documents(
            index_name = embedding_index_name,
            documents=docs,
            embedding=embeddings,
            opensearch_url=os_domain_ep,
            http_auth=auth
        )
        print(f"ingest documents from {start} to {end}", start, end)
else:
    OpenSearchVectorSearch.from_documents(
            index_name = embedding_index_name,
            documents=all_splits,
            embedding=embeddings,
            opensearch_url=os_domain_ep,
            http_auth=auth
        )
    print(f"ingest documents")

retriever = open_search_vector_store.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={
        'k': 5,
        'score_threshold': 0.62
    }
)

qa_with_source = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    return_source_documents=True,
    chain_type_kwargs={"prompt": prompt_template2}
)

print("Question is:" + question)
result = qa_with_source({"query": question})

print("result:" + result["result"])
print("\n\n===========================")
print("\nsource documents:")
for doc in result["source_documents"]:
    print(doc)
    print("---------------------------\n")

source documents:
page_content="If your minimum storage requirement exceeds 1 PB, see Petabyte scale in Amazon OpenSearch Service.\n\nIf you have rolling indexes and want to use a hot-warm architecture, see\n                            UltraWarm storage for Amazon OpenSearch Service.\n\nChoosing the number of shards\n\nAfter you understand your storage requirements, you can investigate your indexing\n                strategy. By default in OpenSearch Service, each index is divided into five primary shards and one\n                replica (total of 10 shards). This behavior differs from open source OpenSearch,\n                which defaults to one primary and one replica shard. Because you can't easily change\n                the number of primary shards for an existing index, you should decide about shard\n                count before indexing your first document." metadata={'source': 'https://docs.aws.amazon.com/opensearch-service/latest/developerguide/sizing-domains.html'}
---------------------------

from uuid import uuid4
from typing import Dict
from langchain.memory import ConversationBufferMemory
from langchain.memory import DynamoDBChatMessageHistory
from langchain.memory import ConversationBufferWindowMemory
from langchain import PromptTemplate, SagemakerEndpoint
from langchain.llms.sagemaker_endpoint import LLMContentHandler
from langchain.chains import RetrievalQA

（中略）

ddb_table_name = "conversation-history-memory"
session_id = str(uuid4())
chat_memory = DynamoDBChatMessageHistory(
        table_name=ddb_table_name,
        session_id=session_id
    )

messages = chat_memory.messages

# Maintains immutable sessions
# If previous session was present, create
# a new session and copy messages, and 
# generate a new session_id 
if messages:
    session_id = str(uuid4())
    chat_memory = DynamoDBChatMessageHistory(
        table_name=ddb_table_name,
        session_id=session_id
    )
    # This is a workaround at the moment. Ideally, this should
    # be added to the DynamoDBChatMessageHistory class
    try:
        messages = messages_to_dict(messages)
        chat_memory.table.put_item(
            Item={"SessionId": session_id, "History": messages}
        )
    except Exception as e:
        print(e)

memory = ConversationBufferMemory(chat_memory=chat_memory, memory_key="chat_history", return_messages=True)

Customize pipelines for retrieval augmented generation
Enhancing the conversational search functionality introduced in OpenSearch 2.10 as an experimental toolkit, this release introduces several new parameters that can be used to customize retrieval augmented generation (RAG) pipelines. These optional parameters provide core logic that allows you to adapt the way OpenSearch interacts with large language models (LLMs) as part of generative artificial intelligence (GenAI) applications. See the documentation for this feature to explore the available parameters.