feat: integration models support added

Author: csgulati09

portkey_ai/__init__.py

@@ -127,6 +127,8 @@
     AsyncIntegrations,
     IntegrationsWorkspaces,
     AsyncIntegrationsWorkspaces,
+    IntegrationsModels,
+    AsyncIntegrationsModels,
 )
 
 from portkey_ai.version import VERSION
@@ -273,4 +275,6 @@
     "AsyncIntegrations",
     "IntegrationsWorkspaces",
     "AsyncIntegrationsWorkspaces",
+    "IntegrationsModels",
+    "AsyncIntegrationsModels",
 ]

portkey_ai/api_resources/__init__.py

@@ -115,6 +115,8 @@
     AsyncIntegrations,
     IntegrationsWorkspaces,
     AsyncIntegrationsWorkspaces,
+    IntegrationsModels,
+    AsyncIntegrationsModels,
 )
 from .utils import (
     Modes,
@@ -265,4 +267,6 @@
     "AsyncIntegrations",
     "IntegrationsWorkspaces",
     "AsyncIntegrationsWorkspaces",
+    "IntegrationsModels",
+    "AsyncIntegrationsModels",
 ]

portkey_ai/api_resources/apis/__init__.py

@@ -144,6 +144,8 @@
     AsyncIntegrations,
     IntegrationsWorkspaces,
     AsyncIntegrationsWorkspaces,
+    IntegrationsModels,
+    AsyncIntegrationsModels,
 )
 
 sys.modules["openai"] = vendored_openai  # For pydantic v1 and v2 compatibility
@@ -274,4 +276,6 @@
     "AsyncIntegrations",
     "IntegrationsWorkspaces",
     "AsyncIntegrationsWorkspaces",
+    "IntegrationsModels",
+    "AsyncIntegrationsModels",
 ]

portkey_ai/api_resources/apis/integrations.py

@@ -10,6 +10,7 @@ class Integrations(APIResource):
     def __init__(self, client: APIClient) -> None:
         super().__init__(client)
         self.workspaces = IntegrationsWorkspaces(client)
+        self.models = IntegrationsModels(client)
 
     def create(
         self,
@@ -174,10 +175,54 @@ def update(
         )
 
 
+class IntegrationsModels(APIResource):
+    def __init__(self, client: APIClient) -> None:
+        super().__init__(client)
+
+    def list(
+        self,
+        *,
+        provider_integration_id: Optional[str] = None,
+    ) -> GenericResponse:
+        return self._get(
+            f"{PortkeyApiPaths.INTEGRATIONS_API}/{provider_integration_id}/models",
+            params=None,
+            body=None,
+            cast_to=GenericResponse,
+            stream=False,
+            stream_cls=None,
+            headers={},
+        )
+
+    def update(
+        self,
+        *,
+        provider_integration_id: Optional[str] = None,
+        allow_all_models: Optional[bool] = None,
+        models: Optional[List[Dict[str, Any]]] = None,
+        **kwargs: Any,
+    ) -> GenericResponse:
+        body = {
+            "allow_all_models": allow_all_models,
+            "models": models,
+            **kwargs,
+        }
+        return self._put(
+            f"{PortkeyApiPaths.INTEGRATIONS_API}/{provider_integration_id}/models",
+            body=body,
+            params=None,
+            cast_to=GenericResponse,
+            stream=False,
+            stream_cls=None,
+            headers={},
+        )
+
+
 class AsyncIntegrations(AsyncAPIResource):
     def __init__(self, client: AsyncAPIClient) -> None:
         super().__init__(client)
         self.workspaces = AsyncIntegrationsWorkspaces(client)
+        self.models = AsyncIntegrationsModels(client)
 
     async def create(
         self,
@@ -340,3 +385,46 @@ async def update(
             stream_cls=None,
             headers={},
         )
+
+
+class AsyncIntegrationsModels(AsyncAPIResource):
+    def __init__(self, client: AsyncAPIClient) -> None:
+        super().__init__(client)
+
+    async def list(
+        self,
+        *,
+        provider_integration_id: Optional[str] = None,
+    ) -> GenericResponse:
+        return await self._get(
+            f"{PortkeyApiPaths.INTEGRATIONS_API}/{provider_integration_id}/models",
+            params=None,
+            body=None,
+            cast_to=GenericResponse,
+            stream=False,
+            stream_cls=None,
+            headers={},
+        )
+
+    async def update(
+        self,
+        *,
+        provider_integration_id: Optional[str] = None,
+        allow_all_models: Optional[bool] = None,
+        models: Optional[List[Dict[str, Any]]] = None,
+        **kwargs: Any,
+    ) -> GenericResponse:
+        body = {
+            "allow_all_models": allow_all_models,
+            "models": models,
+            **kwargs,
+        }
+        return await self._put(
+            f"{PortkeyApiPaths.INTEGRATIONS_API}/{provider_integration_id}/models",
+            body=body,
+            params=None,
+            cast_to=GenericResponse,
+            stream=False,
+            stream_cls=None,
+            headers={},
+        )