
from langchain.agents import (
  initialize_agent,
  Tool,
  AgentType
)
from llama_index.callbacks import (
  CallbackManager,
  LlamaDebugHandler
)
from llama_index.node_parser.simple import SimpleNodeParser
from llama_index import (
    VectorStoreIndex,
    SummaryIndex,
    SimpleDirectoryReader,
    ServiceContext,
    StorageContext,
)

import os
import openai
import logging
import sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)


def init_llm_from_env(temperature=0.1, max_tokens=1024):
  llm_type = os.getenv("LLM")

  if llm_type == 'openai':
      from langchain.chat_models import ChatOpenAI

      openai.api_key = os.getenv("OPENAI_API_KEY")
      llm = ChatOpenAI(temperature=temperature,
                       model_name="gpt-3.5-turbo", 
                       max_tokens=max_tokens)

  elif llm_type == 'xinference':
      from langchain.llms import Xinference

      llm = Xinference(
          server_url=os.getenv("XINFERENCE_SERVER_ENDPOINT"), 
          model_uid=os.getenv("XINFERENCE_LLM_MODEL_UID")
        )
  else:
      raise ValueError(f"Unknown LLM type {llm_type}")

  return llm

def init_embedding_from_env(temperature=0.1, max_tokens=1024):
  embedding_type = os.getenv("EMBEDDING")

  if embedding_type == 'openai':
      from llama_index.embeddings import OpenAIEmbedding

      openai.api_key = os.getenv("OPENAI_API_KEY")
      embedding = OpenAIEmbedding()

  elif embedding_type == 'xinference':
      from langchain.embeddings import XinferenceEmbeddings
      from llama_index.embeddings import LangchainEmbedding

      embedding = LangchainEmbedding(
         XinferenceEmbeddings(
            server_url=os.getenv("XINFERENCE_SERVER_ENDPOINT"),
            model_uid=os.getenv("XINFERENCE_EMBEDDING_MODEL_UID")
         )
      )
  else:
      raise ValueError(f"Unknown EMBEDDING type {embedding_type}")

  return embedding

def get_service_context(callback_handlers):
    callback_manager = CallbackManager(callback_handlers)
    node_parser = SimpleNodeParser.from_defaults(
        chunk_size=512,
        chunk_overlap=128,
        callback_manager=callback_manager,
    )
    return ServiceContext.from_defaults(
        embed_model=init_embedding_from_env(),
        callback_manager=callback_manager,
        llm=init_llm_from_env(),
        chunk_size=512,
        node_parser=node_parser
    )

def get_storage_context():
    return StorageContext.from_defaults()

def get_langchain_agent_from_index(summary_index, vector_index):
    list_query_engine = summary_index.as_query_engine(
        response_mode="tree_summarize",
        use_async=True,
    )
    vector_query_engine = vector_index.as_query_engine(
        similarity_top_k=3
    )        
    tools = [
        Tool(
            name="Summary Tool",
            func=lambda q: str(list_query_engine.query(q)),
            description="useful for when you want to get summarizations",
            return_direct=True,
        ),
        Tool(
            name="Lookup Tool",
            func=lambda q: str(vector_query_engine.query(q)),
            description="useful for when you want to lookup detailed information",
            return_direct=True,
        ),
    ]

    agent_chain = initialize_agent(
        tools, 
        init_llm_from_env(), 
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, 
        verbose=True
    )
    return agent_chain


def get_query_engine_from_index(index):
    return index.as_query_engine(
        similarity_top_k=3
    )

def get_chat_engine_from_index(index):
    return index.as_chat_engine(chat_mode="condense_question", verbose=True)
class ChatEngine:

    def __init__(self, file_path):
        llama_debug = LlamaDebugHandler(print_trace_on_end=True)

        service_context = get_service_context([llama_debug])
        storage_context = get_storage_context()

        documents = SimpleDirectoryReader(input_files=[file_path], filename_as_id=True).load_data()
        logging.info(f"Loaded {len(documents)} documents from {file_path}")

        nodes = service_context.node_parser.get_nodes_from_documents(documents)        
        storage_context.docstore.add_documents(nodes)
        logging.info(f"Adding {len(nodes)} nodes to storage")
    
        self.summary_index = SummaryIndex(nodes, storage_context=storage_context, 
                                          service_context=service_context)
        self.vector_index = VectorStoreIndex(nodes, storage_context=storage_context,
                                             service_context=service_context)  

    # def conversational_chat(self, query, callback_handler):
    #     """
    #     Start a conversational chat with a agent
    #     """
    #     response = self.agent_chain.run(input=query, callbacks=[callback_handler])
    #     return response

    def conversational_chat(self, query, callback_handler):
        """
        Start a conversational chat with a agent
        """
        return get_chat_engine_from_index(self.vector_index).chat(query).response
langchain.chat_models.ChatOpenAI
[((398, 456), 'logging.basicConfig', 'logging.basicConfig', ([], {'stream': 'sys.stdout', 'level': 'logging.INFO'}), '(stream=sys.stdout, level=logging.INFO)\n', (417, 456), False, 'import logging\n'), ((529, 545), 'os.getenv', 'os.getenv', (['"""LLM"""'], {}), "('LLM')\n", (538, 545), False, 'import os\n'), ((1217, 1239), 'os.getenv', 'os.getenv', (['"""EMBEDDING"""'], {}), "('EMBEDDING')\n", (1226, 1239), False, 'import os\n'), ((1961, 1995), 'llama_index.callbacks.CallbackManager', 'CallbackManager', (['callback_handlers'], {}), '(callback_handlers)\n', (1976, 1995), False, 'from llama_index.callbacks import CallbackManager, LlamaDebugHandler\n'), ((2014, 2118), 'llama_index.node_parser.simple.SimpleNodeParser.from_defaults', 'SimpleNodeParser.from_defaults', ([], {'chunk_size': '(512)', 'chunk_overlap': '(128)', 'callback_manager': 'callback_manager'}), '(chunk_size=512, chunk_overlap=128,\n    callback_manager=callback_manager)\n', (2044, 2118), False, 'from llama_index.node_parser.simple import SimpleNodeParser\n'), ((2411, 2441), 'llama_index.StorageContext.from_defaults', 'StorageContext.from_defaults', ([], {}), '()\n', (2439, 2441), False, 'from llama_index import VectorStoreIndex, SummaryIndex, SimpleDirectoryReader, ServiceContext, StorageContext\n'), ((649, 676), 'os.getenv', 'os.getenv', (['"""OPENAI_API_KEY"""'], {}), "('OPENAI_API_KEY')\n", (658, 676), False, 'import os\n'), ((689, 780), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': 'temperature', 'model_name': '"""gpt-3.5-turbo"""', 'max_tokens': 'max_tokens'}), "(temperature=temperature, model_name='gpt-3.5-turbo', max_tokens=\n    max_tokens)\n", (699, 780), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1355, 1382), 'os.getenv', 'os.getenv', (['"""OPENAI_API_KEY"""'], {}), "('OPENAI_API_KEY')\n", (1364, 1382), False, 'import os\n'), ((1401, 1418), 'llama_index.embeddings.OpenAIEmbedding', 'OpenAIEmbedding', ([], {}), '()\n', (1416, 1418), False, 'from llama_index.embeddings import OpenAIEmbedding\n'), ((3697, 3739), 'llama_index.callbacks.LlamaDebugHandler', 'LlamaDebugHandler', ([], {'print_trace_on_end': '(True)'}), '(print_trace_on_end=True)\n', (3714, 3739), False, 'from llama_index.callbacks import CallbackManager, LlamaDebugHandler\n'), ((4266, 4356), 'llama_index.SummaryIndex', 'SummaryIndex', (['nodes'], {'storage_context': 'storage_context', 'service_context': 'service_context'}), '(nodes, storage_context=storage_context, service_context=\n    service_context)\n', (4278, 4356), False, 'from llama_index import VectorStoreIndex, SummaryIndex, SimpleDirectoryReader, ServiceContext, StorageContext\n'), ((4423, 4517), 'llama_index.VectorStoreIndex', 'VectorStoreIndex', (['nodes'], {'storage_context': 'storage_context', 'service_context': 'service_context'}), '(nodes, storage_context=storage_context, service_context=\n    service_context)\n', (4439, 4517), False, 'from llama_index import VectorStoreIndex, SummaryIndex, SimpleDirectoryReader, ServiceContext, StorageContext\n'), ((3871, 3938), 'llama_index.SimpleDirectoryReader', 'SimpleDirectoryReader', ([], {'input_files': '[file_path]', 'filename_as_id': '(True)'}), '(input_files=[file_path], filename_as_id=True)\n', (3892, 3938), False, 'from llama_index import VectorStoreIndex, SummaryIndex, SimpleDirectoryReader, ServiceContext, StorageContext\n'), ((947, 986), 'os.getenv', 'os.getenv', (['"""XINFERENCE_SERVER_ENDPOINT"""'], {}), "('XINFERENCE_SERVER_ENDPOINT')\n", (956, 986), False, 'import os\n'), ((1009, 1046), 'os.getenv', 'os.getenv', (['"""XINFERENCE_LLM_MODEL_UID"""'], {}), "('XINFERENCE_LLM_MODEL_UID')\n", (1018, 1046), False, 'import os\n'), ((1672, 1711), 'os.getenv', 'os.getenv', (['"""XINFERENCE_SERVER_ENDPOINT"""'], {}), "('XINFERENCE_SERVER_ENDPOINT')\n", (1681, 1711), False, 'import os\n'), ((1735, 1778), 'os.getenv', 'os.getenv', (['"""XINFERENCE_EMBEDDING_MODEL_UID"""'], {}), "('XINFERENCE_EMBEDDING_MODEL_UID')\n", (1744, 1778), False, 'import os\n')]