poma 0.1.4__py3-none-any.whl → 0.2.1__py3-none-any.whl

poma/client.py

@@ -8,18 +8,15 @@ import time
 from pathlib import Path
 from typing import Any
 
-from poma.exceptions import AuthenticationError, InvalidInputError, RemoteServerError, InvalidResponseError
-from poma.retrieval import generate_cheatsheets, generate_single_cheatsheet
+from poma.exceptions import (
+    AuthenticationError,
+    RemoteServerError,
+    InvalidResponseError,
+)
+from poma.retrieval import generate_cheatsheets
 
-USER_AGENT = "poma-ai-sdk/0.1.0"
 
-ALLOWED_FILE_EXTENSIONS: set[str] = {
-    ".txt",
-    ".md",
-    ".html",
-    ".htm",
-    ".pdf",
-}
+USER_AGENT = "poma-ai-sdk/0.1.0"
 
 API_BASE_URL = "https://api.poma-ai.com/api/v1"
 
@@ -51,10 +48,13 @@ class Poma:
         # Override API base URL if environment variable is set
         if os.environ.get("API_BASE_URL"):
             api_base_url = os.environ.get("API_BASE_URL")
+        if not api_base_url:
+            raise ValueError("API base URL cannot be empty.")
 
         self.base_api_url = api_base_url.rstrip("/")
-        self._client = client or httpx.Client(timeout=timeout,
-                                                headers={"user-agent": USER_AGENT})
+        self._client = client or httpx.Client(
+            timeout=timeout, headers={"user-agent": USER_AGENT}
+        )
         if not (api_key := api_key or os.environ.get("POMA_API_KEY", "")):
             raise Exception("POMA_API_KEY environment variable not set.")
         self._client.headers.update({"Authorization": f"Bearer {api_key}"})
@@ -77,11 +77,6 @@ class Poma:
         """
         if not file_path or not isinstance(file_path, os.PathLike):
             raise ValueError("file_path must be a non-empty os.PathLike.")
-        file_extension = Path(file_path).suffix.lower()
-        if file_extension not in ALLOWED_FILE_EXTENSIONS:
-            raise InvalidInputError(
-                f"File extension of {file_path} is not allowed; use one of the following types: {', '.join(sorted(ALLOWED_FILE_EXTENSIONS))}."
-            )
         payload = {}
         if base_url:
             payload["base_url"] = base_url
@@ -97,12 +92,18 @@ class Poma:
         except httpx.HTTPStatusError as error:
             status = error.response.status_code
             if status in (401, 403):
-                raise AuthenticationError(response.text) from error
-            raise RemoteServerError(f"{status}: {response.text}") from error
+                raise AuthenticationError(
+                    f"Failed to submit file '{file_path}': authentication error"
+                ) from error
+            raise RemoteServerError(
+                f"Failed to submit file '{file_path}': {status}"
+            ) from error
         try:
             data = response.json()
         except ValueError as error:
-            raise InvalidResponseError("Server returned non-JSON or empty body") from error
+            raise InvalidResponseError(
+                "Server returned non-JSON or empty body"
+            ) from error
         return data
 
     def get_chunk_result(
@@ -149,19 +150,25 @@ class Poma:
                     download = data.get("download", {})
                     download_url = download.get("download_url", "")
                     if not download_url:
-                        raise RuntimeError("Failed to receive download URL from server.")
+                        raise RuntimeError(
+                            "Failed to receive download URL from server."
+                        )
 
                     if download_dir is None:
                         # Return bytes content instead of saving to file
                         file_bytes = self.download_bytes(download_url)
-                        return self.extract_chunks_and_chunksets_from_poma_archive(poma_archive_data=file_bytes)
+                        return self.extract_chunks_and_chunksets_from_poma_archive(
+                            poma_archive_data=file_bytes
+                        )
                     else:
                         # Save downloaded file to directory
                         filename = download.get("filename", "downloaded_file.poma")
-                        downloaded_file_path = self.download_file(download_url,
-                                                                    filename,
-                                                                    save_directory=download_dir)
-                        return self.extract_chunks_and_chunksets_from_poma_archive(poma_archive_path=downloaded_file_path)
+                        downloaded_file_path = self.download_file(
+                            download_url, filename, save_directory=download_dir
+                        )
+                        return self.extract_chunks_and_chunksets_from_poma_archive(
+                            poma_archive_path=downloaded_file_path
+                        )
                 elif status == "failed":
                     error_code = data.get("code", "unknown")
                     error_details = data.get("error", "No details provided.")
@@ -197,33 +204,29 @@ class Poma:
         Returns:
             dict: A dictionary containing the chunks and chunksets.
         """
-        # Sanity check for parameters
-        if not poma_archive_data and not poma_archive_path:
-            raise ValueError("Either poma_archive_data or poma_archive_path must be provided.")
 
         # Load the chunks and chunksets from POMA archive
         chunks = None
         chunksets = None
         if poma_archive_path:
             with zipfile.ZipFile(poma_archive_path, "r") as zip_ref:
-                chunks = zip_ref.read('chunks.json')
-                chunksets = zip_ref.read('chunksets.json')
-        else:
+                chunks = zip_ref.read("chunks.json")
+                chunksets = zip_ref.read("chunksets.json")
+        elif poma_archive_data:
             with zipfile.ZipFile(io.BytesIO(poma_archive_data), "r") as zip_ref:
-                chunks = zip_ref.read('chunks.json')
-                chunksets = zip_ref.read('chunksets.json')
+                chunks = zip_ref.read("chunks.json")
+                chunksets = zip_ref.read("chunksets.json")
+        else:
+            raise ValueError(
+                "Either poma_archive_data or poma_archive_path must be provided."
+            )
 
         # Sanity check
         if not chunks or not chunksets:
-            raise KeyError(
-                "Result must contain 'chunks' and 'chunksets' keys."
-            )
+            raise KeyError("Result must contain 'chunks' and 'chunksets' keys.")
 
         # Load the chunks and chunksets
-        json_result = {
-            "chunks": json.loads(chunks),
-            "chunksets": json.loads(chunksets)
-        }
+        json_result = {"chunks": json.loads(chunks), "chunksets": json.loads(chunksets)}
         return json_result
 
     def create_cheatsheet(
@@ -234,14 +237,24 @@ class Poma:
         """
         Generates a single cheatsheet for one single document
         from relevant chunksets (relevant for a certain query)
-        and from all available chunks (which must contain the textual content).
+        and from all chunks of that document (providing the textual content).
         Args:
             relevant_chunksets (list[dict]): A list of chunksets, each containing a "chunks" key with a list of chunk IDs.
-            all_chunks (list[dict]): A list of all available chunk dictionaries, each representing a chunk of content.
+            all_chunks (list[dict]): A list of all chunk dictionaries of the same document, each representing a chunk of content.
         Returns:
             str: The textual content of the generated cheatsheet.
         """
-        return generate_single_cheatsheet(relevant_chunksets, all_chunks)
+        cheatsheets = generate_cheatsheets(relevant_chunksets, all_chunks)
+        if (
+            not cheatsheets
+            or not isinstance(cheatsheets, list)
+            or len(cheatsheets) == 0
+            or "content" not in cheatsheets[0]
+        ):
+            raise Exception(
+                "Unknown error; cheatsheet could not be created from input chunks."
+            )
+        return cheatsheets[0]["content"]
 
     def create_cheatsheets(
         self,
@@ -250,14 +263,14 @@ class Poma:
     ) -> list[dict[str, Any]]:
         """
         Generates cheatsheets from relevant chunksets (relevant for a certain query)
-        and from all available chunks (which must contain the textual content).
-        One cheatsheet is created for each document tag found in the chunks.
+        and from all the chunks of all affected documents (providing the textual content).
+        One cheatsheet is created for each document found in the chunks (tagged with file_id).
         Args:
             relevant_chunksets (list[dict]): A list of chunksets, each containing a "chunks" key with a list of chunk IDs.
-            all_chunks (list[dict]): A list of all available chunk dictionaries, each representing a chunk of content.
+            all_chunks (list[dict]): A list of all available chunk dictionaries of affected documents, each representing a chunk of content.
         Returns:
             list[dict]: A list of dictionaries representing the generated cheatsheets, each containing:
-                - 'tag': The tag associated with the respective document.
+                - 'file_id': The tag associated with the respective document.
                 - 'content': The textual content of the generated cheatsheet.
         """
         return generate_cheatsheets(relevant_chunksets, all_chunks)
@@ -283,7 +296,7 @@ class Poma:
         """
         if not download_url:
             raise ValueError("download_url cannot be empty")
-        
+
         # Determine filename
         if not filename:
             filename = Path(download_url).name or "downloaded_file"
@@ -303,7 +316,8 @@ class Poma:
         # Save the file
         with open(save_path, "wb") as f:
             f.write(content)
-        return save_path
+
+        return str(save_path)
 
     def download_bytes(
         self,
@@ -319,27 +333,27 @@ class Poma:
         """
         if not download_url:
             raise ValueError("download_url cannot be empty")
-        
+
         # Construct the full URL if it's a relative path
         if download_url.startswith("/"):
             full_url = f"{self.base_api_url}{download_url}"
         else:
             full_url = download_url
-        
+
         print("Downloading file from:", full_url)
         try:
-            # Download the file
             response = self._client.get(full_url)
             response.raise_for_status()
-            
-            # Return the bytes content
             return response.content
-            
         except httpx.HTTPStatusError as error:
             status = error.response.status_code
             if status in (401, 403):
-                raise AuthenticationError(f"Authentication failed when downloading file: {response.text}") from error
-            raise RemoteServerError(f"Failed to download file: {status} {response.text}") from error
+                raise AuthenticationError(
+                    f"Failed to download '{download_url}': authentication error"
+                ) from error
+            raise RemoteServerError(
+                f"Failed to download '{download_url}': {status}"
+            ) from error
         except Exception as error:
             raise RuntimeError(f"File download failed: {error}") from error
 

poma/exceptions.py

@@ -16,5 +16,6 @@ class RemoteServerError(PomaSDKError):
 class InvalidInputError(PomaSDKError):
     """Raised when an unsupported *Content‑Type* is given to ``chunk_text``."""
 
+
 class InvalidResponseError(PomaSDKError):
     """Raised when the server returns non-JSON or empty body."""

poma/integrations/langchain_poma.py

@@ -18,9 +18,8 @@ from langchain_text_splitters import TextSplitter
 from pydantic import Field, PrivateAttr
 
 from poma import Poma
-from poma.client import ALLOWED_FILE_EXTENSIONS
 from poma.exceptions import InvalidInputError
-from poma.retrieval import _cheatsheets_from_chunks
+from poma.retrieval import chunks_from_dicts, _cheatsheets_from_chunks
 
 __all__ = ["PomaFileLoader", "PomaChunksetSplitter", "PomaCheatsheetRetrieverLC"]
 
@@ -52,10 +51,6 @@ class PomaFileLoader(BaseLoader):
             nonlocal skipped, documents
             if not file_path.is_file():
                 return
-            file_extension = file_path.suffix.lower()
-            if not file_extension or file_extension not in ALLOWED_FILE_EXTENSIONS:
-                skipped += 1
-                return
             file_bytes = file_path.read_bytes()
             file_hash = hashlib.md5(file_bytes).hexdigest()
             if file_path.suffix.lower() == ".pdf":
@@ -84,13 +79,10 @@ class PomaFileLoader(BaseLoader):
         else:
             raise FileNotFoundError(f"Unsupported path type (not file/dir): {path}")
 
-        allowed = ", ".join(sorted(ALLOWED_FILE_EXTENSIONS))
         if not documents:
-            raise InvalidInputError(f"No supported files found. Allowed: {allowed}")
+            raise InvalidInputError(f"No supported files found.")
         if skipped > 0:
-            print(
-                f"Skipped {skipped} file(s) due to unsupported or unreadable type. Allowed: {allowed}"
-            )
+            print(f"Skipped {skipped} file(s) due to unsupported or unreadable type.")
         return documents
 
 
@@ -330,7 +322,7 @@ class PomaCheatsheetRetrieverLC(BaseRetriever):
 
     def _create_cheatsheet_langchain(self, chunked_docs: list[Document]) -> str:
         """Generate a single deduplicated cheatsheet from chunked documents."""
-        all_chunks = []
+        all_chunk_dicts = []
         seen = set()
         for doc in chunked_docs:
             doc_id = doc.metadata.get("doc_id", "unknown_doc")
@@ -341,11 +333,9 @@ class PomaCheatsheetRetrieverLC(BaseRetriever):
                 if chunk_index not in seen:
                     seen.add(chunk_index)
                     chunk["tag"] = doc_id
-                    all_chunks.append(chunk)
-        sorted_chunks = sorted(
-            all_chunks, key=lambda chunk: (chunk["tag"], chunk["chunk_index"])
-        )
-        cheatsheets = _cheatsheets_from_chunks(sorted_chunks)
+                    all_chunk_dicts.append(chunk)
+        all_chunks = chunks_from_dicts(all_chunk_dicts)
+        cheatsheets = _cheatsheets_from_chunks(all_chunks)
         if (
             not cheatsheets
             or not isinstance(cheatsheets, list)

poma/integrations/llamaindex_poma.py

@@ -23,9 +23,8 @@ from llama_index.core.schema import (
 from pydantic import PrivateAttr
 
 from poma import Poma
-from poma.client import ALLOWED_FILE_EXTENSIONS
-from poma.retrieval import _cheatsheets_from_chunks
 from poma.exceptions import InvalidInputError
+from poma.retrieval import chunks_from_dicts, _cheatsheets_from_chunks
 
 __all__ = ["PomaFileReader", "PomaChunksetNodeParser", "PomaCheatsheetRetrieverLI"]
 
@@ -54,9 +53,6 @@ class PomaFileReader(BaseReader):
             if not file_path.is_file():
                 return
             file_extension = file_path.suffix.lower()
-            if not file_extension or file_extension not in ALLOWED_FILE_EXTENSIONS:
-                skipped += 1
-                return
             file_bytes = file_path.read_bytes()
             file_hash = hashlib.md5(file_bytes).hexdigest()
             if file_extension == ".pdf":
@@ -87,13 +83,10 @@ class PomaFileReader(BaseReader):
         else:
             raise FileNotFoundError(f"Unsupported path type (not file/dir): {path}")
 
-        allowed = ", ".join(sorted(ALLOWED_FILE_EXTENSIONS))
         if not documents:
-            raise InvalidInputError(f"No supported files found. Allowed: {allowed}")
+            raise InvalidInputError(f"No supported files found.")
         if skipped > 0:
-            print(
-                f"Skipped {skipped} file(s) due to unsupported or unreadable type. Allowed: {allowed}"
-            )
+            print(f"Skipped {skipped} file(s) due to unsupported or unreadable type.")
         return documents
 
 
@@ -329,7 +322,7 @@ class PomaCheatsheetRetrieverLI(BaseRetriever):
 
     def _create_cheatsheet_llamaindex(self, chunked_nodes: list[NodeWithScore]) -> str:
         """Generate a single deduplicated cheatsheet from chunked nodes."""
-        all_chunks = []
+        all_chunk_dicts = []
         seen = set()
         for node in chunked_nodes:
             doc_id = node.metadata.get("doc_id", "unknown_doc")
@@ -344,11 +337,9 @@ class PomaCheatsheetRetrieverLI(BaseRetriever):
                     continue
                 seen.add(chunk_index)
                 chunk["tag"] = doc_id
-                all_chunks.append(chunk)
-        sorted_chunks = sorted(
-            all_chunks, key=lambda chunk: (chunk["tag"], chunk["chunk_index"])
-        )
-        cheatsheets = _cheatsheets_from_chunks(sorted_chunks)
+                all_chunk_dicts.append(chunk)
+        all_chunks = chunks_from_dicts(all_chunk_dicts)
+        cheatsheets = _cheatsheets_from_chunks(all_chunks)
         if (
             not cheatsheets
             or not isinstance(cheatsheets, list)

poma/retrieval.py

@@ -2,62 +2,206 @@
 from collections import defaultdict
 from itertools import chain
 from typing import Any
+import warnings
+
+
+def deprecated(replacement: str):
+    def decorator(func):
+        msg = (
+            f"{func.__name__}() is deprecated and will be removed in a future version. "
+            f"Use {replacement} instead."
+        )
+
+        def wrapper(*args, **kwargs):
+            warnings.warn(msg, DeprecationWarning, stacklevel=2)
+            return func(*args, **kwargs)
+
+        wrapper.__name__ = func.__name__
+        wrapper.__doc__ = (func.__doc__ or "") + f"\n\nDEPRECATED: {msg}\n"
+        return wrapper
+
+    return decorator
 
 
 def generate_cheatsheets(
-    relevant_chunksets: list[dict[str, Any]], all_chunks: list[dict[str, Any]]
+    relevant_chunksets: list[dict[str, Any]],
+    all_chunks: list[dict[str, Any]],
 ) -> list[dict[str, Any]]:
-    chunk_ids = [cs["chunks"] for cs in relevant_chunksets if "chunks" in cs]
-    chunk_ids = list(chain.from_iterable(chunk_ids))  # flatten the list
-    relevant_chunks = _get_relevant_chunks_for_ids(chunk_ids, all_chunks)
-    sorted_chunks = sorted(
-        relevant_chunks, key=lambda chunk: (chunk["tag"], chunk["chunk_index"])
-    )
-    return _cheatsheets_from_chunks(sorted_chunks)
+    # get chunks grouped by document file_id
+    doc_chunks = defaultdict(list)
+    for chunk in all_chunks:
+        file_id = chunk.get("file_id") or chunk.get("tag") or "single_doc"
+        chunk["file_id"] = file_id  # update
+        doc_chunks[file_id].append(chunk)
 
+    # Check for duplicate chunk_index values
+    # (necessary when file_id was not set in chunks)
+    for file_id, chunks in doc_chunks.items():
+        chunk_indices = [c["chunk_index"] for c in chunks]
+        if len(chunk_indices) != len(set(chunk_indices)):
+            raise ValueError(f"Duplicate chunk_index found for file_id: {file_id}")
 
-def generate_single_cheatsheet(
-    relevant_chunksets: list[dict[str, Any]], all_chunks: list[dict[str, Any]]
-) -> str:
+    # get relevant chunksets grouped by document file_id
+    relevant_chunksets_per_doc = defaultdict(list)
+    for chunkset in relevant_chunksets:
+        file_id = chunkset.get("file_id") or chunkset.get("tag") or "single_doc"
+        chunkset["file_id"] = file_id  # update
+        if "chunks" not in chunkset:
+            raise ValueError(
+                "Chunkset not valid; must contain a 'chunks' key with a list of chunk IDs."
+            )
+        relevant_chunksets_per_doc[file_id].append(chunkset)
 
-    def prepare_single_doc_chunks(
-        chunk_dicts: list[dict[str, Any]],
-    ) -> list[dict[str, Any]]:
-        # Make sure there are no duplicate chunk_index values
-        check_dict = defaultdict(set)
-        has_duplicates = any(
-            chunk["chunk_index"] in check_dict[chunk["tag"]]
-            or check_dict[chunk["tag"]].add(chunk["chunk_index"])
-            for chunk in chunk_dicts
-        )
-        if has_duplicates:
+    # Ensure that chunksets and chunks correspond to the same file_ids
+    for file_id in relevant_chunksets_per_doc.keys():
+        if file_id not in doc_chunks:
             raise ValueError(
-                "Duplicate chunk indices found in single document mode. "
-                "Each chunk must have a unique index."
+                f"Chunksets contain file_id '{file_id}' which is not present in the chunks."
             )
-        # Use a fixed tag for chunks from single documents
-        for chunk_dict in chunk_dicts:
-            chunk_dict["tag"] = "single_doc"
-        return chunk_dicts
-
-    chunk_ids = [cs["chunks"] for cs in relevant_chunksets if "chunks" in cs]
-    chunk_ids = list(chain.from_iterable(chunk_ids))  # flatten the list
-    relevant_chunks = _get_relevant_chunks_for_ids(chunk_ids, all_chunks)
-    relevant_chunks = prepare_single_doc_chunks(relevant_chunks)
-    sorted_chunks = sorted(
-        relevant_chunks, key=lambda chunk: (chunk["tag"], chunk["chunk_index"])
+
+    # retrieve relevant chunks with content per document
+    relevant_content_chunks: list[RetrievalChunk] = []
+    for file_id, chunksets_per_doc in relevant_chunksets_per_doc.items():
+        chunk_ids = list(  # flattened list
+            chain.from_iterable(chunkset["chunks"] for chunkset in chunksets_per_doc)
+        )
+        relevant_chunks_dict = _get_relevant_chunks_for_ids(
+            chunk_ids, doc_chunks[file_id]
+        )
+        relevant_chunks: list[RetrievalChunk] = chunks_from_dicts(relevant_chunks_dict)
+        relevant_content_chunks.extend(relevant_chunks)
+
+    return _cheatsheets_from_chunks(relevant_content_chunks)
+
+
+@deprecated("generate_cheatsheets(relevant_chunksets, all_chunks)")
+def generate_single_cheatsheet(
+    relevant_chunksets: list[dict[str, Any]],
+    all_chunks: list[dict[str, Any]],
+) -> str:
+    cheatsheets = generate_cheatsheets(
+        relevant_chunksets=relevant_chunksets,
+        all_chunks=all_chunks,
     )
-    cheatsheets = _cheatsheets_from_chunks(sorted_chunks)
-    if (
-        not cheatsheets
-        or not isinstance(cheatsheets, list)
-        or len(cheatsheets) == 0
-        or "content" not in cheatsheets[0]
+    return cheatsheets[0].get("content", "") if cheatsheets else ""
+
+
+########################
+# RetrievalChunk Class #
+########################
+
+
+class RetrievalChunk:
+    """
+    Represents a chunk of text with associated metadata.
+    Attributes:
+        index (int): The index of the chunk within a sequence.
+        file_id (str): The id associating the chunk with a document.
+        content (str): The textual content of the chunk.
+        depth_rebased (int, optional): The hierarchical depth of the chunk content.
+            In cheatsheets, this affects indentation for certain text parts.
+            Currently only used for code blocks.
+    """
+
+    def __init__(
+        self,
+        index: int,
+        file_id: str,
+        content: str,
+        depth_rebased: int | None,
+    ):
+        self.index = index
+        self.file_id = file_id
+        self.content = content
+        self.depth_rebased = depth_rebased
+
+    @classmethod
+    def from_chunk_dict(
+        cls,
+        chunk_dict: dict,
+        block_min_depth: int | None,
     ):
-        raise Exception(
-            "Unknown error; cheatsheet could not be created from input chunks."
+        if block_min_depth is not None:
+            depth = int(chunk_dict["depth"])
+            depth_rebased = cls._rebase_depth(depth, block_min_depth)
+        else:
+            depth_rebased = None
+        return cls(
+            index=int(chunk_dict["chunk_index"]),
+            file_id=str(
+                chunk_dict.get("file_id") or chunk_dict.get("tag") or "single_doc"
+            ),
+            content=str(chunk_dict["content"]),
+            depth_rebased=depth_rebased,
+        )
+
+    @staticmethod
+    def _rebase_depth(depth: int, min_depth: int, base_unit: int = 0) -> int | None:
+        rebased = depth - min_depth + base_unit
+        return max(0, rebased)
+
+    def __repr__(self):
+        return f"RetrievalChunk(index={self.index}, file_id={self.file_id}, content={self.content}), depth_rebased={self.depth_rebased}"
+
+
+def chunks_from_dicts(chunk_dicts: list[dict]) -> list[RetrievalChunk]:
+    """
+    Converts a list of chunk dictionaries into a list of Chunk objects.
+    File_ids are needed to identify chunks from different documents;
+    if is_single_doc is True, all chunks are assumed to come from a single document
+    and file_id is optional.
+    Args:
+        chunk_dicts (list[dict]): A list of dictionaries, each representing a chunk with required keys:
+            - "chunk_index": The index of the chunk within the document.
+            - "file_id": The identifier of the document.
+            - "content": The textual content of the chunk.
+            - "depth": The depth or level of the chunk.
+    Returns:
+        list[Chunk]: A list of Chunk objects with the textual content needed for the cheatsheets.
+    """
+
+    # Determine the minimum depth per code block
+    min_depth_per_code_block: dict[str, int] = {}
+    for chunk_dict in chunk_dicts:
+        block_id = chunk_dict.get("code")
+        if block_id is None:
+            continue
+        depth = int(chunk_dict["depth"])
+        current = min_depth_per_code_block.get(block_id)
+        min_depth_per_code_block[block_id] = (
+            depth if current is None else min(current, depth)
+        )
+
+    # Create Chunk objects
+    all_chunks: list[RetrievalChunk] = []
+    for chunk_dict in chunk_dicts:
+        code_id = chunk_dict.get("code")
+        if bool(code_id):
+            block_min_depth = min_depth_per_code_block.get(str(code_id))
+        else:
+            block_min_depth = None
+        chunk = RetrievalChunk.from_chunk_dict(chunk_dict, block_min_depth)
+        all_chunks.append(chunk)
+
+    # Sanity check: Make sure there are no duplicate chunk_index values
+    check_dict = defaultdict(set)
+    has_duplicates = any(
+        chunk.index in check_dict[chunk.file_id]
+        or check_dict[chunk.file_id].add(chunk.index)
+        for chunk in all_chunks
+    )
+    if has_duplicates:
+        raise ValueError(
+            "Duplicate chunk indices found in single document mode. "
+            "Each chunk must have a unique index."
         )
-    return cheatsheets[0]["content"]
+
+    return all_chunks
+
+
+###################
+# Private Methods #
+###################
 
 
 def _get_relevant_chunks_for_ids(
@@ -138,39 +282,58 @@ def _get_relevant_chunks_for_ids(
 
 
 def _cheatsheets_from_chunks(
-    content_chunks: list[dict[str, Any]],
+    content_chunks: list[RetrievalChunk],
 ) -> list[dict[str, Any]]:
+    if not content_chunks:
+        return []
+
+    if isinstance(content_chunks[0], dict):
+        raise ValueError(
+            "Input to _cheatsheets_from_chunks must be a list of RetrievalChunk objects, not dicts."
+            "Use chunks_from_dicts() to convert dicts to RetrievalChunk objects first."
+        )
+
+    def _format_chunk_content(chunk: "RetrievalChunk") -> str:
+        if not getattr(chunk, "depth_rebased", False):
+            return chunk.content
+        else:
+            indent = " " * 4 * (chunk.depth_rebased or 0)
+            return f"{indent}{chunk.content}"
+
     cheatsheets: list[dict] = []
 
     compressed_data = {}
+    content_chunks = sorted(content_chunks, key=lambda c: (c.file_id, c.index))
     for chunk in content_chunks:
-        if chunk["tag"] not in compressed_data:
-            # If there is data stored for a previous tag, save it to the cheatsheets list
+        if chunk.file_id not in compressed_data:
+            # If there is data stored for a previous file_id, save it to the cheatsheets list
             if compressed_data:
                 for key, value in compressed_data.items():
-                    cheatsheets.append({"tag": key, "content": value["content"]})
-            # Clear the compressed_data for the current tag
+                    cheatsheets.append(
+                        {"file_id": key, "tag": key, "content": value["content"]}
+                    )
+            # Clear the compressed_data for the current file_id
             compressed_data.clear()
-            # Start a new entry for the current tag
-            compressed_data[chunk["tag"]] = {
-                "content": chunk["content"],
-                "last_chunk": chunk["chunk_index"],
+            # Start a new entry for the current file_id
+            compressed_data[chunk.file_id] = {
+                "content": _format_chunk_content(chunk),
+                "last_chunk": chunk.index,
             }
         else:
+            chunk_content = _format_chunk_content(chunk)
             # Check if chunks are consecutive
-            if (
-                chunk["chunk_index"]
-                == int(compressed_data[chunk["tag"]]["last_chunk"]) + 1
-            ):
-                compressed_data[chunk["tag"]]["content"] += "\n" + chunk["content"]
+            if chunk.index == int(compressed_data[chunk.file_id]["last_chunk"]) + 1:
+                compressed_data[chunk.file_id]["content"] += "\n" + chunk_content
             else:
-                compressed_data[chunk["tag"]]["content"] += "\n[…]\n" + chunk["content"]
+                compressed_data[chunk.file_id]["content"] += "\n[…]\n" + chunk_content
             # Update the last chunk index
-            compressed_data[chunk["tag"]]["last_chunk"] = chunk["chunk_index"]
+            compressed_data[chunk.file_id]["last_chunk"] = chunk.index
 
     # Save the last processed entry to the cheatsheets list
     if compressed_data:
         for key, value in compressed_data.items():
-            cheatsheets.append({"tag": key, "content": value["content"]})
+            cheatsheets.append(
+                {"file_id": key, "tag": key, "content": value["content"]}
+            )
 
     return cheatsheets

{poma-0.1.4.dist-info → poma-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: poma
-Version: 0.1.4
+Version: 0.2.1
 Summary: Official Python SDK for the Poma document-processing API
 Author-email: "POMA AI GmbH, Berlin" <sdk@poma-ai.com>
 License-Expression: MPL-2.0
@@ -40,15 +40,15 @@ pip install poma
 
 For integrations into LangChain and LlamaIndex:
 ```bash
-pip install poma[integrations]
+pip install 'poma[integrations]'
 # Or LangChain/LlamaIndex including example extras:
-pip install poma[integration-examples]
+pip install 'poma[integration-examples]'
 ```
 
 
 - You may also want: `pip install python-dotenv` to load API keys from a .env file.
 - API keys required (POMA_API_KEY) for the POMA AI client via environment variables.
-- **To request a POMA_API_KEY, please contact us at api@poma-ai.com**
+- **To request a POMA_API_KEY, please contact us at sdk@poma-ai.com**
 
 
 ### Example Implementations — all examples, integrations, and additional information can be found in our GitHub repository: [poma-ai/poma](https://github.com/poma-ai/)

poma-0.2.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+poma/__init__.py,sha256=SARVBTJw2pkIXR2_OYMPYjB7W335er_2-9j4yhzVTZI,266
+poma/client.py,sha256=sAaW7lz6s6WQDhkjMYfhOtd_d8baVflToc7q5l9BTYo,14337
+poma/exceptions.py,sha256=CsawBf0pwtZvm_0kPxLQ5WsIU_vCzNoMEMF_o5OJkBQ,519
+poma/retrieval.py,sha256=KcKa1pKjgWVieHflcPC6oiZO8ZZECHEpzqVlHC2RPDk,12475
+poma/integrations/__init__.py,sha256=xrrJluggTLtrKs4jLOZUWkFENqWSHSnhCqYQYY51kq0,405
+poma/integrations/langchain_poma.py,sha256=wDVrPJde7gRJR4ml3wS4qiL_dM4w2RDKVZYM_FrYgHM,14202
+poma/integrations/llamaindex_poma.py,sha256=biSAwcbKShsTdnjtCtM63GCzD5Nig4-85ZBquB_CG8w,14516
+poma-0.2.1.dist-info/licenses/LICENSE,sha256=YRzZ4sQOWV3ut0G4LHZJ2hT90shzZufGlXoIx4LWFEo,15254
+poma-0.2.1.dist-info/METADATA,sha256=vVhHvee7kq5DV4Yu7iCjk_abu8i3vW-I6PIoaDq_J5w,3083
+poma-0.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+poma-0.2.1.dist-info/top_level.txt,sha256=f_3c5Y6SojNnH0iiiE898fIKF6R2LqWyAw-BGi-72YI,5
+poma-0.2.1.dist-info/RECORD,,

poma-0.1.4.dist-info/RECORD

@@ -1,12 +0,0 @@
-poma/__init__.py,sha256=SARVBTJw2pkIXR2_OYMPYjB7W335er_2-9j4yhzVTZI,266
-poma/client.py,sha256=HOlNVfSPYzq1RjLVUtIqbObH5QSLRbi7KnbCnBjpUh4,14201
-poma/exceptions.py,sha256=5d8SdIIRFotKUJJAy9mct2q44oEmAsR15OVEmkLDfkQ,518
-poma/retrieval.py,sha256=bm68_1QscJXa76sxVuAkcwdeGsvEZaQkrY3-3uUxrIg,6730
-poma/integrations/__init__.py,sha256=xrrJluggTLtrKs4jLOZUWkFENqWSHSnhCqYQYY51kq0,405
-poma/integrations/langchain_poma.py,sha256=IL3pWWGCEK_O0JagpnKPFRwKclyNTwPcaTTdKJkYfYY,14608
-poma/integrations/llamaindex_poma.py,sha256=n3M71QXGVA2RTsUC24ZTt__VHEgsTbIW9BVwEn1Xxbg,14868
-poma-0.1.4.dist-info/licenses/LICENSE,sha256=YRzZ4sQOWV3ut0G4LHZJ2hT90shzZufGlXoIx4LWFEo,15254
-poma-0.1.4.dist-info/METADATA,sha256=LBg2pxFEr3zGgMahJS0kQSKAwS003fvMLQfHLpiZsX0,3079
-poma-0.1.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-poma-0.1.4.dist-info/top_level.txt,sha256=f_3c5Y6SojNnH0iiiE898fIKF6R2LqWyAw-BGi-72YI,5
-poma-0.1.4.dist-info/RECORD,,