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

clarifai/rag/utils.py

@@ -7,7 +7,7 @@ import requests
 
 ## TODO: Make this token-aware.
 def convert_messages_to_str(messages: List[dict]) -> str:
-  """convert messages in OpenAI API format into a single string.
+    """convert messages in OpenAI API format into a single string.
 
     Args:
         messages List[dict]: A list of dictionary in the following format:
@@ -19,100 +19,107 @@ def convert_messages_to_str(messages: List[dict]) -> str:
         ]
         ```
     """
-  final_str = ""
-  for msg in messages:
-    if "role" in msg and "content" in msg:
-      role = msg.get("role", "")
-      content = msg.get("content", "")
-      final_str += f"\n\n{role}: {content}"
-  return final_str
+    final_str = ""
+    for msg in messages:
+        if "role" in msg and "content" in msg:
+            role = msg.get("role", "")
+            content = msg.get("content", "")
+            final_str += f"\n\n{role}: {content}"
+    return final_str
 
 
 def format_assistant_message(raw_text: str) -> dict:
-  return {"role": "assistant", "content": raw_text}
+    return {"role": "assistant", "content": raw_text}
 
 
 def load_documents(file_path: str = None, folder_path: str = None, url: str = None) -> List[any]:
-  """Loads documents from a local directory or public url or local filename.
+    """Loads documents from a local directory or public url or local filename.
 
-  Args:
-      file_path (str): The path to the filename.
-      folder_path (str): The path to the folder.
-      url (str): The url to the file.
-  """
-  #check import packages
-  try:
-    from llama_index.core import Document, SimpleDirectoryReader
-    from llama_index.core.readers.download import download_loader
-  except ImportError:
-    raise ImportError("Could not import llama index package. "
-                      "Please install it with `pip install llama-index-core==0.10.1`.")
-  #document loaders for filepath
-  if file_path:
-    if file_path.endswith(".pdf"):
-      PDFReader = download_loader("PDFReader")
-      loader = PDFReader()
-      documents = loader.load_data(file=Path(file_path))
-    elif file_path.endswith(".docx"):
-      docReader = download_loader("DocxReader")
-      loader = docReader()
-      documents = loader.load_data(file=Path(file_path))
-    elif file_path.endswith(".txt"):
-      with open(file_path, 'r') as file:
-        text_content = file.read()
-      documents = [Document(text=text_content)]
-    else:
-      raise ValueError("Only .pdf, .docx, and .txt files are supported.")
+    Args:
+        file_path (str): The path to the filename.
+        folder_path (str): The path to the folder.
+        url (str): The url to the file.
+    """
+    # check import packages
+    try:
+        from llama_index.core import Document, SimpleDirectoryReader
+        from llama_index.core.readers.download import download_loader
+    except ImportError:
+        raise ImportError(
+            "Could not import llama index package. "
+            "Please install it with `pip install llama-index-core==0.10.1`."
+        )
+    # document loaders for filepath
+    if file_path:
+        if file_path.endswith(".pdf"):
+            PDFReader = download_loader("PDFReader")
+            loader = PDFReader()
+            documents = loader.load_data(file=Path(file_path))
+        elif file_path.endswith(".docx"):
+            docReader = download_loader("DocxReader")
+            loader = docReader()
+            documents = loader.load_data(file=Path(file_path))
+        elif file_path.endswith(".txt"):
+            with open(file_path, 'r') as file:
+                text_content = file.read()
+            documents = [Document(text=text_content)]
+        else:
+            raise ValueError("Only .pdf, .docx, and .txt files are supported.")
 
-  #document loaders for folderpath
-  if folder_path:
-    documents = SimpleDirectoryReader(
-        input_dir=Path(folder_path), required_exts=[".pdf", ".docx", ".txt"]).load_data()
+    # document loaders for folderpath
+    if folder_path:
+        documents = SimpleDirectoryReader(
+            input_dir=Path(folder_path), required_exts=[".pdf", ".docx", ".txt"]
+        ).load_data()
 
-  #document loaders for url
-  if url:
-    response = requests.get(url)
-    if response.status_code != 200:
-      raise ValueError(f"Invalid url {url}.")
-    #for text files
-    try:
-      documents = [Document(text=response.content)]
-    #for pdf files
-    except Exception:
-      #check import packages
-      try:
-        from pypdf import PdfReader
-      except ImportError:
-        raise ImportError("Could not import pypdf package. "
-                          "Please install it with `pip install pypdf==3.17.4`.")
-      documents = []
-      pdf_file = PdfReader(io.BytesIO(response.content))
-      num_pages = len(pdf_file.pages)
-      for page in range(num_pages):
-        page_text = pdf_file.pages[page].extract_text()
-        documents.append(Document(text=page_text))
-    else:
-      raise ValueError(f"Invalid url {url}.")
+    # document loaders for url
+    if url:
+        response = requests.get(url)
+        if response.status_code != 200:
+            raise ValueError(f"Invalid url {url}.")
+        # for text files
+        try:
+            documents = [Document(text=response.content)]
+        # for pdf files
+        except Exception:
+            # check import packages
+            try:
+                from pypdf import PdfReader
+            except ImportError:
+                raise ImportError(
+                    "Could not import pypdf package. "
+                    "Please install it with `pip install pypdf==3.17.4`."
+                )
+            documents = []
+            pdf_file = PdfReader(io.BytesIO(response.content))
+            num_pages = len(pdf_file.pages)
+            for page in range(num_pages):
+                page_text = pdf_file.pages[page].extract_text()
+                documents.append(Document(text=page_text))
+        else:
+            raise ValueError(f"Invalid url {url}.")
 
-  return documents
+    return documents
 
 
 def split_document(text: str, chunk_size: int, chunk_overlap: int, **kwargs) -> List[str]:
-  """Splits a document into chunks of text.
+    """Splits a document into chunks of text.
 
-  Args:
-      text (str): The text to split.
-      chunk_size (int): The size of each chunk.
-      chunk_overlap (int): The amount of overlap between each chunk.
-      **kwargs: Additional keyword arguments for the SentenceSplitter.
-  """
-  #check import packages
-  try:
-    from llama_index.core.node_parser.text import SentenceSplitter
-  except ImportError:
-    raise ImportError("Could not import llama index package. "
-                      "Please install it with `pip install llama-index-core==0.10.24`.")
-  #document
-  text_parser = SentenceSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap, **kwargs)
-  text_chunks = text_parser.split_text(text)
-  return text_chunks
+    Args:
+        text (str): The text to split.
+        chunk_size (int): The size of each chunk.
+        chunk_overlap (int): The amount of overlap between each chunk.
+        **kwargs: Additional keyword arguments for the SentenceSplitter.
+    """
+    # check import packages
+    try:
+        from llama_index.core.node_parser.text import SentenceSplitter
+    except ImportError:
+        raise ImportError(
+            "Could not import llama index package. "
+            "Please install it with `pip install llama-index-core==0.10.24`."
+        )
+    # document
+    text_parser = SentenceSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap, **kwargs)
+    text_chunks = text_parser.split_text(text)
+    return text_chunks