code: #!/usr/bin/env python3
  # Copyright (c) 2023-2024 Steve Castellotti
  # This file is part of Urcuchillay and is released under the MIT License.
  # See LICENSE file in the project root for full license information.
  import logging
  import os
  import sys
  
  import config
  
  try:
      import chromadb
      import llama_index
      import llama_index.vector_stores
      import transformers
      import utils
  except ModuleNotFoundError as e:
      print('\nError importing Python module(s)')
      print('If installed using setup.sh it may be necessary to run:\n')
      print('pyenv activate urcuchillay-env\n')
      sys.exit(1)
  
  
  class Client:
      def __init__(self, args):
  
          self.debug = args.debug
  
          llama_debug = llama_index.callbacks.LlamaDebugHandler(print_trace_on_end=self.debug)
          self.callback_manager = llama_index.callbacks.CallbackManager([llama_debug])
  
          # Fallback settings for api_base, api_key, and api_version
          os.environ['OPENAI_API_BASE'] = config.APIConfig.get_openai_api_base(host=args.api_host, port=args.api_port)
          os.environ['OPENAI_API_KEY'] = config.APIConfig.OPENAI_API_KEY
          os.environ['OPENAI_API_VERSION'] = config.APIConfig.OPENAI_API_VERSION
  
          # Set Parallel Iterator
          os.environ['TOKENIZERS_PARALLELISM'] = 'true' if config.Config.TOKENIZERS_PARALLELISM else 'false'
  
          # ChromaDB Settings
          self.db = None
          self.chromadb_settings = chromadb.config.Settings(
              anonymized_telemetry=config.Config.ANONYMIZED_TELEMETRY,
              allow_reset=config.Config.ALLOW_RESET,
          )
  
          self.llm = None
          self.service_context = None
          self.index = None
  
      def get_llm(self, args):
          return llama_index.llms.OpenAI(
              model='text-davinci-003',
              temperature=args.temperature,
              max_tokens=args.context,
              api_base=config.APIConfig.get_openai_api_base(host=args.api_host, port=args.api_port),
              api_key=config.APIConfig.OPENAI_API_KEY,
              api_version=config.APIConfig.OPENAI_API_VERSION,
              max_retries=args.max_retries,
              timeout=args.timeout,
              callback_manager=self.callback_manager,
          )
  
      def get_service_context(self, llm, args):
          embed_model = config.Config.EMBED_MODEL_NAME
          if hasattr(args, 'embed_model_name'):
              if args.embed_model_name == 'default' or args.embed_model_name == 'local':
                  embed_model = args.embed_model_name
              else:
                  if hasattr(args, 'embed_model_provider'):
                      # use Huggingface embeddings
                      embed_model = llama_index.embeddings.HuggingFaceEmbedding(
                          model_name=args.embed_model_provider + '/' + args.embed_model_name)
  
          return llama_index.ServiceContext.from_defaults(
              llm=llm,
              embed_model=embed_model,
              callback_manager=self.callback_manager,
              context_window=args.context,
              num_output=args.max_new_tokens,
          )
  
      def get_index(self, service_context, args, storage_type=config.Config.STORAGE_TYPE):
          if storage_type == 'json' and self.index:
              return self.get_index(service_context, args)
          elif storage_type == 'chromadb':
              return self.get_index_chroma(service_context, args)
          else:
              return None
  
      def save_index(self, args, storage_type=config.Config.STORAGE_TYPE):
          if storage_type == 'json' and self.index:
              self.index.storage_context.persist(persist_dir=args.storage)
          elif storage_type == 'chromadb':
              # For ChromaDB, storage is already written to disk
              # as part of the loading data process
              pass
  
      def reset_index(self, args):
          logging.warning('resetting index')
          if config.Config.STORAGE_TYPE == 'json':
              utils.storage_reset(storage_path=args.storage)
          elif config.Config.STORAGE_TYPE == 'chromadb':
              if not self.db:
                  self.db = chromadb.PersistentClient(
                      settings=self.chromadb_settings,
                      path=args.storage
                  )
              self.db.reset()
              self.reset_chroma_collection()
              self.service_context = self.get_service_context(self.llm, args)
              self.index = self.get_index(self.service_context, args)
  
      @staticmethod
      def get_index_json(service_context, args):
          if args.load and all(os.path.exists(os.path.join(args.storage, filename))
                               for filename in config.Config.STORAGE_FILES):
              # load vector index from storage
              storage_context = llama_index.StorageContext.from_defaults(persist_dir=args.storage)
              return llama_index.load_index_from_storage(storage_context, service_context=service_context)
          else:
              if not os.path.exists(args.data) or not os.listdir(args.data):
                  # Create a temporary empty file for the index if a missing or empty data directory was supplied
                  temp_file = utils.create_temporary_empty_file()
                  documents = llama_index.SimpleDirectoryReader(input_files=[temp_file]).load_data()
                  index = llama_index.VectorStoreIndex.from_documents(documents, service_context=service_context)
                  os.remove(temp_file)
                  return index
              else:
                  documents = llama_index.SimpleDirectoryReader(args.data).load_data()
                  return llama_index.VectorStoreIndex.from_documents(
                      documents, service_context=service_context
                  )
  
      def get_index_chroma(self, service_context, args):
  
          if not self.db:
              self.db = chromadb.PersistentClient(
                  settings=self.chromadb_settings,
                  path=args.storage
              )
  
          chroma_collection = self.db.get_or_create_collection('quickstart')
  
          # set up ChromaVectorStore and load in data
          vector_store = llama_index.vector_stores.ChromaVectorStore(chroma_collection=chroma_collection)
  
          if args.load:
              # noinspection PyTypeChecker
              index = llama_index.VectorStoreIndex.from_vector_store(
                  vector_store,
                  service_context=service_context,
              )
          else:
              storage_context = llama_index.storage.storage_context.StorageContext.from_defaults(
                  vector_store=vector_store)
  
              documents = llama_index.SimpleDirectoryReader(args.data).load_data()
              index = llama_index.VectorStoreIndex.from_documents(
                  documents, storage_context=storage_context, service_context=service_context
              )
  
          return index
  
      @staticmethod
      def reset_chroma_collection():
          chroma_client = chromadb.EphemeralClient()
          for collection in chroma_client.list_collections():
              chroma_client.delete_collection(name=collection.name)
  
apis: llama_index.VectorStoreIndex.from_documents llama_index.callbacks.LlamaDebugHandler llama_index.ServiceContext.from_defaults llama_index.vector_stores.ChromaVectorStore llama_index.StorageContext.from_defaults llama_index.storage.storage_context.StorageContext.from_defaults llama_index.SimpleDirectoryReader llama_index.VectorStoreIndex.from_vector_store llama_index.embeddings.HuggingFaceEmbedding llama_index.callbacks.CallbackManager llama_index.load_index_from_storage
extract_api: [((590, 601), 'sys.exit', 'sys.exit', (['(1)'], {}), '(1)\n', (598, 601), False, 'import sys\n'), ((704, 774), 'llama_index.callbacks.LlamaDebugHandler', 'llama_index.callbacks.LlamaDebugHandler', ([], {'print_trace_on_end': 'self.debug'}), '(print_trace_on_end=self.debug)\n', (743, 774), False, 'import llama_index\n'), ((807, 859), 'llama_index.callbacks.CallbackManager', 'llama_index.callbacks.CallbackManager', (['[llama_debug]'], {}), '([llama_debug])\n', (844, 859), False, 'import llama_index\n'), ((968, 1044), 'config.APIConfig.get_openai_api_base', 'config.APIConfig.get_openai_api_base', ([], {'host': 'args.api_host', 'port': 'args.api_port'}), '(host=args.api_host, port=args.api_port)\n', (1004, 1044), False, 'import config\n'), ((1420, 1545), 'chromadb.config.Settings', 'chromadb.config.Settings', ([], {'anonymized_telemetry': 'config.Config.ANONYMIZED_TELEMETRY', 'allow_reset': 'config.Config.ALLOW_RESET'}), '(anonymized_telemetry=config.Config.\n    ANONYMIZED_TELEMETRY, allow_reset=config.Config.ALLOW_RESET)\n', (1444, 1545), False, 'import chromadb\n'), ((2798, 2981), 'llama_index.ServiceContext.from_defaults', 'llama_index.ServiceContext.from_defaults', ([], {'llm': 'llm', 'embed_model': 'embed_model', 'callback_manager': 'self.callback_manager', 'context_window': 'args.context', 'num_output': 'args.max_new_tokens'}), '(llm=llm, embed_model=embed_model,\n    callback_manager=self.callback_manager, context_window=args.context,\n    num_output=args.max_new_tokens)\n', (2838, 2981), False, 'import llama_index\n'), ((3795, 3829), 'logging.warning', 'logging.warning', (['"""resetting index"""'], {}), "('resetting index')\n", (3810, 3829), False, 'import logging\n'), ((6041, 6126), 'llama_index.vector_stores.ChromaVectorStore', 'llama_index.vector_stores.ChromaVectorStore', ([], {'chroma_collection': 'chroma_collection'}), '(chroma_collection=chroma_collection\n    )\n', (6084, 6126), False, 'import llama_index\n'), ((6853, 6879), 'chromadb.EphemeralClient', 'chromadb.EphemeralClient', ([], {}), '()\n', (6877, 6879), False, 'import chromadb\n'), ((3891, 3937), 'utils.storage_reset', 'utils.storage_reset', ([], {'storage_path': 'args.storage'}), '(storage_path=args.storage)\n', (3910, 3937), False, 'import utils\n'), ((4696, 4762), 'llama_index.StorageContext.from_defaults', 'llama_index.StorageContext.from_defaults', ([], {'persist_dir': 'args.storage'}), '(persist_dir=args.storage)\n', (4736, 4762), False, 'import llama_index\n'), ((4782, 4872), 'llama_index.load_index_from_storage', 'llama_index.load_index_from_storage', (['storage_context'], {'service_context': 'service_context'}), '(storage_context, service_context=\n    service_context)\n', (4817, 4872), False, 'import llama_index\n'), ((5765, 5842), 'chromadb.PersistentClient', 'chromadb.PersistentClient', ([], {'settings': 'self.chromadb_settings', 'path': 'args.storage'}), '(settings=self.chromadb_settings, path=args.storage)\n', (5790, 5842), False, 'import chromadb\n'), ((6206, 6303), 'llama_index.VectorStoreIndex.from_vector_store', 'llama_index.VectorStoreIndex.from_vector_store', (['vector_store'], {'service_context': 'service_context'}), '(vector_store,\n    service_context=service_context)\n', (6252, 6303), False, 'import llama_index\n'), ((6391, 6487), 'llama_index.storage.storage_context.StorageContext.from_defaults', 'llama_index.storage.storage_context.StorageContext.from_defaults', ([], {'vector_store': 'vector_store'}), '(vector_store\n    =vector_store)\n', (6455, 6487), False, 'import llama_index\n'), ((6602, 6727), 'llama_index.VectorStoreIndex.from_documents', 'llama_index.VectorStoreIndex.from_documents', (['documents'], {'storage_context': 'storage_context', 'service_context': 'service_context'}), '(documents, storage_context=\n    storage_context, service_context=service_context)\n', (6645, 6727), False, 'import llama_index\n'), ((1871, 1947), 'config.APIConfig.get_openai_api_base', 'config.APIConfig.get_openai_api_base', ([], {'host': 'args.api_host', 'port': 'args.api_port'}), '(host=args.api_host, port=args.api_port)\n', (1907, 1947), False, 'import config\n'), ((5097, 5132), 'utils.create_temporary_empty_file', 'utils.create_temporary_empty_file', ([], {}), '()\n', (5130, 5132), False, 'import utils\n'), ((5256, 5348), 'llama_index.VectorStoreIndex.from_documents', 'llama_index.VectorStoreIndex.from_documents', (['documents'], {'service_context': 'service_context'}), '(documents, service_context=\n    service_context)\n', (5299, 5348), False, 'import llama_index\n'), ((5360, 5380), 'os.remove', 'os.remove', (['temp_file'], {}), '(temp_file)\n', (5369, 5380), False, 'import os\n'), ((5536, 5628), 'llama_index.VectorStoreIndex.from_documents', 'llama_index.VectorStoreIndex.from_documents', (['documents'], {'service_context': 'service_context'}), '(documents, service_context=\n    service_context)\n', (5579, 5628), False, 'import llama_index\n'), ((2645, 2761), 'llama_index.embeddings.HuggingFaceEmbedding', 'llama_index.embeddings.HuggingFaceEmbedding', ([], {'model_name': "(args.embed_model_provider + '/' + args.embed_model_name)"}), "(model_name=args.\n    embed_model_provider + '/' + args.embed_model_name)\n", (2688, 2761), False, 'import llama_index\n'), ((4047, 4124), 'chromadb.PersistentClient', 'chromadb.PersistentClient', ([], {'settings': 'self.chromadb_settings', 'path': 'args.storage'}), '(settings=self.chromadb_settings, path=args.storage)\n', (4072, 4124), False, 'import chromadb\n'), ((4901, 4926), 'os.path.exists', 'os.path.exists', (['args.data'], {}), '(args.data)\n', (4915, 4926), False, 'import os\n'), ((4934, 4955), 'os.listdir', 'os.listdir', (['args.data'], {}), '(args.data)\n', (4944, 4955), False, 'import os\n'), ((6525, 6569), 'llama_index.SimpleDirectoryReader', 'llama_index.SimpleDirectoryReader', (['args.data'], {}), '(args.data)\n', (6558, 6569), False, 'import llama_index\n'), ((4508, 4544), 'os.path.join', 'os.path.join', (['args.storage', 'filename'], {}), '(args.storage, filename)\n', (4520, 4544), False, 'import os\n'), ((5161, 5219), 'llama_index.SimpleDirectoryReader', 'llama_index.SimpleDirectoryReader', ([], {'input_files': '[temp_file]'}), '(input_files=[temp_file])\n', (5194, 5219), False, 'import llama_index\n'), ((5456, 5500), 'llama_index.SimpleDirectoryReader', 'llama_index.SimpleDirectoryReader', (['args.data'], {}), '(args.data)\n', (5489, 5500), False, 'import llama_index\n')]