docs: add docstrings

Author: ankit-v2-3

videodb/client.py

@@ -148,6 +148,13 @@ def get_invoices(self) -> List[dict]:
         return self.get(path=f"{ApiPath.billing}/{ApiPath.invoices}")
 
     def create_event(self, event_prompt: str, label: str):
+        """Create an rtstream event.
+
+        :param str event_prompt: Prompt for the event
+        :param str label: Label for the event
+        :return: Event ID
+        :rtype: str
+        """
         event_data = self.post(
             f"{ApiPath.rtstream}/{ApiPath.event}",
             data={"event_prompt": event_prompt, "label": label},
@@ -156,6 +163,11 @@ def create_event(self, event_prompt: str, label: str):
         return event_data.get("event_id")
 
     def list_events(self):
+        """List all rtstream events.
+
+        :return: List of events
+        :rtype: list[dict]
+        """
         event_data = self.get(f"{ApiPath.rtstream}/{ApiPath.event}")
         return event_data.get("events", [])
 

videodb/collection.py

@@ -168,6 +168,13 @@ def delete_image(self, image_id: str) -> None:
     def connect_rtstream(
         self, url: str, name: str, sample_rate: int = None
     ) -> RTStream:
+        """Connect to an rtstream.
+
+        :param str url: URL of the rtstream
+        :param str name: Name of the rtstream
+        :param int sample_rate: Sample rate of the rtstream default is 1fps
+        :return: :class:`RTStream <RTStream>` object
+        """
         rtstream_data = self._connection.post(
             path=f"{ApiPath.rtstream}",
             data={
@@ -180,19 +187,31 @@ def connect_rtstream(
         return RTStream(self._connection, **rtstream_data)
 
     def get_rtstream(self, id: str) -> RTStream:
+        """Get an rtstream by its ID.
+
+        :param str id: ID of the rtstream
+        :return: :class:`RTStream <RTStream>` object
+        :rtype: :class:`videodb.rtstream.RTStream`
+        """
         rtstream_data = self._connection.get(
             path=f"{ApiPath.rtstream}/{id}",
         )
         return RTStream(self._connection, **rtstream_data)
 
     def list_rtstreams(self) -> List[RTStream]:
+        """List all rtstreams in the collection.
+
+        :return: List of :class:`RTStream <RTStream>` objects
+        :rtype: List[:class:`videodb.rtstream.RTStream`]
+        """
         rtstreams_data = self._connection.get(
             path=f"{ApiPath.rtstream}",
         )
         return [
             RTStream(self._connection, **rtstream)
             for rtstream in rtstreams_data.get("results")
         ]
+
     def generate_image(
         self,
         prompt: str,

videodb/rtstream.py

@@ -5,6 +5,17 @@
 
 
 class RTStreamSceneIndex:
+    """RTStreamSceneIndex class to interact with the rtstream scene index
+
+    :ivar str rtstream_index_id: Unique identifier for the rtstream scene index
+    :ivar str rtstream_id: ID of the rtstream this scene index belongs to
+    :ivar str extraction_type: Type of extraction
+    :ivar dict extraction_config: Configuration for extraction
+    :ivar str prompt: Prompt for scene extraction
+    :ivar str name: Name of the scene index
+    :ivar str status: Status of the scene index
+    """
+
     def __init__(
         self, _connection, rtstream_index_id: str, rtstream_id, **kwargs
     ) -> None:
@@ -21,6 +32,7 @@ def __repr__(self) -> str:
         return (
             f"RTStreamSceneIndex("
             f"rtstream_index_id={self.rtstream_index_id}, "
+            f"rtstream_id={self.rtstream_id}, "
             f"extraction_type={self.extraction_type}, "
             f"extraction_config={self.extraction_config}, "
             f"prompt={self.prompt}, "
@@ -29,6 +41,15 @@ def __repr__(self) -> str:
         )
 
     def get_scenes(self, start: int = None, end: int = None, page=1, page_size=100):
+        """Get rtstream scene index scenes.
+
+        :param int start: Start time of the scenes
+        :param int end: End time of the scenes
+        :param int page: Page number
+        :param int page_size: Number of scenes per page
+        :return: List of scenes
+        :rtype: List[dict]
+        """
         params = {"page": page, "page_size": page_size}
         if start and end:
             params["start"] = start
@@ -46,20 +67,37 @@ def get_scenes(self, start: int = None, end: int = None, page=1, page_size=100):
         }
 
     def start(self):
+        """Start the scene index.
+
+        :return: None
+        :rtype: None
+        """
         self._connection.patch(
             f"{ApiPath.rtstream}/{self.rtstream_id}/{ApiPath.index}/{ApiPath.scene}/{self.rtstream_index_id}/{ApiPath.status}",
             data={"action": "start"},
         )
         self.status = "connected"
 
     def stop(self):
+        """Stop the scene index.
+
+        :return: None
+        :rtype: None
+        """
         self._connection.patch(
             f"{ApiPath.rtstream}/{self.rtstream_id}/{ApiPath.index}/{ApiPath.scene}/{self.rtstream_index_id}/{ApiPath.status}",
             data={"action": "stop"},
         )
         self.status = "stopped"
 
     def create_alert(self, event_id, callback_url) -> str:
+        """Create an event alert.
+
+        :param str event_id: ID of the event
+        :param str callback_url: URL to receive the alert callback
+        :return: Alert ID
+        :rtype: str
+        """
         alert_data = self._connection.post(
             f"{ApiPath.rtstream}/{self.rtstream_id}/{ApiPath.index}/{self.rtstream_index_id}/{ApiPath.alert}",
             data={
@@ -70,25 +108,52 @@ def create_alert(self, event_id, callback_url) -> str:
         return alert_data.get("alert_id", None)
 
     def list_alerts(self):
+        """List all alerts for the rtstream scene index.
+
+        :return: List of alerts
+        :rtype: List[dict]
+        """
         alert_data = self._connection.get(
             f"{ApiPath.rtstream}/{self.rtstream_id}/{ApiPath.index}/{self.rtstream_index_id}/{ApiPath.alert}"
         )
         return alert_data.get("alerts", [])
 
     def enable_alert(self, alert_id):
+        """Enable an alert.
+
+        :param str alert_id: ID of the alert
+        :return: None
+        :rtype: None
+        """
         self._connection.patch(
             f"{ApiPath.rtstream}/{self.rtstream_id}/{ApiPath.index}/{self.rtstream_index_id}/{ApiPath.alert}/{alert_id}/{ApiPath.status}",
             data={"action": "enable"},
         )
 
     def disable_alert(self, alert_id):
+        """Disable an alert.
+
+        :param str alert_id: ID of the alert
+        :return: None
+        :rtype: None
+        """
         self._connection.patch(
             f"{ApiPath.rtstream}/{self.rtstream_id}/{ApiPath.index}/{self.rtstream_index_id}/{ApiPath.alert}/{alert_id}/{ApiPath.status}",
             data={"action": "disable"},
         )
 
 
 class RTStream:
+    """RTStream class to interact with the RTStream
+
+    :ivar str id: Unique identifier for the rtstream
+    :ivar str name: Name of the rtstream
+    :ivar str collection_id: ID of the collection this rtstream belongs to
+    :ivar str created_at: Timestamp of the rtstream creation
+    :ivar int sample_rate: Sample rate of the rtstream
+    :ivar str status: Status of the rtstream
+    """
+
     def __init__(self, _connection, id: str, **kwargs) -> None:
         self._connection = _connection
         self.id = id
@@ -110,20 +175,37 @@ def __repr__(self) -> str:
         )
 
     def start(self):
+        """Connect to the rtstream.
+
+        :return: None
+        :rtype: None
+        """
         self._connection.patch(
             f"{ApiPath.rtstream}/{self.id}/{ApiPath.status}",
             data={"action": "start"},
         )
         self.status = "connected"
 
     def stop(self):
+        """Disconnect from the rtstream.
+
+        :return: None
+        :rtype: None
+        """
         self._connection.patch(
             f"{ApiPath.rtstream}/{self.id}/{ApiPath.status}",
             data={"action": "stop"},
         )
         self.status = "stopped"
 
     def generate_stream(self, start, end):
+        """Generate a stream from the rtstream.
+
+        :param int start: Start time of the stream
+        :param int end: End time of the stream
+        :return: Stream URL
+        :rtype: str
+        """
         stream_data = self._connection.get(
             f"{ApiPath.rtstream}/{self.id}/{ApiPath.stream}",
             params={"start": start, "end": end},
@@ -139,6 +221,17 @@ def index_scenes(
         model_config={},
         name=None,
     ):
+        """Index scenes from the rtstream.
+
+        :param str extraction_type: Type of extraction
+        :param dict extraction_config: Configuration for extraction
+        :param str prompt: Prompt for scene extraction
+        :param str model_name: Name of the model
+        :param dict model_config: Configuration for the model
+        :param str name: Name of the scene index
+        :return: Scene index, :class:`RTStreamSceneIndex <RTStreamSceneIndex>` object
+        :rtype: :class:`videodb.rtstream.RTStreamSceneIndex`
+        """
         index_data = self._connection.post(
             f"{ApiPath.rtstream}/{self.id}/{ApiPath.index}/{ApiPath.scene}",
             data={
@@ -164,6 +257,11 @@ def index_scenes(
         )
 
     def list_scene_indexes(self):
+        """List all scene indexes for the rtstream.
+
+        :return: List of :class:`RTStreamSceneIndex <RTStreamSceneIndex>` objects
+        :rtype: List[:class:`videodb.rtstream.RTStreamSceneIndex`]
+        """
         index_data = self._connection.get(
             f"{ApiPath.rtstream}/{self.id}/{ApiPath.index}/{ApiPath.scene}"
         )
@@ -182,6 +280,12 @@ def list_scene_indexes(self):
         ]
 
     def get_scene_index(self, index_id: str) -> RTStreamSceneIndex:
+        """Get a scene index by its ID.
+
+        :param str index_id: ID of the scene index
+        :return: Scene index, :class:`RTStreamSceneIndex <RTStreamSceneIndex>` object
+        :rtype: :class:`videodb.rtstream.RTStreamSceneIndex`
+        """
         index_data = self._connection.get(
             f"{ApiPath.rtstream}/{self.id}/{ApiPath.index}/{index_id}"
         )