import sys

import pandas as pd
from llama_index import Document, set_global_service_context, StorageContext, load_index_from_storage, VectorStoreIndex
from llama_index.indices.base import BaseIndex
from llama_index.storage.docstore import SimpleDocumentStore
from llama_index.storage.index_store import SimpleIndexStore
from llama_index.vector_stores import SimpleVectorStore

from config import (
    API_KEY,
    DEPLOYMENT_NAME,
    MODEL_NAME,
    API_BASE,
    API_VERSION,
    EMBEDDING_MODEL,
    EMBEDDING_DEPLOYMENT,
)

class LlamaQueryEngine:

    def __init__(
            self,
            api_key=API_KEY,
            deployment_name=DEPLOYMENT_NAME,
            model_name=MODEL_NAME,
            api_base=API_BASE,
            api_version=API_VERSION,
            embedding_model=EMBEDDING_MODEL,
            embedding_deployment=EMBEDDING_DEPLOYMENT,
    ):
        import openai
        import logging
        import os

        from langchain.embeddings import OpenAIEmbeddings
        from llama_index.llms import AzureOpenAI
        from llama_index import LangchainEmbedding
        from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext

        logging.basicConfig(stream=sys.stdout, level=logging.INFO)
        openai.api_type = "azure"
        openai.api_base = api_base
        openai.api_version = api_version
        os.environ["OPENAI_API_KEY"] = api_key
        openai.api_key = os.getenv("OPENAI_API_KEY")

        llm = AzureOpenAI(
            deployment_name=deployment_name,
            model=model_name,
            temperature=0,
            engine="gpt35",
            max_tokens=2048
        )

        embedding_llm = LangchainEmbedding(
            OpenAIEmbeddings(
                model=embedding_model,
                deployment=embedding_deployment,
                openai_api_key=openai.api_key,
                openai_api_base=openai.api_base,
                openai_api_type=openai.api_type,
                openai_api_version=openai.api_version,
            ),
            embed_batch_size=1,
        )

        service_context = ServiceContext.from_defaults(
            llm=llm,
            embed_model=embedding_llm,
        )

        set_global_service_context(service_context)
        # index = VectorStoreIndex.from_documents(documents)
        # self.index = index
        # self.query_engine = index.as_query_engine()
        self.index = None
        self.query_engine = None

    def load_doc_from_csv(self, csv_path, text_column="decoded_readme", max_docs=20, is_persist=False, has_persist=False, persist_dir="app/data/persist"):
        if has_persist:
            self.retrieve_index(persist_dir)
            return
        df = pd.read_csv(csv_path)
        text_list = df[text_column].tolist()
        text_list = text_list[:max_docs]
        documents = [Document(text=t) for t in text_list]
        index = VectorStoreIndex.from_documents(documents)
        self.index = index

        from llama_index.indices.postprocessor import SimilarityPostprocessor
        from llama_index.query_engine import RetrieverQueryEngine
        from llama_index.indices.vector_store import VectorIndexRetriever
        from llama_index import get_response_synthesizer

        # configure retriever
        retriever = VectorIndexRetriever(
            index=index,
            similarity_top_k=2,
        )

        # configure response synthesizer
        response_synthesizer = get_response_synthesizer()

        # assemble query engine
        query_engine = RetrieverQueryEngine(
            retriever=retriever,
            response_synthesizer=response_synthesizer,
            node_postprocessors=[
                SimilarityPostprocessor(similarity_cutoff=0.7)
            ]
        )
        self.query_engine = query_engine
        # self.query_engine = index.as_query_engine()
        if is_persist:
            self.persist_index(persist_dir)

    def retrieve_index(self, persist_dir):
        storage_context = StorageContext.from_defaults(
            docstore=SimpleDocumentStore.from_persist_dir(persist_dir=persist_dir),
            vector_store=SimpleVectorStore.from_persist_dir(persist_dir=persist_dir),
            index_store=SimpleIndexStore.from_persist_dir(persist_dir=persist_dir),
        )
        self.index = load_index_from_storage(storage_context)
        self.query_engine = self.index.as_query_engine()

    def persist_index(self, persist_dir):
        self.index.storage_context.persist(persist_dir=persist_dir)

    def query(self, query_text):
        if not self.query_engine:
            raise Exception("No query engine loaded")
        return self.query_engine.query(query_text)

    def get_index(self):
        return self.index

langchain.embeddings.OpenAIEmbeddings
[((1194, 1252), 'logging.basicConfig', 'logging.basicConfig', ([], {'stream': 'sys.stdout', 'level': 'logging.INFO'}), '(stream=sys.stdout, level=logging.INFO)\n', (1213, 1252), False, 'import logging\n'), ((1435, 1462), 'os.getenv', 'os.getenv', (['"""OPENAI_API_KEY"""'], {}), "('OPENAI_API_KEY')\n", (1444, 1462), False, 'import os\n'), ((1478, 1593), 'llama_index.llms.AzureOpenAI', 'AzureOpenAI', ([], {'deployment_name': 'deployment_name', 'model': 'model_name', 'temperature': '(0)', 'engine': '"""gpt35"""', 'max_tokens': '(2048)'}), "(deployment_name=deployment_name, model=model_name, temperature=\n    0, engine='gpt35', max_tokens=2048)\n", (1489, 1593), False, 'from llama_index.llms import AzureOpenAI\n'), ((2106, 2170), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm': 'llm', 'embed_model': 'embedding_llm'}), '(llm=llm, embed_model=embedding_llm)\n', (2134, 2170), False, 'from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext\n'), ((2215, 2258), 'llama_index.set_global_service_context', 'set_global_service_context', (['service_context'], {}), '(service_context)\n', (2241, 2258), False, 'from llama_index import Document, set_global_service_context, StorageContext, load_index_from_storage, VectorStoreIndex\n'), ((2719, 2740), 'pandas.read_csv', 'pd.read_csv', (['csv_path'], {}), '(csv_path)\n', (2730, 2740), True, 'import pandas as pd\n'), ((2901, 2943), 'llama_index.VectorStoreIndex.from_documents', 'VectorStoreIndex.from_documents', (['documents'], {}), '(documents)\n', (2932, 2943), False, 'from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext\n'), ((3298, 3351), 'llama_index.indices.vector_store.VectorIndexRetriever', 'VectorIndexRetriever', ([], {'index': 'index', 'similarity_top_k': '(2)'}), '(index=index, similarity_top_k=2)\n', (3318, 3351), False, 'from llama_index.indices.vector_store import VectorIndexRetriever\n'), ((3460, 3486), 'llama_index.get_response_synthesizer', 'get_response_synthesizer', ([], {}), '()\n', (3484, 3486), False, 'from llama_index import get_response_synthesizer\n'), ((4321, 4361), 'llama_index.load_index_from_storage', 'load_index_from_storage', (['storage_context'], {}), '(storage_context)\n', (4344, 4361), False, 'from llama_index import Document, set_global_service_context, StorageContext, load_index_from_storage, VectorStoreIndex\n'), ((1716, 1932), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'model': 'embedding_model', 'deployment': 'embedding_deployment', 'openai_api_key': 'openai.api_key', 'openai_api_base': 'openai.api_base', 'openai_api_type': 'openai.api_type', 'openai_api_version': 'openai.api_version'}), '(model=embedding_model, deployment=embedding_deployment,\n    openai_api_key=openai.api_key, openai_api_base=openai.api_base,\n    openai_api_type=openai.api_type, openai_api_version=openai.api_version)\n', (1732, 1932), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((2848, 2864), 'llama_index.Document', 'Document', ([], {'text': 't'}), '(text=t)\n', (2856, 2864), False, 'from llama_index import Document, set_global_service_context, StorageContext, load_index_from_storage, VectorStoreIndex\n'), ((4057, 4118), 'llama_index.storage.docstore.SimpleDocumentStore.from_persist_dir', 'SimpleDocumentStore.from_persist_dir', ([], {'persist_dir': 'persist_dir'}), '(persist_dir=persist_dir)\n', (4093, 4118), False, 'from llama_index.storage.docstore import SimpleDocumentStore\n'), ((4145, 4204), 'llama_index.vector_stores.SimpleVectorStore.from_persist_dir', 'SimpleVectorStore.from_persist_dir', ([], {'persist_dir': 'persist_dir'}), '(persist_dir=persist_dir)\n', (4179, 4204), False, 'from llama_index.vector_stores import SimpleVectorStore\n'), ((4230, 4288), 'llama_index.storage.index_store.SimpleIndexStore.from_persist_dir', 'SimpleIndexStore.from_persist_dir', ([], {'persist_dir': 'persist_dir'}), '(persist_dir=persist_dir)\n', (4263, 4288), False, 'from llama_index.storage.index_store import SimpleIndexStore\n'), ((3703, 3749), 'llama_index.indices.postprocessor.SimilarityPostprocessor', 'SimilarityPostprocessor', ([], {'similarity_cutoff': '(0.7)'}), '(similarity_cutoff=0.7)\n', (3726, 3749), False, 'from llama_index.indices.postprocessor import SimilarityPostprocessor\n')]