Updated model list endpoint

Author: orangetin

src/together/api.py

@@ -2,6 +2,7 @@
 from typing import Any, Dict, List, Optional, cast
 
 import requests
+import urllib.parse
 
 from together.files import Files
 from together.finetune import Finetune
@@ -39,21 +40,22 @@ def __init__(
         self.endpoint_url = endpoint_url
         self.supply_endpoint_url = supply_endpoint_url
 
-    def get_supply(self) -> Dict[str, Any]:
+    def get_all_models(self) -> Dict[str, Any]:
+        model_url = urllib.parse.urljoin(self.endpoint_url, "models/info?=")
+        headers = {
+            "Authorization": f"Bearer {self.together_api_key}",
+        }
         try:
             response = requests.get(
-                self.supply_endpoint_url,
-                json={
-                    "method": "together_getDepth",
-                    "id": 1,
-                },
+                model_url,
+                headers=headers,
             )
         except requests.exceptions.RequestException as e:
             self.logger.critical(f"Response error raised: {e}")
             exit_1(self.logger)
 
         try:
-            response_json = dict(response.json())
+            response_json = list(response.json())
         except Exception as e:
             self.logger.critical(
                 f"JSON Error raised: {e}\nResponse status code = {response.status_code}"
@@ -62,13 +64,6 @@ def get_supply(self) -> Dict[str, Any]:
 
         return response_json
 
-    def get_all_models(self) -> List[str]:
-        models = cast(List[str], self.get_supply()["result"].keys())
-
-        models = [str(sub[:-1]) for sub in models]  # remove the ? after the model names
-
-        return models
-
     def get_available_models(self) -> List[str]:
         res = self.get_supply()
         names = res["result"].keys()

src/together/commands/api.py

@@ -17,7 +17,6 @@ def add_parser(
     child_parsers = parser.add_subparsers(required=True)
 
     _add_list(child_parsers, parents=parents)
-    _add_raw(child_parsers, parents=parents)
 
 
 def _add_list(
@@ -26,35 +25,22 @@ def _add_list(
 ) -> None:
     list_model_subparser = parser.add_parser("list", parents=parents)
     list_model_subparser.add_argument(
-        "--all",
-        "-a",
-        help="List all models (available and unavailable)",
+        "--raw",
+        help="Raw details of all models",
         default=False,
         action="store_true",
     )
     list_model_subparser.set_defaults(func=_run_list)
 
 
-def _add_raw(
-    parser: argparse._SubParsersAction[argparse.ArgumentParser],
-    parents: List[argparse.ArgumentParser],
-) -> None:
-    list_parser = parser.add_parser("raw-supply", parents=parents)
-    list_parser.set_defaults(func=_run_raw)
-
-
 def _run_list(args: argparse.Namespace) -> None:
     api = API(endpoint_url=args.endpoint, log_level=args.log)
-
-    if args.all:
-        response = api.get_all_models()
+    response = api.get_all_models()
+    if args.raw:
+        print(json.dumps(response, indent=4))
     else:
-        response = api.get_available_models()
-
-    print(json.dumps(response, indent=4))
+        models = []
+        for i in response:
+            models.append(i["name"])
+        print(json.dumps(models, indent=4))
 
-
-def _run_raw(args: argparse.Namespace) -> None:
-    api = API(endpoint_url=args.endpoint, log_level=args.log)
-    response = api.get_supply()
-    print(json.dumps(response, indent=4))