code: import argparse
  from typing import Optional, Any
  from mlx_lm import load, generate
  from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, Settings
  from llama_index.core.llms.callbacks import llm_completion_callback 
  from llama_index.core.llms import CustomLLM, CompletionResponse, CompletionResponseGen, LLMMetadata
  from pydantic import BaseModel
  
  class OurLLM(CustomLLM, BaseModel):
      model: Optional[Any] = None
      tokenizer: Optional[Any] = None
  
      def __init__(self, model_name: str, **data):
          super().__init__(**data)  # Initialize BaseModel part with data
          # Directly load the model and tokenizer
          self.model, self.tokenizer = load(model_name)
      context_window: int = 2096
      max_tokens : int = 100
      model_name: str = "custom"
      
      @property
      def metadata(self) -> LLMMetadata:
          """Get LLM metadata."""
          return LLMMetadata(
              context_window = self.context_window,
              model_name=self.model_name,
              max_tokens=self.max_tokens
          )
  
      def process_generated_text(self, text: str) -> str:
          token_pos = text.find("\n\n")
          if token_pos != -1:
              # Truncate text at the first occurrence of two new lines
              return text[:token_pos]
          return text
  
      @llm_completion_callback()
      def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         # Remove 'formatted' argument if present
          kwargs.pop('formatted', None)
      
          generated_text = generate(self.model, self.tokenizer, prompt=prompt, verbose=False, **kwargs)
          processed_text = self.process_generated_text(generated_text)
          return CompletionResponse(text=processed_text)
  
  
      @llm_completion_callback()
      def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
          generated_text = generate(self.model, self.tokenizer, prompt=prompt, verbose=False, **kwargs)
          processed_text = self.process_generated_text(generated_text)
          for char in processed_text:  
              yield CompletionResponse(text=char, delta=char)
  
  
  def main():
      parser = argparse.ArgumentParser(description="Query a document collection with an LLM.")
      parser.add_argument("--model_name", type=str, required=True, help="Model name to use for the LLM.")
      parser.add_argument("--directory", type=str, required=True, help="Directory containing the documents to index.")
      parser.add_argument("--embed_model", type=str, required=True, help="Embed model to use for vectorizing documents.")
      parser.add_argument("--query", type=str, required=True, help="Query to perform on the document collection.")
  
      args = parser.parse_args()
  
      # Convert input from the user into strings
      model_name = str(args.model_name)
      directory = str(args.directory)
      embed_model = str(args.embed_model)
      query_str = str(args.query)
  
      # Setup the LLM and embed model
      Settings.llm = OurLLM(model_name=model_name)
      Settings.embed_model = embed_model
  
      # Load documents and create index
      documents = SimpleDirectoryReader(directory).load_data()
      index = VectorStoreIndex.from_documents(documents, show_progress=True)
  
      # Perform query and print response
      query_engine = index.as_query_engine()
      response = query_engine.query(query_str)
      print(response)
  
  
  if __name__ == "__main__":
      main()
  
  #Usage: python mlx_rag.py --model_name "mlx-community/Mistral-7B-v0.1-hf-4bit-mlx" --directory "data" --embed_model "local:BAAI/bge-base-en-v1.5" --query "Complete the sentence: In all criminal prosecutions, the accused shall enjoy"
  
apis: llama_index.core.VectorStoreIndex.from_documents llama_index.core.SimpleDirectoryReader llama_index.core.llms.CompletionResponse llama_index.core.llms.LLMMetadata llama_index.core.llms.callbacks.llm_completion_callback
extract_api: [((1291, 1316), 'llama_index.core.llms.callbacks.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (1314, 1316), False, 'from llama_index.core.llms.callbacks import llm_completion_callback\n'), ((1715, 1740), 'llama_index.core.llms.callbacks.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (1738, 1740), False, 'from llama_index.core.llms.callbacks import llm_completion_callback\n'), ((2121, 2200), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""Query a document collection with an LLM."""'}), "(description='Query a document collection with an LLM.')\n", (2144, 2200), False, 'import argparse\n'), ((3118, 3180), 'llama_index.core.VectorStoreIndex.from_documents', 'VectorStoreIndex.from_documents', (['documents'], {'show_progress': '(True)'}), '(documents, show_progress=True)\n', (3149, 3180), False, 'from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, Settings\n'), ((674, 690), 'mlx_lm.load', 'load', (['model_name'], {}), '(model_name)\n', (678, 690), False, 'from mlx_lm import load, generate\n'), ((885, 992), 'llama_index.core.llms.LLMMetadata', 'LLMMetadata', ([], {'context_window': 'self.context_window', 'model_name': 'self.model_name', 'max_tokens': 'self.max_tokens'}), '(context_window=self.context_window, model_name=self.model_name,\n    max_tokens=self.max_tokens)\n', (896, 992), False, 'from llama_index.core.llms import CustomLLM, CompletionResponse, CompletionResponseGen, LLMMetadata\n'), ((1507, 1583), 'mlx_lm.generate', 'generate', (['self.model', 'self.tokenizer'], {'prompt': 'prompt', 'verbose': '(False)'}), '(self.model, self.tokenizer, prompt=prompt, verbose=False, **kwargs)\n', (1515, 1583), False, 'from mlx_lm import load, generate\n'), ((1668, 1707), 'llama_index.core.llms.CompletionResponse', 'CompletionResponse', ([], {'text': 'processed_text'}), '(text=processed_text)\n', (1686, 1707), False, 'from llama_index.core.llms import CustomLLM, CompletionResponse, CompletionResponseGen, LLMMetadata\n'), ((1850, 1926), 'mlx_lm.generate', 'generate', (['self.model', 'self.tokenizer'], {'prompt': 'prompt', 'verbose': '(False)'}), '(self.model, self.tokenizer, prompt=prompt, verbose=False, **kwargs)\n', (1858, 1926), False, 'from mlx_lm import load, generate\n'), ((3061, 3093), 'llama_index.core.SimpleDirectoryReader', 'SimpleDirectoryReader', (['directory'], {}), '(directory)\n', (3082, 3093), False, 'from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, Settings\n'), ((2052, 2093), 'llama_index.core.llms.CompletionResponse', 'CompletionResponse', ([], {'text': 'char', 'delta': 'char'}), '(text=char, delta=char)\n', (2070, 2093), False, 'from llama_index.core.llms import CustomLLM, CompletionResponse, CompletionResponseGen, LLMMetadata\n')]