code: """Memory for agents."""
  import os
  from llama_index.core import GPTVectorStoreIndex, GPTListIndex
  from llama_index.core import Document, ServiceContext
  from llama_index.legacy import LLMPredictor
  from llama_index.core import StorageContext, load_index_from_storage
  from langchain_community.chat_models import ChatOpenAI
  
  import logging
  
  logging.basicConfig(level=logging.INFO)
  logger = logging.getLogger(__name__)
  
  # https://gpt-index.readthedocs.io/en/latest/guides/index_guide.html
  INDEX_TYPES = {
      # Good for retrieval, because of top_k and embeddings.
      "vector": GPTVectorStoreIndex,
      # Good for aggregate summaries, but slow.
      "list": GPTListIndex,
  }
  
  LLM_PREDICTOR_TYPES = {
      "gpt-3.5-turbo": ChatOpenAI,
      "gpt-3.5-turbo-16k": ChatOpenAI,
      "gpt-4": ChatOpenAI,
  }
  
  
  class Memory:
      def __init__(self, memory_folder=None, index_type="vector", llm_predictor="gpt-3.5-turbo"):
          assert index_type in INDEX_TYPES, f"Invalid index type: {index_type}"
          assert llm_predictor in LLM_PREDICTOR_TYPES, f"Invalid LLM predictor: {llm_predictor}"
  
          self.texts = []
          llm_kwargs = {"temperature": 0, "model_name": llm_predictor}
          predictor_constructor = LLM_PREDICTOR_TYPES[llm_predictor]
          llm = LLMPredictor(llm=predictor_constructor(**llm_kwargs))
          service_context = ServiceContext.from_defaults(llm_predictor=llm)
  
          if memory_folder and os.path.exists(memory_folder):
              logger.info("Loading memory from disk.")
              storage_context = StorageContext.from_defaults(persist_dir=memory_folder)
              self.index = load_index_from_storage(storage_context)
          else:
              self.index = INDEX_TYPES[index_type].from_documents([], service_context=service_context)
          self.llm_predictor = llm_predictor
  
      def query(self, prompt, similarity_top_k=3):
          query_engine = self.index.as_query_engine(similarity_top_k=similarity_top_k)
          return query_engine.query(prompt)
  
      def add(self, text):
          if text in self.texts:
              logger.info("Skipping duplicate text.")
              return
          self.texts.append(text)
          self.index.insert(Document(text=text))
  
      def save(self, path):
          self.index.storage_context.persist(path)
apis: llama_index.core.StorageContext.from_defaults llama_index.core.load_index_from_storage llama_index.core.Document llama_index.core.ServiceContext.from_defaults
extract_api: [((337, 376), 'logging.basicConfig', 'logging.basicConfig', ([], {'level': 'logging.INFO'}), '(level=logging.INFO)\n', (356, 376), False, 'import logging\n'), ((386, 413), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (403, 413), False, 'import logging\n'), ((1330, 1377), 'llama_index.core.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm_predictor': 'llm'}), '(llm_predictor=llm)\n', (1358, 1377), False, 'from llama_index.core import Document, ServiceContext\n'), ((1408, 1437), 'os.path.exists', 'os.path.exists', (['memory_folder'], {}), '(memory_folder)\n', (1422, 1437), False, 'import os\n'), ((1522, 1577), 'llama_index.core.StorageContext.from_defaults', 'StorageContext.from_defaults', ([], {'persist_dir': 'memory_folder'}), '(persist_dir=memory_folder)\n', (1550, 1577), False, 'from llama_index.core import StorageContext, load_index_from_storage\n'), ((1603, 1643), 'llama_index.core.load_index_from_storage', 'load_index_from_storage', (['storage_context'], {}), '(storage_context)\n', (1626, 1643), False, 'from llama_index.core import StorageContext, load_index_from_storage\n'), ((2165, 2184), 'llama_index.core.Document', 'Document', ([], {'text': 'text'}), '(text=text)\n', (2173, 2184), False, 'from llama_index.core import Document, ServiceContext\n')]