Merge pull request trustyai-explainability#31 from m-misiura/RemoteProviderSpec

Update provider.py to ensure it can be loaded using `RemoteProviderSpec` introduced in lls 0.23.0

Author: m-misiura

llama_stack_provider_trustyai_fms/provider.py

@@ -1,36 +1,56 @@
 import logging
 
-from llama_stack.providers.datatypes import (
-    AdapterSpec,
-    Api,
-    ProviderSpec,
-    remote_provider_spec,
-)
-
-# Add logging at the top
+from llama_stack.providers.datatypes import Api, ProviderSpec
+
 logger = logging.getLogger(__name__)
 
+try:
+    from llama_stack.providers.datatypes import AdapterSpec, remote_provider_spec
+
+    USE_LEGACY = True
+    logger.debug("Using legacy remote_provider_spec")
+except ImportError:
+    from llama_stack.providers.datatypes import RemoteProviderSpec
+
+    USE_LEGACY = False
+    logger.debug("Using new RemoteProviderSpec")
+
 
 def get_provider_spec() -> ProviderSpec:
-    return remote_provider_spec(
-        api=Api.safety,
-        adapter=AdapterSpec(
-            adapter_type="trustyai_fms",
+    if USE_LEGACY:
+        return remote_provider_spec(
+            api=Api.safety,
+            adapter=AdapterSpec(
+                adapter_type="trustyai_fms",
+                config_class="llama_stack_provider_trustyai_fms.config.FMSSafetyProviderConfig",
+                module="llama_stack_provider_trustyai_fms",
+            ),
+        )
+    else:
+        return RemoteProviderSpec(
+            api=Api.safety,
+            provider_type="remote::trustyai_fms",
             config_class="llama_stack_provider_trustyai_fms.config.FMSSafetyProviderConfig",
             module="llama_stack_provider_trustyai_fms",
-        ),
-    )
+            adapter_type="trustyai_fms",
+        )
 
 
 def get_shields_provider_spec() -> ProviderSpec:
-    spec = remote_provider_spec(
-        api=Api.shields,
-        adapter=AdapterSpec(
-            adapter_type="trustyai_fms",
+    if USE_LEGACY:
+        return remote_provider_spec(
+            api=Api.shields,
+            adapter=AdapterSpec(
+                adapter_type="trustyai_fms",
+                config_class="llama_stack_provider_trustyai_fms.config.FMSSafetyProviderConfig",
+                module="llama_stack_provider_trustyai_fms",
+            ),
+        )
+    else:
+        return RemoteProviderSpec(
+            api=Api.shields,
+            provider_type="remote::trustyai_fms",
             config_class="llama_stack_provider_trustyai_fms.config.FMSSafetyProviderConfig",
             module="llama_stack_provider_trustyai_fms",
-        ),
-    )
-    # Add debug logging
-    logger.debug(f"Returning shields provider spec: {spec}")
-    return spec
+            adapter_type="trustyai_fms",
+        )

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "llama-stack-provider-trustyai-fms"
-version = "0.2.2"
+version = "0.2.3"
 description = "Remote safety provider for Llama Stack integrating FMS Guardrails Orchestrator and community detectors"
 authors = [
     {name = "GitHub: m-misiura"}

tests/unit/test_provider.py

@@ -1,4 +1,13 @@
-from llama_stack.providers.datatypes import AdapterSpec, Api, ProviderSpec
+from llama_stack.providers.datatypes import Api, ProviderSpec
+
+try:
+    from llama_stack.providers.datatypes import AdapterSpec
+
+    USE_LEGACY = True
+except ImportError:
+    from llama_stack.providers.datatypes import RemoteProviderSpec
+
+    USE_LEGACY = False
 
 from llama_stack_provider_trustyai_fms.provider import (
     get_provider_spec,
@@ -9,26 +18,46 @@
 class TestProviderFunctions:
     def test_get_provider_spec(self):
         spec = get_provider_spec()
-
         assert isinstance(spec, ProviderSpec)
         assert spec.api == Api.safety
-        assert isinstance(spec.adapter, AdapterSpec)
-        assert spec.adapter.adapter_type == "trustyai_fms"
-        assert (
-            spec.adapter.config_class
-            == "llama_stack_provider_trustyai_fms.config.FMSSafetyProviderConfig"
-        )
-        assert spec.adapter.module == "llama_stack_provider_trustyai_fms"
+
+        if USE_LEGACY:
+            assert isinstance(spec.adapter, AdapterSpec)
+            assert spec.adapter.adapter_type == "trustyai_fms"
+            assert (
+                spec.adapter.config_class
+                == "llama_stack_provider_trustyai_fms.config.FMSSafetyProviderConfig"
+            )
+            assert spec.adapter.module == "llama_stack_provider_trustyai_fms"
+        else:
+            assert isinstance(spec, RemoteProviderSpec)
+            assert spec.provider_type == "remote::trustyai_fms"
+            assert spec.adapter_type == "trustyai_fms"
+            assert (
+                spec.config_class
+                == "llama_stack_provider_trustyai_fms.config.FMSSafetyProviderConfig"
+            )
+            assert spec.module == "llama_stack_provider_trustyai_fms"
 
     def test_get_shields_provider_spec(self):
         spec = get_shields_provider_spec()
-
         assert isinstance(spec, ProviderSpec)
         assert spec.api == Api.shields
-        assert isinstance(spec.adapter, AdapterSpec)
-        assert spec.adapter.adapter_type == "trustyai_fms"
-        assert (
-            spec.adapter.config_class
-            == "llama_stack_provider_trustyai_fms.config.FMSSafetyProviderConfig"
-        )
-        assert spec.adapter.module == "llama_stack_provider_trustyai_fms"
+
+        if USE_LEGACY:
+            assert isinstance(spec.adapter, AdapterSpec)
+            assert spec.adapter.adapter_type == "trustyai_fms"
+            assert (
+                spec.adapter.config_class
+                == "llama_stack_provider_trustyai_fms.config.FMSSafetyProviderConfig"
+            )
+            assert spec.adapter.module == "llama_stack_provider_trustyai_fms"
+        else:
+            assert isinstance(spec, RemoteProviderSpec)
+            assert spec.provider_type == "remote::trustyai_fms"
+            assert spec.adapter_type == "trustyai_fms"
+            assert (
+                spec.config_class
+                == "llama_stack_provider_trustyai_fms.config.FMSSafetyProviderConfig"
+            )
+            assert spec.module == "llama_stack_provider_trustyai_fms"