From 95dbd13a883588fafafb3ffcaec0cdaceb81fbc4 Mon Sep 17 00:00:00 2001 From: Rafid Date: Mon, 9 Dec 2024 14:03:58 -0800 Subject: [PATCH] Renames related to the recent API renames (#58) Addressed few other remaining renames related to [the recent API change](https://github.com/cohere-ai/compass/pull/594). --- cohere/compass/__init__.py | 2 +- cohere/compass/clients/compass.py | 12 ++++++------ cohere/compass/models/documents.py | 2 +- cohere/compass/models/search.py | 12 ++++++------ pyproject.toml | 2 +- tests/test_compass_client.py | 6 +++--- 6 files changed, 18 insertions(+), 18 deletions(-) diff --git a/cohere/compass/__init__.py b/cohere/compass/__init__.py index c048f5d..323dc6b 100644 --- a/cohere/compass/__init__.py +++ b/cohere/compass/__init__.py @@ -34,6 +34,6 @@ class GroupAuthorizationActions(str, Enum): class GroupAuthorizationInput(BaseModel): - doc_ids: List[str] + document_ids: List[str] authorized_groups: List[str] action: GroupAuthorizationActions diff --git a/cohere/compass/clients/compass.py b/cohere/compass/clients/compass.py index a530a95..3331b4c 100644 --- a/cohere/compass/clients/compass.py +++ b/cohere/compass/clients/compass.py @@ -112,7 +112,7 @@ def __init__( "put_documents": self.session.put, "search_documents": self.session.post, "search_chunks": self.session.post, - "add_context": self.session.post, + "add_attributes": self.session.post, "refresh": self.session.post, "upload_documents": self.session.post, "edit_group_authorization": self.session.post, @@ -133,7 +133,7 @@ def __init__( "put_documents": "/api/v1/indexes/{index_name}/documents", "search_documents": "/api/v1/indexes/{index_name}/documents/_search", "search_chunks": "/api/v1/indexes/{index_name}/documents/_search_chunks", - "add_context": "/api/v1/indexes/{index_name}/documents/add_context/{document_id}", + "add_attributes": "/api/v1/indexes/{index_name}/documents/{document_id}/_add_attributes", "refresh": "/api/v1/indexes/{index_name}/_refresh", "upload_documents": "/api/v1/indexes/{index_name}/documents/_upload", "edit_group_authorization": "/api/v1/indexes/{index_name}/group_authorization", @@ -227,7 +227,7 @@ def list_indexes(self): index_name="", ) - def add_context( + def add_attributes( self, *, index_name: str, @@ -247,7 +247,7 @@ def add_context( """ return self._send_request( - api_name="add_context", + api_name="add_attributes", document_id=document_id, data=context, max_retries=max_retries, @@ -289,7 +289,7 @@ def upload_document( filebytes: bytes, content_type: str, document_id: uuid.UUID, - context: Dict[str, Any] = {}, + attributes: Dict[str, Any] = {}, max_retries: int = DEFAULT_MAX_RETRIES, sleep_retry_seconds: int = DEFAULT_SLEEP_RETRY_SECONDS, ) -> Optional[Union[str, Dict[str, Any]]]: @@ -317,7 +317,7 @@ def upload_document( content_type=content_type, content_length_bytes=len(filebytes), content_encoded_bytes=b64, - context=context, + attributes=attributes, ) result = self._send_request( diff --git a/cohere/compass/models/documents.py b/cohere/compass/models/documents.py index f584c1c..daf1933 100644 --- a/cohere/compass/models/documents.py +++ b/cohere/compass/models/documents.py @@ -175,7 +175,7 @@ class ParseableDocument(BaseModel): content_type: str content_length_bytes: PositiveInt # File size must be a non-negative integer content_encoded_bytes: str # Base64-encoded file contents - context: Dict[str, Any] = Field(default_factory=dict) + attributes: Dict[str, Any] = Field(default_factory=dict) class PushDocumentsInput(BaseModel): diff --git a/cohere/compass/models/search.py b/cohere/compass/models/search.py index 7ee42c8..9cbd8b2 100644 --- a/cohere/compass/models/search.py +++ b/cohere/compass/models/search.py @@ -14,7 +14,7 @@ class AssetInfo(BaseModel): class RetrievedChunk(BaseModel): chunk_id: str sort_id: int - parent_doc_id: str + parent_document_id: str content: Dict[str, Any] origin: Optional[Dict[str, Any]] = None assets_info: Optional[list[AssetInfo]] = None @@ -22,9 +22,9 @@ class RetrievedChunk(BaseModel): class RetrievedDocument(BaseModel): - doc_id: str + document_id: str path: str - parent_doc_id: str + parent_document_id: str content: Dict[str, Any] index_fields: Optional[List[str]] = None authorized_groups: Optional[List[str]] = None @@ -32,8 +32,8 @@ class RetrievedDocument(BaseModel): score: float -class RetrieveChunkExtended(RetrievedChunk): - doc_id: str +class RetrievedChunkExtended(RetrievedChunk): + document_id: str path: str index_fields: Optional[List[str]] = None @@ -43,7 +43,7 @@ class SearchDocumentsResponse(BaseModel): class SearchChunksResponse(BaseModel): - hits: List[RetrieveChunkExtended] + hits: List[RetrievedChunkExtended] class SearchFilter(BaseModel): diff --git a/pyproject.toml b/pyproject.toml index 47a8c6a..af2a441 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "compass-sdk" -version = "0.9.0" +version = "0.9.1" authors = [] description = "Compass SDK" readme = "README.md" diff --git a/tests/test_compass_client.py b/tests/test_compass_client.py index 4214915..ae99dbd 100644 --- a/tests/test_compass_client.py +++ b/tests/test_compass_client.py @@ -73,14 +73,14 @@ def test_refresh_is_valid(requests_mock: Mocker): ) -def test_add_context_is_valid(requests_mock: Mocker): +def test_add_attributes_is_valid(requests_mock: Mocker): compass = CompassClient(index_url="http://test.com") - compass.add_context( + compass.add_attributes( index_name="test_index", document_id="test_id", context={"fake": "context"} ) assert requests_mock.request_history[0].method == "POST" assert ( requests_mock.request_history[0].url - == "http://test.com/api/v1/indexes/test_index/documents/add_context/test_id" + == "http://test.com/api/v1/indexes/test_index/documents/test_id/_add_attributes" ) assert requests_mock.request_history[0].body == b'{"fake": "context"}'