clarifai 11.3.0rc2__py3-none-any.whl → 11.4.0__py3-none-any.whl

clarifai/rag/rag.py

@@ -12,8 +12,12 @@ from clarifai.client.user import User
 from clarifai.client.workflow import Workflow
 from clarifai.constants.rag import MAX_UPLOAD_BATCH_SIZE
 from clarifai.errors import UserError
-from clarifai.rag.utils import (convert_messages_to_str, format_assistant_message, load_documents,
-                                split_document)
+from clarifai.rag.utils import (
+    convert_messages_to_str,
+    format_assistant_message,
+    load_documents,
+    split_document,
+)
 from clarifai.utils.constants import CLARIFAI_USER_ID_ENV_VAR
 from clarifai.utils.logging import logger
 from clarifai.utils.misc import get_from_dict_or_env
@@ -22,7 +26,7 @@ DEFAULT_RAG_PROMPT_TEMPLATE = "Context information is below:\n{data.hits}\nGiven
 
 
 class RAG:
-  """
+    """
     RAG is a class for Retrieval Augmented Generation.
 
     Example:
@@ -30,290 +34,310 @@ class RAG:
         >>> rag_agent = RAG(workflow_url=YOUR_WORKFLOW_URL)
         >>> rag_agent.chat(messages=[{"role":"human", "content":"What is Clarifai"}])
     """
-  chat_state_id = None
-
-  def __init__(self,
-               workflow_url: str = None,
-               workflow: Workflow = None,
-               base_url: str = "https://api.clarifai.com",
-               pat: str = None,
-               **kwargs):
-    """Initialize an empty or existing RAG.
-    """
-    self.logger = logger
-    if workflow_url is not None and workflow is None:
-      self.logger.info("workflow_url:%s", workflow_url)
-      w = Workflow(workflow_url, base_url=base_url, pat=pat)
-      self._prompt_workflow = w
-      self._app = App(app_id=w.app_id, user_id=w.user_id, base_url=w.base, pat=w.pat)
-    elif workflow_url is None and workflow is not None:
-      self._prompt_workflow = workflow
-      self._app = App(
-          app_id=workflow.app_id,
-          user_id=workflow.user_id,
-          base_url=workflow.base,
-          pat=workflow.pat)
-
-  @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(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"}])
-    """
-    if not app_url:
-      try:
-        user_id = get_from_dict_or_env(key="user_id", env_key=CLARIFAI_USER_ID_ENV_VAR, **kwargs)
-      except Exception:
-        pass
-
-    now_ts = uuid.uuid4().hex[:10]
-    if user_id and not app_url:
-      user = User(user_id=user_id, base_url=base_url, pat=pat)
-      ## Create an App
-      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)
-      uid = app_url.split(".com/")[1].split("/")[0]
-      user = User(user_id=uid, base_url=base_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(url=llm_url, pat=pat)
-
-    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({
-        "min_score": min_score,
-        "max_results": max_results,
-        "prompt_template": prompt_template
-    })
-    prompter_model_params = {"params": params}
-
-    ## Create rag-prompter model and version
-    model_id = f"prompter-{workflow_id}-{now_ts}" if workflow_id is not None else f"rag-prompter-{now_ts}"
-    prompter_model = app.create_model(model_id=model_id, 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}" if workflow_id is None else workflow_id
-    workflow_dict = {
-        "workflow": {
-            "id":
-                workflow_id,
-            "nodes": [{
-                "id": "rag-prompter",
-                "model": {
-                    "model_id": prompter_model.id,
-                    "model_version_id": prompter_model.model_version.id
-                }
-            }, {
-                "id": "llm",
-                "model": {
-                    "model_id": llm.id,
-                    "user_id": llm.user_id,
-                    "app_id": llm.app_id
-                },
-                "node_inputs": [{
-                    "node_id": "rag-prompter"
-                }]
-            }]
+    chat_state_id = None
+
+    def __init__(
+        self,
+        workflow_url: str = None,
+        workflow: Workflow = None,
+        base_url: str = "https://api.clarifai.com",
+        pat: str = None,
+        **kwargs,
+    ):
+        """Initialize an empty or existing RAG."""
+        self.logger = logger
+        if workflow_url is not None and workflow is None:
+            self.logger.info("workflow_url:%s", workflow_url)
+            w = Workflow(workflow_url, base_url=base_url, pat=pat)
+            self._prompt_workflow = w
+            self._app = App(app_id=w.app_id, user_id=w.user_id, base_url=w.base, pat=w.pat)
+        elif workflow_url is None and workflow is not None:
+            self._prompt_workflow = workflow
+            self._app = App(
+                app_id=workflow.app_id,
+                user_id=workflow.user_id,
+                base_url=workflow.base,
+                pat=workflow.pat,
+            )
+
+    @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(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"}])
+        """
+        if not app_url:
+            try:
+                user_id = get_from_dict_or_env(
+                    key="user_id", env_key=CLARIFAI_USER_ID_ENV_VAR, **kwargs
+                )
+            except Exception:
+                pass
+
+        now_ts = uuid.uuid4().hex[:10]
+        if user_id and not app_url:
+            user = User(user_id=user_id, base_url=base_url, pat=pat)
+            ## Create an App
+            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)
+            uid = app_url.split(".com/")[1].split("/")[0]
+            user = User(user_id=uid, base_url=base_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(url=llm_url, pat=pat)
+
+        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(
+            {
+                "min_score": min_score,
+                "max_results": max_results,
+                "prompt_template": prompt_template,
+            }
+        )
+        prompter_model_params = {"params": params}
+
+        ## Create rag-prompter model and version
+        model_id = (
+            f"prompter-{workflow_id}-{now_ts}"
+            if workflow_id is not None
+            else f"rag-prompter-{now_ts}"
+        )
+        prompter_model = app.create_model(model_id=model_id, 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}" if workflow_id is None else workflow_id
+        workflow_dict = {
+            "workflow": {
+                "id": workflow_id,
+                "nodes": [
+                    {
+                        "id": "rag-prompter",
+                        "model": {
+                            "model_id": prompter_model.id,
+                            "model_version_id": prompter_model.model_version.id,
+                        },
+                    },
+                    {
+                        "id": "llm",
+                        "model": {
+                            "model_id": llm.id,
+                            "user_id": llm.user_id,
+                            "app_id": llm.app_id,
+                        },
+                        "node_inputs": [{"node_id": "rag-prompter"}],
+                    },
+                ],
+            }
         }
-    }
-    with open(workflow_yaml_filename, 'w') as out_file:
-      yaml.dump(workflow_dict, out_file, default_flow_style=False)
-
-    ## Create prompt workflow
-    wf = app.create_workflow(config_filepath=workflow_yaml_filename)
-    del user, llm, prompter_model, prompter_model_params
-    return cls(workflow=wf)
-
-  def upload(self,
-             file_path: str = None,
-             folder_path: str = None,
-             url: str = None,
-             batch_size: int = 128,
-             chunk_size: int = 1024,
-             chunk_overlap: int = 200,
-             dataset_id: str = None,
-             metadata: dict = None,
-             **kwargs) -> None:
-    """Uploads documents to the app.
-        - Read from a local directory or public url or local filename.
-        - Parse the document(s) into chunks.
-        - Ingest chunks into the app with metadata.
-
-    Args:
-        file_path str: File path to the document.
-        folder_path str: Folder path to the documents.
-        url str: Public url to the document.
-        batch_size int: Batch size for uploading.
-        chunk_size int: Chunk size for splitting the document.
-        chunk_overlap int: The token overlap of each chunk when splitting.
-        **kwargs: Additional arguments for the SentenceSplitter. Refer https://docs.llamaindex.ai/en/stable/api/llama_index.node_parser.SentenceSplitter.html
-
-    Example:
-        >>> from clarifai.rag import RAG
-        >>> 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:
-      raise ValueError(f"batch_size cannot be greater than {MAX_UPLOAD_BATCH_SIZE}")
-
-    #check if only one of file_path, folder_path, or url is specified
-    if file_path and (folder_path or url):
-      raise ValueError("Only one of file_path, folder_path, or url can be specified.")
-    if folder_path and (file_path or url):
-      raise ValueError("Only one of file_path, folder_path, or url can be specified.")
-    if url and (file_path or folder_path):
-      raise ValueError("Only one of file_path, folder_path, or url can be specified.")
-
-    #loading documents
-    documents = load_documents(file_path=file_path, folder_path=folder_path, url=url)
-
-    #splitting documents into chunks
-    text_chunks = []
-    metadata_list = []
-
-    #iterate through documents
-    for doc in documents:
-      doc_i = 0
-      cur_text_chunks = split_document(
-          text=doc.text, chunk_size=chunk_size, chunk_overlap=chunk_overlap, **kwargs)
-      text_chunks.extend(cur_text_chunks)
-      metadata_list.extend([doc.metadata for _ in range(len(cur_text_chunks))])
-      #if batch size is reached, upload the batch
-      if len(text_chunks) > batch_size:
-        for idx in range(0, len(text_chunks), batch_size):
-          if idx + batch_size > len(text_chunks):
-            continue
-          batch_texts = text_chunks[0:batch_size]
-          batch_ids = [uuid.uuid4().hex for _ in range(batch_size)]
-          #metadata
-          batch_metadatas = metadata_list[0:batch_size]
-          meta_list = []
-          for meta in batch_metadatas:
-            meta_struct = Struct()
-            meta_struct.update(meta)
-            meta_struct.update({"doc_chunk_no": doc_i})
-            if metadata and isinstance(metadata, dict):
-              meta_struct.update(metadata)
-            meta_list.append(meta_struct)
-            doc_i += 1
-          del batch_metadatas
-          #creating input proto
-          input_batch = [
-              self._app.inputs().get_text_input(
-                  input_id=batch_ids[i],
-                  raw_text=text,
-                  dataset_id=dataset_id,
-                  metadata=meta_list[i],
-              ) for i, text in enumerate(batch_texts)
-          ]
-          #uploading input with metadata
-          self._app.inputs().upload_inputs(inputs=input_batch)
-          #delete uploaded chunks
-          del text_chunks[0:batch_size]
-          del metadata_list[0:batch_size]
-
-    #uploading the remaining chunks
-    if len(text_chunks) > 0:
-      batch_size = len(text_chunks)
-      batch_ids = [uuid.uuid4().hex for _ in range(batch_size)]
-      #metadata
-      batch_metadatas = metadata_list[0:batch_size]
-      meta_list = []
-      for meta in batch_metadatas:
-        meta_struct = Struct()
-        meta_struct.update(meta)
-        meta_struct.update({"doc_chunk_no": doc_i})
-        if metadata and isinstance(metadata, dict):
-          meta_struct.update(metadata)
-        meta_list.append(meta_struct)
-        doc_i += 1
-      del batch_metadatas
-      #creating input proto
-      input_batch = [
-          self._app.inputs().get_text_input(
-              input_id=batch_ids[i],
-              raw_text=text,
-              dataset_id=dataset_id,
-              metadata=meta_list[i],
-          ) for i, text in enumerate(text_chunks)
-      ]
-      #uploading input with metadata
-      self._app.inputs().upload_inputs(inputs=input_batch)
-      del text_chunks
-      del metadata_list
-
-  def chat(self, messages: List[dict], client_manage_state: bool = False) -> List[dict]:
-    """Chat interface in OpenAI API format.
-
-    Args:
-        messages List[dict]: A list of dictionary in the following format:
-        ```
-        [
-          {"role": "user", "content": "Hello there."},
-          {"role": "assistant", "content": "Hi, I'm Claude. How can I help you?"},
-          {"role": "user", "content": "Can you explain LLMs in plain English?"},
-        ]
-        ```
-        client_manage_state (bool): Whether the client will handle chat state management. Default is false.
-
-    This will pass back the workflow state ID for the server to store chat state.
-    """
-    if client_manage_state:
-      single_prompt = convert_messages_to_str(messages)
-      input_proto = Inputs._get_proto("", "", text_pb=resources_pb2.Text(raw=single_prompt))
-      response = self._prompt_workflow.predict([input_proto])
-      messages.append(format_assistant_message(response.results[0].outputs[-1].data.text.raw))
-      return messages
-
-    # server-side state management
-    message = messages[-1].get("content", "")
-    if len(message) == 0:
-      raise UserError("Empty message supplied.")
-
-    # get chat state id
-    chat_state_id = "init" if self.chat_state_id is None else self.chat_state_id
-
-    # call predict
-    input_proto = Inputs._get_proto("", "", text_pb=resources_pb2.Text(raw=message))
-    response = self._prompt_workflow.predict([input_proto], workflow_state_id=chat_state_id)
-
-    # store chat state id
-    self.chat_state_id = response.workflow_state.id
-    return [format_assistant_message(response.results[0].outputs[-1].data.text.raw)]
+        with open(workflow_yaml_filename, 'w') as out_file:
+            yaml.dump(workflow_dict, out_file, default_flow_style=False)
+
+        ## Create prompt workflow
+        wf = app.create_workflow(config_filepath=workflow_yaml_filename)
+        del user, llm, prompter_model, prompter_model_params
+        return cls(workflow=wf)
+
+    def upload(
+        self,
+        file_path: str = None,
+        folder_path: str = None,
+        url: str = None,
+        batch_size: int = 128,
+        chunk_size: int = 1024,
+        chunk_overlap: int = 200,
+        dataset_id: str = None,
+        metadata: dict = None,
+        **kwargs,
+    ) -> None:
+        """Uploads documents to the app.
+            - Read from a local directory or public url or local filename.
+            - Parse the document(s) into chunks.
+            - Ingest chunks into the app with metadata.
+
+        Args:
+            file_path str: File path to the document.
+            folder_path str: Folder path to the documents.
+            url str: Public url to the document.
+            batch_size int: Batch size for uploading.
+            chunk_size int: Chunk size for splitting the document.
+            chunk_overlap int: The token overlap of each chunk when splitting.
+            **kwargs: Additional arguments for the SentenceSplitter. Refer https://docs.llamaindex.ai/en/stable/api/llama_index.node_parser.SentenceSplitter.html
+
+        Example:
+            >>> from clarifai.rag import RAG
+            >>> 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:
+            raise ValueError(f"batch_size cannot be greater than {MAX_UPLOAD_BATCH_SIZE}")
+
+        # check if only one of file_path, folder_path, or url is specified
+        if file_path and (folder_path or url):
+            raise ValueError("Only one of file_path, folder_path, or url can be specified.")
+        if folder_path and (file_path or url):
+            raise ValueError("Only one of file_path, folder_path, or url can be specified.")
+        if url and (file_path or folder_path):
+            raise ValueError("Only one of file_path, folder_path, or url can be specified.")
+
+        # loading documents
+        documents = load_documents(file_path=file_path, folder_path=folder_path, url=url)
+
+        # splitting documents into chunks
+        text_chunks = []
+        metadata_list = []
+
+        # iterate through documents
+        for doc in documents:
+            doc_i = 0
+            cur_text_chunks = split_document(
+                text=doc.text, chunk_size=chunk_size, chunk_overlap=chunk_overlap, **kwargs
+            )
+            text_chunks.extend(cur_text_chunks)
+            metadata_list.extend([doc.metadata for _ in range(len(cur_text_chunks))])
+            # if batch size is reached, upload the batch
+            if len(text_chunks) > batch_size:
+                for idx in range(0, len(text_chunks), batch_size):
+                    if idx + batch_size > len(text_chunks):
+                        continue
+                    batch_texts = text_chunks[0:batch_size]
+                    batch_ids = [uuid.uuid4().hex for _ in range(batch_size)]
+                    # metadata
+                    batch_metadatas = metadata_list[0:batch_size]
+                    meta_list = []
+                    for meta in batch_metadatas:
+                        meta_struct = Struct()
+                        meta_struct.update(meta)
+                        meta_struct.update({"doc_chunk_no": doc_i})
+                        if metadata and isinstance(metadata, dict):
+                            meta_struct.update(metadata)
+                        meta_list.append(meta_struct)
+                        doc_i += 1
+                    del batch_metadatas
+                    # creating input proto
+                    input_batch = [
+                        self._app.inputs().get_text_input(
+                            input_id=batch_ids[i],
+                            raw_text=text,
+                            dataset_id=dataset_id,
+                            metadata=meta_list[i],
+                        )
+                        for i, text in enumerate(batch_texts)
+                    ]
+                    # uploading input with metadata
+                    self._app.inputs().upload_inputs(inputs=input_batch)
+                    # delete uploaded chunks
+                    del text_chunks[0:batch_size]
+                    del metadata_list[0:batch_size]
+
+        # uploading the remaining chunks
+        if len(text_chunks) > 0:
+            batch_size = len(text_chunks)
+            batch_ids = [uuid.uuid4().hex for _ in range(batch_size)]
+            # metadata
+            batch_metadatas = metadata_list[0:batch_size]
+            meta_list = []
+            for meta in batch_metadatas:
+                meta_struct = Struct()
+                meta_struct.update(meta)
+                meta_struct.update({"doc_chunk_no": doc_i})
+                if metadata and isinstance(metadata, dict):
+                    meta_struct.update(metadata)
+                meta_list.append(meta_struct)
+                doc_i += 1
+            del batch_metadatas
+            # creating input proto
+            input_batch = [
+                self._app.inputs().get_text_input(
+                    input_id=batch_ids[i],
+                    raw_text=text,
+                    dataset_id=dataset_id,
+                    metadata=meta_list[i],
+                )
+                for i, text in enumerate(text_chunks)
+            ]
+            # uploading input with metadata
+            self._app.inputs().upload_inputs(inputs=input_batch)
+            del text_chunks
+            del metadata_list
+
+    def chat(self, messages: List[dict], client_manage_state: bool = False) -> List[dict]:
+        """Chat interface in OpenAI API format.
+
+        Args:
+            messages List[dict]: A list of dictionary in the following format:
+            ```
+            [
+              {"role": "user", "content": "Hello there."},
+              {"role": "assistant", "content": "Hi, I'm Claude. How can I help you?"},
+              {"role": "user", "content": "Can you explain LLMs in plain English?"},
+            ]
+            ```
+            client_manage_state (bool): Whether the client will handle chat state management. Default is false.
+
+        This will pass back the workflow state ID for the server to store chat state.
+        """
+        if client_manage_state:
+            single_prompt = convert_messages_to_str(messages)
+            input_proto = Inputs._get_proto("", "", text_pb=resources_pb2.Text(raw=single_prompt))
+            response = self._prompt_workflow.predict([input_proto])
+            messages.append(
+                format_assistant_message(response.results[0].outputs[-1].data.text.raw)
+            )
+            return messages
+
+        # server-side state management
+        message = messages[-1].get("content", "")
+        if len(message) == 0:
+            raise UserError("Empty message supplied.")
+
+        # get chat state id
+        chat_state_id = "init" if self.chat_state_id is None else self.chat_state_id
+
+        # call predict
+        input_proto = Inputs._get_proto("", "", text_pb=resources_pb2.Text(raw=message))
+        response = self._prompt_workflow.predict([input_proto], workflow_state_id=chat_state_id)
+
+        # store chat state id
+        self.chat_state_id = response.workflow_state.id
+        return [format_assistant_message(response.results[0].outputs[-1].data.text.raw)]