Merge pull request #997 from gchq/token-download-method

Added get_download_file_token method to client
gchq · Jan 12, 2024 · 5211f3b · 5211f3b
2 parents 8cc441d + b69fa5a
commit 5211f3b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/lib/python-beta/src/bailo/__init__.py b/lib/python-beta/src/bailo/__init__.py
@@ -1,7 +1,6 @@
 from __future__ import annotations
 
-# SEMANTIC VERSION
-__version__ = "0.3.0"
+__version__ = "0.4.0"
 
 from bailo.core.agent import Agent, PkiAgent, TokenAgent
 from bailo.core.client import Client

diff --git a/lib/python-beta/src/bailo/core/client.py b/lib/python-beta/src/bailo/core/client.py
@@ -4,7 +4,7 @@
 from io import BytesIO
 from typing import Any
 
-from bailo.core.agent import Agent
+from bailo.core.agent import Agent, TokenAgent
 from bailo.core.enums import ModelVisibility, SchemaKind
 from bailo.core.utils import filter_none
 
@@ -311,7 +311,13 @@ def get_download_file(
         :param buffer: BytesIO object for bailo to write to
         :return: The unique file ID
         """
-        req = self.agent.get(f"{self.url}/v2/model/{model_id}/file/{file_id}/download", stream=True, timeout=10_000)
+        if isinstance(self.agent, TokenAgent):
+            req = self.agent.get(
+                f"{self.url}/v2/token/model/{model_id}/file/{file_id}/download", stream=True, timeout=10_000
+            )
+        else:
+            req = self.agent.get(f"{self.url}/v2/model/{model_id}/file/{file_id}/download", stream=True, timeout=10_000)
+
         shutil.copyfileobj(req.raw, buffer)
         return file_id