Publish collection stats (#2)

Author: oeken

README.md

@@ -1,5 +1,8 @@
 # Needle Python Library
 
+[![PyPI - Version](https://img.shields.io/pypi/v/needle-python.svg)](https://pypi.org/project/needle-python)
+[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/needle-python.svg)](https://pypi.org/project/needle-python)
+
 This Python library provides convenient acccess to Needle API. There are various methods and data types which, we believe will help you explore Needle API quickly. There may be some functionality available in REST API earlier than this Python library. In any case, we recommend to take look the the complete [documentation](https://docs.needle-ai.com). Thank you for flying with us. 🚀
 
 ## Installation
@@ -95,3 +98,7 @@ If a request to Needle API fails, `needle.v1.models.Error` object will be thrown
 ## Support 📞
 
 If you have questions you can contact us in our [Discord channel](https://discord.gg/JzJcHgTyZx). 
+
+# License
+
+`needle-python` is distributed under the terms of the MIT license.

needle/v1/collections/__init__.py

@@ -12,6 +12,8 @@
     Collection,
     Error,
     SearchResult,
+    CollectionStats,
+    CollectionDataStats,
 )
 from needle.v1.collections.files import NeedleCollectionsFiles
 
@@ -129,7 +131,13 @@ def list(self):
             for c in body.get("result")
         ]
 
-    def search(self, collection_id: str, text: str):
+    def search(
+        self,
+        collection_id: str,
+        text: str,
+        max_distance: Optional[float] = None,
+        top_k: Optional[int] = None,
+    ):
         """
         Searches within a collection based on the provided parameters.
 
@@ -143,7 +151,11 @@ def search(self, collection_id: str, text: str):
             Error: If the API request fails.
         """
         endpoint = f"{self.search_endpoint}/{collection_id}/search"
-        req_body = {"text": text}
+        req_body = {
+            "text": text,
+            "max_distance": max_distance,
+            "top_k": top_k,
+        }
         resp = self.session.post(endpoint, headers=self.headers, json=req_body)
         body = resp.json()
         if resp.status_code >= 400:
@@ -156,3 +168,39 @@ def search(self, collection_id: str, text: str):
             )
             for r in body.get("result")
         ]
+
+    def get_stats(self, collection_id: str):
+        """
+        Retrieves statistics of a collection.
+
+        Args:
+            collection_id (str): The ID of the collection to retrieve statistics for.
+
+        Returns:
+            dict: The collection statistics.
+
+        Raises:
+            Error: If the API request fails.
+        """
+        endpoint = f"{self.endpoint}/{collection_id}/stats"
+        resp = self.session.get(endpoint, headers=self.headers)
+        body = resp.json()
+        if resp.status_code >= 400:
+            error = body.get("error")
+            raise Error(**error)
+
+        result = body.get("result")
+        data_stats = [
+            CollectionDataStats(
+                status=ds.get("status"),
+                files=ds.get("files"),
+                bytes=ds.get("bytes"),
+            )
+            for ds in result.get("data_stats")
+        ]
+        return CollectionStats(
+            data_stats=data_stats,
+            chunks_count=result.get("chunks_count"),
+            characters=result.get("characters"),
+            users=result.get("users"),
+        )

needle/v1/models.py

@@ -93,6 +93,29 @@ class CollectionFile:
     status: CollectionFileStatus
 
 
+@dataclass(frozen=True)
+class CollectionDataStats:
+    """
+    Represents data statistics of a collection in the Needle API.
+    """
+
+    status: Optional[str]
+    files: int
+    bytes: int
+
+
+@dataclass(frozen=True)
+class CollectionStats:
+    """
+    Represents statistics of a collection in the Needle API.
+    """
+
+    data_stats: list[CollectionDataStats]
+    chunks_count: int
+    characters: int
+    users: int
+
+
 @dataclass(frozen=True)
 class SearchResult:
     """

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "needle-python"
-version = "0.2.0"
+version = "0.3.0"
 description = "Needle client library for Python"
 authors = [
     "Onur Eken <m.onureken@gmail.com>",