clarifai 10.0.0__py3-none-any.whl → 10.1.0__py3-none-any.whl

clarifai/models/model_serving/repo_build/build.py

@@ -0,0 +1,198 @@
+# Copyright 2023 Clarifai, Inc.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Triton python backend inference model controller.
+"""
+
+import inspect
+import os
+import shutil
+import zipfile
+from dataclasses import asdict
+from pathlib import Path
+from typing import Iterable, Literal, Union
+
+import yaml
+from tqdm import tqdm
+
+from ..constants import BUILT_MODEL_EXT
+from ..model_config import MODEL_TYPES, ModelConfigClass, base, get_model_config, load_user_config
+from ..model_config.base import *  # noqa
+from ..model_config.config import parse_config
+from ..model_config.triton.serializer import Serializer
+
+
+def __parse_type_to_class():
+  _t = {}
+  _classes = inspect.getmembers(base, inspect.isclass)
+  for cls_name, cls_obj in _classes:
+    if cls_obj.__base__ is base._BaseClarifaiModel:
+      _t.update({cls_obj._config.clarifai_model.type: cls_name})
+  return _t
+
+
+_TYPE_TO_CLASS = __parse_type_to_class()
+
+
+def _get_static_file_path(relative_path: str):
+  curr_dir = os.path.dirname(__file__)
+  return os.path.join(curr_dir, "static_files", relative_path)
+
+
+def _read_static_file(relative_path: str):
+  path = _get_static_file_path(relative_path)
+  with open(path, "r") as f:
+    return f.read()
+
+
+def copy_folder(src_folder, dest_folder, exclude_items=None):
+
+  if exclude_items is None:
+    exclude_items = set()
+
+  # Ensure the destination folder exists
+  if not os.path.exists(dest_folder):
+    os.makedirs(dest_folder)
+
+  loader = tqdm(os.listdir(src_folder))
+  if exclude_items:
+    print(f"NOTE: skipping {exclude_items}")
+
+  for item in loader:
+    loader.set_description(f"copying {item}...")
+    src_item = os.path.join(src_folder, item)
+    dest_item = os.path.join(dest_folder, item)
+
+    # Skip items in the exclude list
+    if item in exclude_items or item.endswith(BUILT_MODEL_EXT):
+      continue
+
+    # Copy files directly
+    if os.path.isfile(src_item):
+      shutil.copy2(src_item, dest_item)
+
+    # Copy directories using copytree
+    elif os.path.isdir(src_item):
+      shutil.copytree(src_item, dest_item, symlinks=False, ignore=None, dirs_exist_ok=True)
+
+
+def zip_dir(input: Union[Path, str], zip_filename: Union[Path, str]):
+  """
+  Zip folder without compressing
+  """
+  # Convert to Path object
+  dir = Path(input)
+
+  with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_STORED) as zip_file:
+    for entry in dir.rglob("*"):
+      zip_file.write(entry, entry.relative_to(dir))
+
+
+class RepositoryBuilder:
+
+  @staticmethod
+  def init_repository(model_type: str, working_dir: str, backend=Literal['triton'], **kwargs):
+    assert model_type in MODEL_TYPES
+    model_type = model_type
+    default_model_type_config: ModelConfigClass = get_model_config(model_type)
+
+    os.makedirs(working_dir, exist_ok=True)
+
+    def __write_to(filename, data):
+      with open(os.path.join(working_dir, filename), "w") as f:
+        f.write(data)
+
+    # create inference.py
+    _filename = "inference.py"
+    inference_py = _read_static_file(_filename)
+    inference_py = inference_py.replace("InferenceModel()",
+                                        f"InferenceModel({_TYPE_TO_CLASS[model_type]})")
+    inference_py = inference_py.replace("predict_docstring",
+                                        eval(_TYPE_TO_CLASS[model_type]).predict.__doc__)
+    # create config
+    config = asdict(default_model_type_config)
+    if backend == "triton":
+      max_batch_size = kwargs.get("max_batch_size", None)
+      image_shape = kwargs.get("image_shape", None)
+      if max_batch_size:
+        config['serving_backend']['triton']['max_batch_size'] = max_batch_size
+      if image_shape:
+        config['serving_backend']['triton']['image_shape'] = image_shape
+      config = parse_config(config).dump_to_user_config()
+      config_data = yaml.dump(config)
+      sample_yaml = _read_static_file("sample_clarifai_config.yaml")
+      config_data = sample_yaml + "\n\n" + config_data
+      __write_to("clarifai_config.yaml", config_data)
+      #
+    # create inference.py after checking all configs
+    __write_to(_filename, inference_py)
+    # create test.py
+    __write_to("test.py", _read_static_file("test.py"))
+    # create requirements.txt
+    __write_to("requirements.txt", _read_static_file("_requirements.txt"))
+
+  @staticmethod
+  def build(working_dir: str, output_dir: str = None, name: str = None, backend=Literal['triton']):
+    if not output_dir:
+      output_dir = working_dir
+    else:
+      os.makedirs(output_dir, exist_ok=True)
+
+    temp_folder = os.path.join(working_dir, ".cache")
+    os.makedirs(temp_folder, exist_ok=True)
+
+    user_config_file = os.path.join(working_dir, "clarifai_config.yaml")
+    assert os.path.exists(
+        user_config_file
+    ), f"FileNotFound: please make sure `clarifai_config.yaml` exists in {working_dir}"
+    user_config = load_user_config(user_config_file)
+
+    if backend == "triton":
+      triton_1_ver = os.path.join(temp_folder, "1")
+      os.makedirs(triton_1_ver, exist_ok=True)
+      # check if labels exists
+      for output_config in user_config.serving_backend.triton.output:
+        if output_config.label_filename:
+          user_labels = user_config.clarifai_model.labels
+          assert user_labels, f"Model type `{user_config.clarifai_model.type}` requires labels, "\
+          f"but can not found value of `clarifai_model.labels` in {user_config_file}. Please update this attribute to build the model"
+          with open(os.path.join(temp_folder, "labels.txt"), "w") as f:
+            if not isinstance(user_labels, Iterable):
+              user_labels = [user_labels]
+            f.write("\n".join([str(lb) for lb in user_labels]) + "\n")
+
+      # copy model.py
+      shutil.copy(_get_static_file_path("triton/model.py"), triton_1_ver)
+      # copy requirements.txt
+      shutil.copy(os.path.join(working_dir, "requirements.txt"), temp_folder)
+      # copy all other files
+      copy_folder(
+          working_dir, triton_1_ver, exclude_items=["requirements.txt", ".cache", "__pycache__"])
+      # generate config.pbtxt
+      _config_pbtxt_serializer = Serializer(user_config.serving_backend.triton)
+      _config_pbtxt_serializer.to_file(temp_folder)
+
+    else:
+      raise ValueError(f"backend must be ['triton'], got {backend}")
+
+    clarifai_model_name = name or user_config.clarifai_model.clarifai_model_id or "model"
+    clarifai_model_name += BUILT_MODEL_EXT
+    clarifai_model_name = os.path.join(output_dir, clarifai_model_name)
+
+    print(
+        "Model building in progress; the duration may vary depending on the size of checkpoints/assets..."
+    )
+    zip_dir(temp_folder, clarifai_model_name)
+    print(f"Finished. Your model is located at {clarifai_model_name}")
+
+    return clarifai_model_name

clarifai/models/model_serving/repo_build/static_files/_requirements.txt

@@ -0,0 +1,2 @@
+clarifai>9.10.4
+tritonclient[all]

clarifai/models/model_serving/repo_build/static_files/base_test.py

@@ -0,0 +1,169 @@
+import os
+from copy import deepcopy
+from typing import Dict, Iterable, List, Union
+
+import numpy as np
+import yaml
+
+from ...constants import IMAGE_TENSOR_NAME, TEXT_TENSOR_NAME
+from ...model_config import (ClassifierOutput, EmbeddingOutput, ImageOutput, InferParam,
+                             InferParamManager, MasksOutput, ModelTypes, TextOutput,
+                             VisualDetector, load_user_config)
+
+_default_texts = ["Photo of a cat", "A cat is playing around", "Hello, this is test"]
+
+_default_images = [
+    np.zeros((100, 100, 3), dtype='uint8'),  #black
+    np.ones((100, 100, 3), dtype='uint8') * 255,  #white
+    np.random.uniform(0, 255, (100, 100, 3)).astype('uint8')  #noise
+]
+
+
+def _is_valid_logit(x: np.array):
+  return np.all(0 <= x) and np.all(x <= 1)
+
+
+def _is_non_negative(x: np.array):
+  return np.all(x >= 0)
+
+
+def _is_integer(x):
+  return np.all(np.equal(np.mod(x, 1), 0))
+
+
+class BaseTest:
+  init_inference_parameters = {}
+
+  def __init__(self, init_inference_parameters={}) -> None:
+    import sys
+    if 'inference' in sys.modules:
+      del sys.modules['inference']
+    import inference
+    from inference import InferenceModel
+    self.model = InferenceModel()
+    self._base_dir = os.path.dirname(inference.__file__)
+    self.cfg_path = os.path.join(self._base_dir, "clarifai_config.yaml")
+    self.user_config = load_user_config(self.cfg_path)
+    self._user_labels = None
+    # check if labels exists
+    for output_config in self.user_config.serving_backend.triton.output:
+      if output_config.label_filename:
+        self._user_labels = self.user_config.clarifai_model.labels
+        assert self._user_labels, f"Model type `{self.user_config.clarifai_model.type}` requires labels, "\
+        f"but can not found value of `clarifai_model.labels` in {self.cfg_path}. Please update this attribute to build the model"
+
+    # update init vs user_defined params
+    user_defined_infer_params = [
+        InferParam(**each) for each in self.user_config.clarifai_model.inference_parameters
+    ]
+    total_infer_params = []
+    if init_inference_parameters:
+      self.init_inference_parameters = init_inference_parameters
+    for k, v in self.init_inference_parameters.items():
+      _exist = False
+      for user_param in user_defined_infer_params:
+        if user_param.path == k:
+          if user_param.default_value != v:
+            print(f"Warning: Overwrite parameter `{k}` with default value `{v}`")
+          user_param.default_value = v
+          _exist = True
+          total_infer_params.append(user_param)
+          user_defined_infer_params.remove(user_param)
+          break
+      if not _exist:
+        total_infer_params.append(InferParamManager.from_kwargs(**{k: v}).params[0])
+
+    self.infer_param_manager = InferParamManager(
+        params=total_infer_params + user_defined_infer_params)
+    self.user_config.clarifai_model.inference_parameters = self.infer_param_manager.get_list_params(
+    )
+    self._overwrite_cfg()
+
+  @property
+  def user_labels(self):
+    return self._user_labels
+
+  def _overwrite_cfg(self):
+    config = yaml.dump(self.user_config.dump_to_user_config(),)
+    with open(self.cfg_path, "w") as f:
+      f.write(config)
+
+  def predict(self, input_data: Union[List[np.ndarray], List[str], Dict[str, Union[List[
+      np.ndarray], List[str]]]], **inference_parameters) -> Iterable:
+    """
+    Test Prediction method is exact `InferenceModel.predict` method with
+    checking inference paramters.
+
+    Args:
+    -----
+    - input_data: A list of input data item to predict on. The type depends on model input type:
+      * `image`: List[np.ndarray]
+      * `text`: List[str]
+      * `multimodal`:
+        input_data is list of dict where key is input type name e.i. `image`, `text` and value is list.
+        {"image": List[np.ndarray], "text": List[str]}
+
+    - **inference_parameters: keyword args of your inference parameters.
+
+    Returns:
+    --------
+      List of your inference model output type
+    """
+    infer_params = self.infer_param_manager.validate(**inference_parameters)
+    outputs = self.model.predict(input_data=input_data, inference_parameters=infer_params)
+    outputs = self._verify_outputs(outputs)
+    return outputs
+
+  def _verify_outputs(self, outputs: List[Union[ClassifierOutput, VisualDetector, EmbeddingOutput,
+                                                TextOutput, ImageOutput, MasksOutput]]):
+    """Test output value/dims
+
+    Args:
+      outputs (List[Union[ClassifierOutput, VisualDetector, EmbeddingOutput, TextOutput, ImageOutput, MasksOutput]]): Outputs of `predict` method
+    """
+    _outputs = deepcopy(outputs)
+    _output = _outputs[0]
+
+    if isinstance(_output, EmbeddingOutput):
+      # not test
+      pass
+    elif isinstance(_output, ClassifierOutput):
+      for each in _outputs:
+        assert _is_valid_logit(each.predicted_scores), "`predicted_scores` must be in range [0, 1]"
+        assert len(each.predicted_scores) == len(
+            self.user_labels
+        ), f"`predicted_scores` dim must be equal to labels, got {len(each.predicted_scores)} != labels {len(self.user_labels)}"
+    elif isinstance(_output, VisualDetector):
+      for each in _outputs:
+        assert _is_valid_logit(each.predicted_scores), "`predicted_scores` must be in range [0, 1]"
+        assert _is_integer(each.predicted_labels), "`predicted_labels` must be integer"
+        assert np.all(0 <= each.predicted_labels) and np.all(each.predicted_labels < len(
+            self.user_labels)), f"`predicted_labels` must be in [0, {len(self.user_labels) - 1}]"
+        assert _is_non_negative(each.predicted_bboxes), "`predicted_bboxes` must be >= 0"
+    elif isinstance(_output, MasksOutput):
+      for each in _outputs:
+        assert np.all(0 <= each.predicted_mask) and np.all(each.predicted_mask < len(
+            self.user_labels)), f"`predicted_mask` must be in [0, {len(self.user_labels) - 1}]"
+    elif isinstance(_output, TextOutput):
+      pass
+    elif isinstance(_output, ImageOutput):
+      for each in _outputs:
+        assert _is_non_negative(each.image), "`image` must be >= 0"
+    else:
+      pass
+
+    return outputs
+
+  def test_with_default_inputs(self):
+    model_type = self.user_config.clarifai_model.type
+    if model_type == ModelTypes.multimodal_embedder:
+      self.predict(input_data=[{IMAGE_TENSOR_NAME: each} for each in _default_images])
+      self.predict(input_data=[{TEXT_TENSOR_NAME: each} for each in _default_texts])
+      self.predict(input_data=[{
+          TEXT_TENSOR_NAME: text,
+          IMAGE_TENSOR_NAME: img
+      } for text, img in zip(_default_texts, _default_images)])
+    elif model_type.startswith("visual"):
+      self.predict(input_data=_default_images)
+    else:
+      self.predict(input_data=_default_texts)

clarifai/models/model_serving/repo_build/static_files/inference.py

@@ -0,0 +1,26 @@
+# User model inference script.
+
+import os
+from pathlib import Path
+from typing import Dict, Union
+from clarifai.models.model_serving.model_config import *  # noqa
+
+
+class InferenceModel():
+  """User model inference class."""
+
+  def __init__(self) -> None:
+    """
+    Load inference time artifacts that are called frequently .e.g. models, tokenizers, etc.
+    in this method so they are loaded only once for faster inference.
+    """
+    # current directory
+    self.base_path: Path = os.path.dirname(__file__)
+
+  def predict(self,
+              input_data: list,
+              inference_parameters: Dict[str, Union[bool, str, float, int]] = {}) -> list:
+    """predict_docstring
+    """
+
+    raise NotImplementedError()

clarifai/models/model_serving/repo_build/static_files/sample_clarifai_config.yaml

@@ -0,0 +1,25 @@
+# Sample config of inference_parameters and labels
+# For detail, please refer to docs
+# --------------------
+# inference_parameters:
+#   - path: boolean_var
+#     default_value: true
+#     field_type: 1
+#     description: a boolean variable
+#   - path: string_var
+#     default_value: "a string"
+#     field_type: 2
+#     description: a string variable
+#   - path: number_var
+#     default_value: 1
+#     field_type: 3
+#     description: a number variable
+#   - path: secret_string_var
+#     default_value: "YOUR_SECRET"
+#     field_type: 21
+#     description: a string variable contains secret like API key
+# labels:
+#   - concept1
+#   - concept2
+#   - concept3
+#   - concept4

clarifai/models/model_serving/repo_build/static_files/test.py

@@ -0,0 +1,40 @@
+import unittest
+
+from clarifai.models.model_serving.repo_build import BaseTest
+
+
+class CustomTest(unittest.TestCase):
+  """
+  BaseTest loads the InferenceModel from the inference.py file in the current working directory.
+  To execute the predict method of the InferenceModel, use the predict method in BaseTest.
+  It takes the exact same inputs and inference parameters, returning the same outputs as InferenceModel.predict.
+  The difference is that BaseTest.predict verifies your_infer_parameters against config.clarifai_models.inference_parameters and checks the output values.
+
+  For example, test input value of visual-classifier
+
+  def test_input(self):
+    import cv2
+    path = "path/to/image"
+    img = cv2.imread(path)
+    outputs = self.model.predict([img], infer_param1=..., infer_param2=...)
+    print(outputs)
+    assert outputs
+
+  """
+
+  def setUp(self) -> None:
+    your_infer_parameter = dict(
+    )  # for example dict(float_var=0.12, string_var="test", _secret_string_var="secret")
+    self.model = BaseTest(your_infer_parameter)
+
+  def test_default_cases(self):
+    """Test your model with dummy inputs.
+    In general, you only need to run this test to check your InferneceModel implementation.
+    In case the default inputs makes your model failed for some reason (not because of assert in `test_with_default_inputs`),
+    you can comment out this test.
+    """
+    self.model.test_with_default_inputs()
+
+  def test_specific_case1(self):
+    """ Implement your test case"""
+    pass

clarifai/models/model_serving/{models/pb_model.py → repo_build/static_files/triton/model.py}

@@ -19,8 +19,6 @@ try:
   import triton_python_backend_utils as pb_utils
 except ModuleNotFoundError:
   pass
-from google.protobuf import text_format
-from tritonclient.grpc.model_config_pb2 import ModelConfig
 from clarifai.models.model_serving.model_config.inference_parameter import parse_req_parameters
 
 
@@ -33,18 +31,13 @@ class TritonPythonModel:
     """
     Triton server init.
     """
-    args["model_repository"] = args["model_repository"].replace("/1/model.py", "")
     sys.path.append(os.path.dirname(__file__))
     from inference import InferenceModel
 
     self.inference_obj = InferenceModel()
 
     # Read input_name from config file
-    self.config_msg = ModelConfig()
-    with open(os.path.join(args["model_repository"], "config.pbtxt"), "r") as f:
-      cfg = f.read()
-    text_format.Merge(cfg, self.config_msg)
-    self.input_names = [inp.name for inp in self.config_msg.input]
+    self.input_names = [inp.name for inp in self.inference_obj.config.serving_backend.triton.input]
 
   def execute(self, requests):
     """
@@ -53,22 +46,30 @@ class TritonPythonModel:
     responses = []
 
     for request in requests:
-      parameters = request.parameters()
+      try:
+        parameters = request.parameters()
+      except Exception:
+        print(
+            "It seems this triton version does not support `parameters()` in request. "
+            "Please upgrade tritonserver version otherwise can not use `inference_parameters`. Error message: {e}"
+        )
+        parameters = None
+
       parameters = parse_req_parameters(parameters) if parameters else {}
 
       if len(self.input_names) == 1:
         in_batch = pb_utils.get_input_tensor_by_name(request, self.input_names[0])
         in_batch = in_batch.as_numpy()
-        inference_response = self.inference_obj.get_predictions(in_batch, **parameters)
+        data = in_batch
       else:
-        multi_in_batch_dict = {}
+        data = {}
         for input_name in self.input_names:
           in_batch = pb_utils.get_input_tensor_by_name(request, input_name)
           in_batch = in_batch.as_numpy() if in_batch is not None else []
-          multi_in_batch_dict.update({input_name: in_batch})
-
-        inference_response = self.inference_obj.get_predictions(multi_in_batch_dict, **parameters)
+          data.update({input_name: in_batch})
 
+      inference_response = self.inference_obj._tritonserver_predict(
+          input_data=data, inference_parameters=parameters)
       responses.append(inference_response)
 
     return responses

clarifai/models/model_serving/utils.py

@@ -0,0 +1,21 @@
+import os
+
+from .constants import CLARIFAI_PAT_PATH
+
+
+def _persist_pat(pat: str):
+  """ Write down pat to CLARIFAI_PAT_PATH """
+  with open(CLARIFAI_PAT_PATH, "w") as f:
+    f.write(pat)
+
+
+def _read_pat():
+  if not os.path.exists(CLARIFAI_PAT_PATH):
+    return None
+  with open(CLARIFAI_PAT_PATH, "r") as f:
+    return f.read().replace("\n", "").replace("\r", "").strip()
+
+
+def login(pat=None):
+  """ if pat provided, set pat to CLARIFAI_PAT otherwise read pat from file"""
+  os.environ["CLARIFAI_PAT"] = pat or _read_pat()

clarifai/rag/rag.py

@@ -17,6 +17,8 @@ from clarifai.rag.utils import (convert_messages_to_str, format_assistant_messag
                                 split_document)
 from clarifai.utils.logging import get_logger
 
+DEFAULT_RAG_PROMPT_TEMPLATE = "Context information is below:\n{data.hits}\nGiven the context information and not prior knowledge, answer the query.\nQuery: {data.text.raw}\nAnswer: "
+
 
 class RAG:
   """
@@ -24,7 +26,8 @@ class RAG:
 
     Example:
         >>> from clarifai.rag import RAG
-        >>> rag_agent = RAG()
+        >>> rag_agent = RAG(workflow_url=YOUR_WORKFLOW_URL)
+        >>> rag_agent.chat(messages=[{"role":"human", "content":"What is Clarifai"}])
     """
   chat_state_id = None
 
@@ -49,40 +52,69 @@ class RAG:
   @classmethod
   def setup(cls,
             user_id: str = None,
+            app_url: str = None,
             llm_url: str = "https://clarifai.com/mistralai/completion/models/mistral-7B-Instruct",
             base_workflow: str = "Text",
             workflow_yaml_filename: str = 'prompter_wf.yaml',
+            workflow_id: str = None,
             base_url: str = "https://api.clarifai.com",
             pat: str = None,
             **kwargs):
     """Creates an app with `Text` as base workflow, create prompt model, create prompt workflow.
 
+    **kwargs: Additional keyword arguments to be passed to rag-promter model.
+          - min_score (float): The minimum score for search hits.
+          - max_results (float): The maximum number of search hits.
+          - prompt_template (str): The prompt template used. Must contain {data.hits} for the search hits and {data.text.raw} for the query string.
+
     Example:
         >>> from clarifai.rag import RAG
-        >>> rag_agent = RAG.setup()
+        >>> rag_agent = RAG.setup(user_id=YOUR_USER_ID)
+        >>> rag_agent.chat(messages=[{"role":"human", "content":"What is Clarifai"}])
+
+    Or if you already have an existing app with ingested data:
+        >>> rag_agent = RAG.setup(app_url=YOUR_APP_URL)
+        >>> rag_agent.chat(messages=[{"role":"human", "content":"What is Clarifai"}])
     """
-    user = User(user_id=user_id, base_url=base_url, pat=pat)
+
+    if user_id and not app_url:
+      user = User(user_id=user_id, base_url=base_url, pat=pat)
+      ## Create an App
+      now_ts = str(int(datetime.now().timestamp()))
+      app_id = f"rag_app_{now_ts}"
+      app = user.create_app(app_id=app_id, base_workflow=base_workflow)
+
+    if not user_id and app_url:
+      app = App(url=app_url, pat=pat)
+
+    if user_id and app_url:
+      raise UserError("Must provide one of user_id or app_url, not both.")
+
+    if not user_id and not app_url:
+      raise UserError(
+          "user_id or app_url must be provided. The user_id can be found at https://clarifai.com/settings."
+      )
+
     llm = Model(llm_url)
 
+    min_score = kwargs.get("min_score", 0.95)
+    max_results = kwargs.get("max_results", 5)
+    prompt_template = kwargs.get("prompt_template", DEFAULT_RAG_PROMPT_TEMPLATE)
     params = Struct()
     params.update({
-        "prompt_template":
-            "Context information is below:\n{data.hits}\nGiven the context information and not prior knowledge, answer the query.\nQuery: {data.text.raw}\nAnswer: "
+        "min_score": min_score,
+        "max_results": max_results,
+        "prompt_template": prompt_template
     })
     prompter_model_params = {"params": params}
 
-    ## Create an App
-    now_ts = str(int(datetime.now().timestamp()))
-    app_id = f"rag_app_{now_ts}"
-    app = user.create_app(app_id=app_id, base_workflow=base_workflow)
-
     ## Create rag-prompter model and version
     prompter_model = app.create_model(
         model_id=f"rag_prompter_{now_ts}", model_type_id="rag-prompter")
     prompter_model = prompter_model.create_version(output_info=prompter_model_params)
 
     ## Generate a tmp yaml file for workflow creation
-    workflow_id = f"rag-wf-{now_ts}"
+    workflow_id = f"rag-wf-{now_ts}" if workflow_id is None else workflow_id
     workflow_dict = {
         "workflow": {
             "id":
@@ -138,9 +170,10 @@ class RAG:
 
     Example:
         >>> from clarifai.rag import RAG
-        >>> rag_agent = RAG.setup()
+        >>> rag_agent = RAG.setup(user_id=YOUR_USER_ID)
         >>> rag_agent.upload(folder_path = "~/work/docs")
         >>> rag_agent.upload(file_path = "~/work/docs/manual.pdf")
+        >>> rag_agent.chat(messages=[{"role":"human", "content":"What is Clarifai"}])
     """
     #set batch size
     if batch_size > MAX_UPLOAD_BATCH_SIZE:

clarifai/rag/utils.py

@@ -3,8 +3,9 @@ from pathlib import Path
 from typing import List
 
 import requests
-from llama_index import Document, SimpleDirectoryReader, download_loader
-from llama_index.node_parser.text import SentenceSplitter
+from llama_index.core import Document, SimpleDirectoryReader
+from llama_index.core.node_parser.text import SentenceSplitter
+from llama_index.core.readers.download import download_loader
 from pypdf import PdfReader