from approaches.index.store.cosmos_index_store import CosmosIndexStore
from llama_index import StorageContext
from approaches.index.store.cosmos_doc_store import CosmosDocumentStore
from llama_index import load_index_from_storage
import os
import openai
from langchain.chat_models  import AzureChatOpenAI
from langchain.embeddings import OpenAIEmbeddings
from llama_index import LangchainEmbedding
from llama_index.vector_stores import QdrantVectorStore
from llama_index import (
    LLMPredictor,
    ServiceContext
)
from llama_index.node_parser import SimpleNodeParser
from llama_index.langchain_helpers.text_splitter import TokenTextSplitter
from llama_index import SimpleDirectoryReader, Document
from llama_index.indices.knowledge_graph.base import GPTKnowledgeGraphIndex
import qdrant_client

from dotenv import load_dotenv

load_dotenv()

AZURE_INDEX_STORAGE_CONNECTION_STRING = os.environ.get("AZURE_INDEX_STORAGE_CONNECTION_STRING") or None
AZURE_QDRANT_HOST = os.environ.get("AZURE_QDRANT_HOST") or None
AZURE_OPENAI_API_BASE = os.environ.get("AZURE_OPENAI_BASE")
AZURE_OPENAI_API_KEY = os.environ.get("AZURE_OPENAI_API_KEY_SOUTH_CENTRAL_US")
AZURE_OPENAI_CHATGPT_DEPLOYMENT = os.environ.get("AZURE_OPENAI_CHATGPT_DEPLOYMENT")

openai.api_type = "azure"
openai.api_base = AZURE_OPENAI_API_BASE
openai.api_version = "2023-03-15-preview"
os.environ["OPENAI_API_KEY"] = str(AZURE_OPENAI_API_KEY)
openai.api_key = AZURE_OPENAI_API_KEY

class GPTKGIndexer:

    def __init__(self):
        if AZURE_INDEX_STORAGE_CONNECTION_STRING is None or AZURE_QDRANT_HOST is None:
            return
        self._connection_string = AZURE_INDEX_STORAGE_CONNECTION_STRING
        self._index_store = CosmosIndexStore.from_uri(uri=str(self._connection_string), db_name="kg_index")
        self._doc_store = CosmosDocumentStore.from_uri(uri=str(self._connection_string), db_name = "doc_store")
        self._storage_context = StorageContext.from_defaults(
            docstore=self._doc_store,
            index_store=self._index_store)
        self._llm = AzureChatOpenAI(deployment_name=str(AZURE_OPENAI_CHATGPT_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,
            temperature=0.0
        )
        llm_predictor = LLMPredictor(llm=self._llm)

        self._embedding_llm = LangchainEmbedding(
            OpenAIEmbeddings(
                model="text-embedding-ada-002",
                deployment="text-embedding-ada-002",
                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,
        )
        self._service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, chunk_size_limit=1024)
        try:
            print("Loading index from storage")
            self.index = load_index_from_storage(storage_context=self._storage_context, service_context = self._service_context)
            print("Index loaded from storage")
        except:
            print("Initializing new index")
            self.index = self._init_index()
            print("Initialized new index")
        

    def add_document(self, fileContent: str):
        text_splitter = TokenTextSplitter(separator=" ", chunk_size=2048, chunk_overlap=20)
        text_chunks = text_splitter.split_text(fileContent)
        doc_chunks = [Document(t) for t in text_chunks]
        for doc_chunk in doc_chunks:
            self.index.insert(doc_chunk)

    def query(self, question: str):
        query_engine = self.index.as_query_engine(
            include_text=False, 
            response_mode="tree_summarize"
        )
        response = query_engine.query(question)
        return response

    def _init_index(self):
        self.index = GPTKnowledgeGraphIndex(
            [],
            service_context=self._service_context,
            storage_context=self._storage_context
        )
langchain.embeddings.OpenAIEmbeddings
[((832, 845), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (843, 845), False, 'from dotenv import load_dotenv\n'), ((1039, 1074), 'os.environ.get', 'os.environ.get', (['"""AZURE_OPENAI_BASE"""'], {}), "('AZURE_OPENAI_BASE')\n", (1053, 1074), False, 'import os\n'), ((1098, 1153), 'os.environ.get', 'os.environ.get', (['"""AZURE_OPENAI_API_KEY_SOUTH_CENTRAL_US"""'], {}), "('AZURE_OPENAI_API_KEY_SOUTH_CENTRAL_US')\n", (1112, 1153), False, 'import os\n'), ((1188, 1237), 'os.environ.get', 'os.environ.get', (['"""AZURE_OPENAI_CHATGPT_DEPLOYMENT"""'], {}), "('AZURE_OPENAI_CHATGPT_DEPLOYMENT')\n", (1202, 1237), False, 'import os\n'), ((887, 942), 'os.environ.get', 'os.environ.get', (['"""AZURE_INDEX_STORAGE_CONNECTION_STRING"""'], {}), "('AZURE_INDEX_STORAGE_CONNECTION_STRING')\n", (901, 942), False, 'import os\n'), ((971, 1006), 'os.environ.get', 'os.environ.get', (['"""AZURE_QDRANT_HOST"""'], {}), "('AZURE_QDRANT_HOST')\n", (985, 1006), False, 'import os\n'), ((1918, 2008), 'llama_index.StorageContext.from_defaults', 'StorageContext.from_defaults', ([], {'docstore': 'self._doc_store', 'index_store': 'self._index_store'}), '(docstore=self._doc_store, index_store=self.\n    _index_store)\n', (1946, 2008), False, 'from llama_index import StorageContext\n'), ((2366, 2393), 'llama_index.LLMPredictor', 'LLMPredictor', ([], {'llm': 'self._llm'}), '(llm=self._llm)\n', (2378, 2393), False, 'from llama_index import LLMPredictor, ServiceContext\n'), ((2866, 2951), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm_predictor': 'llm_predictor', 'chunk_size_limit': '(1024)'}), '(llm_predictor=llm_predictor, chunk_size_limit=1024\n    )\n', (2894, 2951), False, 'from llama_index import LLMPredictor, ServiceContext\n'), ((3411, 3478), 'llama_index.langchain_helpers.text_splitter.TokenTextSplitter', 'TokenTextSplitter', ([], {'separator': '""" """', 'chunk_size': '(2048)', 'chunk_overlap': '(20)'}), "(separator=' ', chunk_size=2048, chunk_overlap=20)\n", (3428, 3478), False, 'from llama_index.langchain_helpers.text_splitter import TokenTextSplitter\n'), ((3968, 4076), 'llama_index.indices.knowledge_graph.base.GPTKnowledgeGraphIndex', 'GPTKnowledgeGraphIndex', (['[]'], {'service_context': 'self._service_context', 'storage_context': 'self._storage_context'}), '([], service_context=self._service_context,\n    storage_context=self._storage_context)\n', (3990, 4076), False, 'from llama_index.indices.knowledge_graph.base import GPTKnowledgeGraphIndex\n'), ((2457, 2691), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'model': '"""text-embedding-ada-002"""', 'deployment': '"""text-embedding-ada-002"""', '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='text-embedding-ada-002', deployment=\n    'text-embedding-ada-002', openai_api_key=openai.api_key,\n    openai_api_base=openai.api_base, openai_api_type=openai.api_type,\n    openai_api_version=openai.api_version)\n", (2473, 2691), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((3033, 3138), 'llama_index.load_index_from_storage', 'load_index_from_storage', ([], {'storage_context': 'self._storage_context', 'service_context': 'self._service_context'}), '(storage_context=self._storage_context,\n    service_context=self._service_context)\n', (3056, 3138), False, 'from llama_index import load_index_from_storage\n'), ((3561, 3572), 'llama_index.Document', 'Document', (['t'], {}), '(t)\n', (3569, 3572), False, 'from llama_index import SimpleDirectoryReader, Document\n')]