Fix argument names to document_id

cohere-ai · Dec 6, 2024 · 3bf3747 · 3bf3747
1 parent 4917353
commit 3bf3747
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 27 deletions.
diff --git a/cohere/compass/clients/compass.py b/cohere/compass/clients/compass.py
@@ -570,7 +570,9 @@ def _get_request_blocks(
         num_chunks = 0
         for _, doc in enumerate(docs, 1):
             if doc.status != CompassDocumentStatus.Success:
-                logger.error(f"Document {doc.metadata.document_id} has errors: {doc.errors}")
+                logger.error(
+                    f"Document {doc.metadata.document_id} has errors: {doc.errors}"
+                )
                 for error in doc.errors:
                     errors.append({doc.metadata.document_id: list(error.values())[0]})
             else:

diff --git a/cohere/compass/clients/parser.py b/cohere/compass/clients/parser.py
@@ -270,36 +270,36 @@ def process_file(
         return docs
 
     @staticmethod
-    def _adapt_doc_id_compass_doc(doc: Dict[Any,Any]) -> CompassDocument:
+    def _adapt_doc_id_compass_doc(doc: Dict[Any, Any]) -> CompassDocument:
         """
         Adapt the doc_id to document_id
         """
 
-        metadata = doc['metadata']
-        if not 'document_id' in metadata:
-            metadata['document_id'] = metadata.pop('doc_id')
-            metadata['parent_document_id'] = metadata.pop('parent_doc_id')
+        metadata = doc["metadata"]
+        if "document_id" not in metadata:
+            metadata["document_id"] = metadata.pop("doc_id")
+            metadata["parent_document_id"] = metadata.pop("parent_doc_id")
 
-        chunks = doc['chunks']
+        chunks = doc["chunks"]
         for chunk in chunks:
-            if not 'parent_document_id' in chunk:
-                chunk['parent_document_id'] = chunk.pop('parent_doc_id')
-            if not 'document_id' in chunk:
-                chunk['document_id'] = chunk.pop('doc_id')
-            if not 'path' in chunk:
-                chunk['path'] = doc['metadata']['filename']
-                
+            if "parent_document_id" not in chunk:
+                chunk["parent_document_id"] = chunk.pop("parent_doc_id")
+            if "document_id" not in chunk:
+                chunk["document_id"] = chunk.pop("doc_id")
+            if "path" not in chunk:
+                chunk["path"] = doc["metadata"]["filename"]
+
         res = CompassDocument(
-            filebytes=doc['filebytes'], 
+            filebytes=doc["filebytes"],
             metadata=metadata,
-            content=doc['content'],
-            content_type=doc['content_type'],
-            elements=doc['elements'],
+            content=doc["content"],
+            content_type=doc["content_type"],
+            elements=doc["elements"],
             chunks=chunks,
-            index_fields=doc['index_fields'],
-            errors=doc['errors'],
-            ignore_metadata_errors=doc['ignore_metadata_errors'],
-            markdown=doc['markdown']
+            index_fields=doc["index_fields"],
+            errors=doc["errors"],
+            ignore_metadata_errors=doc["ignore_metadata_errors"],
+            markdown=doc["markdown"],
         )
 
-        return res
+        return res
diff --git a/cohere/compass/models/documents.py b/cohere/compass/models/documents.py
@@ -147,7 +147,6 @@ class Chunk(BaseModel):
     origin: Optional[Dict[str, Any]] = None
     assets: Optional[List[DocumentChunkAsset]] = None
     asset_ids: Optional[List[str]] = None
-
 
 
 class Document(BaseModel):

diff --git a/tests/test_compass_client.py b/tests/test_compass_client.py
@@ -6,7 +6,7 @@
 
 def test_delete_url_formatted_with_doc_and_index(requests_mock: Mocker):
     compass = CompassClient(index_url="http://test.com")
-    compass.delete_document(index_name="test_index", doc_id="test_id")
+    compass.delete_document(index_name="test_index", document_id="test_id")
     assert (
         requests_mock.request_history[0].url
         == "http://test.com/api/v1/indexes/test_index/documents/test_id"
@@ -55,7 +55,7 @@ def test_list_indices_is_valid(requests_mock: Mocker):
 
 def test_get_documents_is_valid(requests_mock: Mocker):
     compass = CompassClient(index_url="http://test.com")
-    compass.get_document(index_name="test_index", doc_id="test_id")
+    compass.get_document(index_name="test_index", document_id="test_id")
     assert requests_mock.request_history[0].method == "GET"
     assert (
         requests_mock.request_history[0].url
@@ -76,7 +76,7 @@ def test_refresh_is_valid(requests_mock: Mocker):
 def test_add_context_is_valid(requests_mock: Mocker):
     compass = CompassClient(index_url="http://test.com")
     compass.add_context(
-        index_name="test_index", doc_id="test_id", context={"fake": "context"}
+        index_name="test_index", document_id="test_id", context={"fake": "context"}
     )
     assert requests_mock.request_history[0].method == "POST"
     assert (