# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import sys
from typing import List

from langchain.chains import RetrievalQA
from langchain.chains.conversational_retrieval.base import (
    BaseConversationalRetrievalChain,
)
from langchain.llms.vertexai import VertexAI
from langchain.memory import ConversationBufferMemory
from langchain.tools import BaseTool

current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)

from MyVertexAIEmbedding import MyVertexAIEmbedding  # noqa: E402
from VertexMatchingEngine import MatchingEngine, MatchingEngineUtils  # noqa: E402

# https://cdn.cloudflare.steamstatic.com/steam/apps/597180/manuals/Old_World-Official_User_Manual.pdf?t=1653279974
"""
Matching Engine As Retriever
"""

ME_REGION = os.getenv("GOOGLE_CLOUD_REGIN")
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
ME_INDEX_NAME = f"{PROJECT_ID}-chatbot-vme"
ME_DIMENSIONS = 768
ME_EMBEDDING_DIR = f"gs://{PROJECT_ID}-chatbot-embeddings"
REQUESTS_PER_MINUTE = 15

mengine = MatchingEngineUtils(
    project_id=PROJECT_ID, region=ME_REGION, index_name=ME_INDEX_NAME
)
embedding = MyVertexAIEmbedding()

llm = VertexAI()
memory = ConversationBufferMemory()


def create_PDFQA_chain_me_RetrievalQA() -> BaseConversationalRetrievalChain:
    mengine = MatchingEngineUtils(
        project_id=PROJECT_ID, region=ME_REGION, index_name=ME_INDEX_NAME
    )
    ME_INDEX_ID, ME_INDEX_ENDPOINT_ID = mengine.get_index_and_endpoint()

    me = MatchingEngine.from_components(
        project_id=PROJECT_ID,
        region=ME_REGION,
        gcs_bucket_name=f'gs://{ME_EMBEDDING_DIR.split("/")[2]}',
        embedding=embedding,
        index_id=ME_INDEX_ID,
        endpoint_id=ME_INDEX_ENDPOINT_ID,
    )

    retriever = me.as_retriever()
    doc_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=False,
        verbose=True,
    )

    return doc_chain


class VIAI_INFO_ME(BaseTool):
    name = "VIAI_INFO_ME"
    description = """
    Use this tool to get information regarding the solution "Visual Inspection AI Edge", or "VIAI Edge".
    The Tool Input is the user's question, the user may reference to previous convsation,
    add context to the question when needed.
    The Output is the result
    """

    def _run(self, query: str) -> str:
        if query == "":
            query = "summarize"
        chat_history: List[str] = []
        print("Running tool:{}".format(query))

        qa = create_PDFQA_chain_me_RetrievalQA()
        result = qa(
            {"query": query, "chat_history": chat_history}, return_only_outputs=False
        )
        return result

    async def _arun(self, query: str) -> str:
        """Use the tool asynchronously."""
        print(f"*** Invoking MockTool with query '{query}'")
        return f"Answer of '{query}' is 'Michael Chi'"

langchain.llms.vertexai.VertexAI langchain.memory.ConversationBufferMemory langchain.chains.RetrievalQA.from_chain_type
[((957, 985), 'sys.path.append', 'sys.path.append', (['current_dir'], {}), '(current_dir)\n', (972, 985), False, 'import sys\n'), ((1302, 1333), 'os.getenv', 'os.getenv', (['"""GOOGLE_CLOUD_REGIN"""'], {}), "('GOOGLE_CLOUD_REGIN')\n", (1311, 1333), False, 'import os\n'), ((1347, 1380), 'os.getenv', 'os.getenv', (['"""GOOGLE_CLOUD_PROJECT"""'], {}), "('GOOGLE_CLOUD_PROJECT')\n", (1356, 1380), False, 'import os\n'), ((1540, 1631), 'VertexMatchingEngine.MatchingEngineUtils', 'MatchingEngineUtils', ([], {'project_id': 'PROJECT_ID', 'region': 'ME_REGION', 'index_name': 'ME_INDEX_NAME'}), '(project_id=PROJECT_ID, region=ME_REGION, index_name=\n    ME_INDEX_NAME)\n', (1559, 1631), False, 'from VertexMatchingEngine import MatchingEngine, MatchingEngineUtils\n'), ((1645, 1666), 'MyVertexAIEmbedding.MyVertexAIEmbedding', 'MyVertexAIEmbedding', ([], {}), '()\n', (1664, 1666), False, 'from MyVertexAIEmbedding import MyVertexAIEmbedding\n'), ((1674, 1684), 'langchain.llms.vertexai.VertexAI', 'VertexAI', ([], {}), '()\n', (1682, 1684), False, 'from langchain.llms.vertexai import VertexAI\n'), ((1694, 1720), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {}), '()\n', (1718, 1720), False, 'from langchain.memory import ConversationBufferMemory\n'), ((930, 955), 'os.path.abspath', 'os.path.abspath', (['__file__'], {}), '(__file__)\n', (945, 955), False, 'import os\n'), ((1814, 1905), 'VertexMatchingEngine.MatchingEngineUtils', 'MatchingEngineUtils', ([], {'project_id': 'PROJECT_ID', 'region': 'ME_REGION', 'index_name': 'ME_INDEX_NAME'}), '(project_id=PROJECT_ID, region=ME_REGION, index_name=\n    ME_INDEX_NAME)\n', (1833, 1905), False, 'from VertexMatchingEngine import MatchingEngine, MatchingEngineUtils\n'), ((2311, 2438), 'langchain.chains.RetrievalQA.from_chain_type', 'RetrievalQA.from_chain_type', ([], {'llm': 'llm', 'chain_type': '"""stuff"""', 'retriever': 'retriever', 'return_source_documents': '(False)', 'verbose': '(True)'}), "(llm=llm, chain_type='stuff', retriever=\n    retriever, return_source_documents=False, verbose=True)\n", (2338, 2438), False, 'from langchain.chains import RetrievalQA\n')]