From 7357fb98e6c8045a599e559c497234095fa20a48 Mon Sep 17 00:00:00 2001
From: artemiyatcohere <artemiy@cohere.com>
Date: Tue, 10 Dec 2024 17:22:46 +0000
Subject: [PATCH] Add asset_type field to the assets structures

---
 cohere/compass/models/documents.py | 13 +++++++++++++
 cohere/compass/models/search.py    |  3 +++
 2 files changed, 16 insertions(+)

diff --git a/cohere/compass/models/documents.py b/cohere/compass/models/documents.py
index daf1933..585e589 100644
--- a/cohere/compass/models/documents.py
+++ b/cohere/compass/models/documents.py
@@ -22,7 +22,20 @@ class CompassDocumentMetadata(ValidatedModel):
     parent_document_id: str = ""
 
 
+class AssetType(str, Enum):
+    def __str__(self) -> str:
+        return self.value
+
+    # A page that has been rendered as an image
+    PAGE_IMAGE = "page_image"
+    # A Markdown representation of a page's content
+    PAGE_MARKDOWN = "page_markdown"
+    # A dump of the text extracted from a document
+    DOCUMENT_TEXT = "document_text"
+
+
 class CompassDocumentChunkAsset(BaseModel):
+    asset_type: AssetType
     content_type: str
     asset_data: str
 
diff --git a/cohere/compass/models/search.py b/cohere/compass/models/search.py
index 9cbd8b2..8fe847a 100644
--- a/cohere/compass/models/search.py
+++ b/cohere/compass/models/search.py
@@ -5,8 +5,11 @@
 # 3rd party imports
 from pydantic import BaseModel
 
+from cohere.compass.models.documents import AssetType
+
 
 class AssetInfo(BaseModel):
+    asset_type: AssetType
     content_type: str
     presigned_url: str