digitalkin 0.2.6__py3-none-any.whl → 0.2.7__py3-none-any.whl

digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.2.6"
+    __version__ = "0.2.7"

digitalkin/grpc_servers/module_servicer.py

@@ -244,7 +244,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         # Get input schema if available
         try:
             # Convert schema to proto format
-            input_schema_proto = self.module_class.get_input_format(request.llm_format)
+            input_schema_proto = self.module_class.get_input_format(llm_format=request.llm_format)
             input_format_struct = json_format.Parse(
                 text=input_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -280,7 +280,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         # Get output schema if available
         try:
             # Convert schema to proto format
-            output_schema_proto = self.module_class.get_output_format(request.llm_format)
+            output_schema_proto = self.module_class.get_output_format(llm_format=request.llm_format)
             output_format_struct = json_format.Parse(
                 text=output_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -316,7 +316,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         # Get setup schema if available
         try:
             # Convert schema to proto format
-            setup_schema_proto = self.module_class.get_setup_format(request.llm_format)
+            setup_schema_proto = self.module_class.get_setup_format(llm_format=request.llm_format)
             setup_format_struct = json_format.Parse(
                 text=setup_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -332,3 +332,39 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
             success=True,
             setup_schema=setup_format_struct,
         )
+
+    def GetModuleSecret(  # noqa: N802
+        self,
+        request: information_pb2.GetModuleSecretRequest,
+        context: grpc.ServicerContext,
+    ) -> information_pb2.GetModuleSecretResponse:
+        """Get information about the module's secrets.
+
+        Args:
+            request: The get module secret request.
+            context: The gRPC context.
+
+        Returns:
+            A response with the module's secret schema.
+        """
+        logger.info("GetModuleSecret called for module: '%s'", self.module_class.__name__)
+
+        # Get secret schema if available
+        try:
+            # Convert schema to proto format
+            secret_schema_proto = self.module_class.get_secret_format(llm_format=request.llm_format)
+            secret_format_struct = json_format.Parse(
+                text=secret_schema_proto,
+                message=struct_pb2.Struct(),  # pylint: disable=no-member
+                ignore_unknown_fields=True,
+            )
+        except NotImplementedError as e:
+            logger.warning(e)
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_details(e)
+            return information_pb2.GetModuleSecretResponse()
+
+        return information_pb2.GetModuleSecretResponse(
+            success=True,
+            secret_schema=secret_format_struct,
+        )

digitalkin/modules/_base_module.py

@@ -17,6 +17,7 @@ from digitalkin.services.registry.registry_strategy import RegistryStrategy
 from digitalkin.services.services_config import ServicesConfig, ServicesStrategy
 from digitalkin.services.snapshot.snapshot_strategy import SnapshotStrategy
 from digitalkin.services.storage.storage_strategy import StorageStrategy
+from digitalkin.utils.llm_ready_schema import llm_ready_schema
 
 
 class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretModelT]):
@@ -73,7 +74,7 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         return self._status
 
     @classmethod
-    def get_secret_format(cls, llm_format: bool) -> str:  # noqa: FBT001
+    def get_secret_format(cls, *, llm_format: bool) -> str:
         """Get the JSON schema of the secret format model.
 
         Raises:
@@ -84,13 +85,13 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         """
         if cls.secret_format is not None:
             if llm_format:
-                return json.dumps(cls.secret_format, indent=2)
+                return json.dumps(llm_ready_schema(cls.secret_format), indent=2)
             return json.dumps(cls.secret_format.model_json_schema(), indent=2)
         msg = f"{cls.__name__}' class does not define a 'secret_format'."
         raise NotImplementedError(msg)
 
     @classmethod
-    def get_input_format(cls, llm_format: bool) -> str:  # noqa: FBT001
+    def get_input_format(cls, *, llm_format: bool) -> str:
         """Get the JSON schema of the input format model.
 
         Raises:
@@ -101,13 +102,13 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         """
         if cls.input_format is not None:
             if llm_format:
-                return json.dumps(cls.input_format, indent=2)
+                return json.dumps(llm_ready_schema(cls.input_format), indent=2)
             return json.dumps(cls.input_format.model_json_schema(), indent=2)
         msg = f"{cls.__name__}' class does not define an 'input_format'."
         raise NotImplementedError(msg)
 
     @classmethod
-    def get_output_format(cls, llm_format: bool) -> str:  # noqa: FBT001
+    def get_output_format(cls, *, llm_format: bool) -> str:
         """Get the JSON schema of the output format model.
 
         Raises:
@@ -118,13 +119,13 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         """
         if cls.output_format is not None:
             if llm_format:
-                return json.dumps(cls.output_format, indent=2)
+                return json.dumps(llm_ready_schema(cls.output_format), indent=2)
             return json.dumps(cls.output_format.model_json_schema(), indent=2)
         msg = "'%s' class does not define an 'output_format'."
         raise NotImplementedError(msg)
 
     @classmethod
-    def get_setup_format(cls, llm_format: bool) -> str:  # noqa: FBT001
+    def get_setup_format(cls, *, llm_format: bool) -> str:
         """Gets the JSON schema of the setup format model.
 
         Raises:
@@ -135,7 +136,7 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         """
         if cls.setup_format is not None:
             if llm_format:
-                return json.dumps(cls.setup_format, indent=2)
+                return json.dumps(llm_ready_schema(cls.setup_format), indent=2)
             return json.dumps(cls.setup_format.model_json_schema(), indent=2)
         msg = "'%s' class does not define an 'setup_format'."
         raise NotImplementedError(msg)

digitalkin/utils/llm_ready_schema.py

@@ -0,0 +1,75 @@
+"""LLM format schema for Pydantic models.
+
+This module provides functionality to generate JSON schemas for Pydantic models ready for LLMs.
+"""
+
+import copy
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
+
+
+class CustomOrderSchema(GenerateJsonSchema):
+    """Custom schema generator to sort keys in a specific order."""
+
+    def sort(self, value: JsonSchemaValue, parent_key: str | None = None) -> JsonSchemaValue:  # noqa: ARG002
+        """Sort the keys of the schema in a specific order.
+
+        Args:
+            value: The schema value to sort.
+            parent_key: The parent key of the schema value.
+
+        Returns:
+            The sorted schema value.
+        """
+        if isinstance(value, dict):
+            # Define your preferred order
+            preferred = ["title", "description", "type", "examples", "properties"]
+            # Collect all keys, putting preferred ones first
+            keys = preferred + [k for k in value if k not in preferred]
+            # Recurse for each value
+            return {k: self.sort(value[k], k) for k in keys if k in value}
+        if isinstance(value, list):
+            return [self.sort(v) for v in value]
+        return value
+
+
+def inline_refs(schema: dict) -> dict:
+    """Recursively resolve and inline all $ref in the schema.
+
+    Args:
+        schema: The JSON schema to inline.
+
+    Returns:
+        The inlined JSON schema.
+    """
+    schema = copy.deepcopy(schema)
+    defs = schema.pop("$defs", {})
+
+    def _resolve(obj: Any) -> Any:  # noqa: ANN401
+        if isinstance(obj, dict):
+            if "$ref" in obj:
+                ref = obj["$ref"]
+                if ref.startswith("#/$defs/"):
+                    key = ref.split("/")[-1]
+                    return _resolve(defs[key])
+            return {k: _resolve(v) for k, v in obj.items()}
+        if isinstance(obj, list):
+            return [_resolve(item) for item in obj]
+        return obj
+
+    return _resolve(schema)
+
+
+def llm_ready_schema(model: type[BaseModel]) -> dict:
+    """Convert a Pydantic model to a JSON schema ready for LLMs.
+
+    Args:
+        model: The Pydantic model to convert.
+
+    Returns:
+        The JSON schema as a dictionary.
+    """
+    schema = model.model_json_schema(schema_generator=CustomOrderSchema)
+    return inline_refs(schema)

{digitalkin-0.2.6.dist-info → digitalkin-0.2.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.2.6
+Version: 0.2.7
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International

{digitalkin-0.2.6.dist-info → digitalkin-0.2.7.dist-info}/RECORD

@@ -7,13 +7,13 @@ base_server/mock/__init__.py,sha256=YZFT-F1l_TpvJYuIPX-7kTeE1CfOjhx9YmNRXVoi-jQ,
 base_server/mock/mock_pb2.py,sha256=sETakcS3PAAm4E-hTCV1jIVaQTPEAIoVVHupB8Z_k7Y,1843
 base_server/mock/mock_pb2_grpc.py,sha256=BbOT70H6q3laKgkHfOx1QdfmCS_HxCY4wCOX84YAdG4,3180
 digitalkin/__init__.py,sha256=7LLBAba0th-3SGqcpqFO-lopWdUkVLKzLZiMtB-mW3M,162
-digitalkin/__version__.py,sha256=oM4jNCAYWlE0yUk_rpWiJCYE3Pyq4cZrXpraEmj9p-E,190
+digitalkin/__version__.py,sha256=sq3ZYGvIa5rczwcukAFcGJdR0MZ8h0cIPYmDlhbjOP4,190
 digitalkin/logger.py,sha256=9cDgyJV2QXXT8F--xRODFlZyDgjuTTXNdpCU3GdqCsk,382
 digitalkin/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 digitalkin/grpc_servers/__init__.py,sha256=0cJBlwipSmFdXkyH3T0i6OJ1WpAtNsZgYX7JaSnkbtg,804
 digitalkin/grpc_servers/_base_server.py,sha256=ec4xmgAuOMVg45a63O_PEa2T7mI4tJ6boxcXauFyZ5g,18649
 digitalkin/grpc_servers/module_server.py,sha256=0jC5MN-H4MGA9O9J87gnlM8EHcbrkb0chuCm8yZh7gI,10232
-digitalkin/grpc_servers/module_servicer.py,sha256=byqRpp9m8HS3NUUtKBWFXOncpUWGDW_xsyddA2qwaW4,12294
+digitalkin/grpc_servers/module_servicer.py,sha256=gm1dhmHhBGqY2xlhAp5i15O6waOA1mDHbFJxsT2RUT8,13662
 digitalkin/grpc_servers/registry_server.py,sha256=PmWaH4Xmg5Sj7NtFVLBNTOzkOfqo7dw_qyVBnaW4jy4,2238
 digitalkin/grpc_servers/registry_servicer.py,sha256=mCAjNhdMq5DozZMEPsJK__DIxePEYxSWV-gAq-Xctk4,16469
 digitalkin/grpc_servers/utils/exceptions.py,sha256=I00OM8p8up20He4dU1fiHsvdLj1DymjR_UmoeUm2MSA,785
@@ -29,7 +29,7 @@ digitalkin/models/services/__init__.py,sha256=HsW7MUGFPvH7Ri28WN4BHHBfEQk5dzU_9F
 digitalkin/models/services/cost.py,sha256=QTEuFD6xz62nob0z4ksE-INJWcZ-iFiuNW5mvXhpFes,1599
 digitalkin/models/services/storage.py,sha256=cYTVIriGKiprF9OerhSxmc_jM6fUTVwmeon1yQCinkE,143
 digitalkin/modules/__init__.py,sha256=ppYARmhvdVi55ofC0QZerIempSlcJYDeCXhcl4qXObw,278
-digitalkin/modules/_base_module.py,sha256=iVPRlooAGIW4_3lN9loUPe1q6uf62-vSsEE5THvn8eU,9007
+digitalkin/modules/_base_module.py,sha256=WU2oUW4-FjFtieuyRfXv3R3P4N1xT78zzMWrSTdhmAg,9090
 digitalkin/modules/archetype_module.py,sha256=T2Ehj7EpAC2MO9WQbJv39hqRw7rh3exhVZTEL3JPM8U,421
 digitalkin/modules/job_manager.py,sha256=QHcrm3F99I9tlNzXefkyhkkwPITf_H1bEb_obMDbUic,6177
 digitalkin/modules/tool_module.py,sha256=86g0M1wHZ1ReIc7AkKfyjnlGN2QYJBGxrEQpKVlyrZI,421
@@ -67,12 +67,13 @@ digitalkin/services/storage/grpc_storage.py,sha256=eSadiI3JuveItx8LdiC3GbjbjnB8G
 digitalkin/services/storage/storage_strategy.py,sha256=vGo4aYkEp_GZV11m7vd-xY_Z3gVa5K0gMTzbj2Au_3o,6600
 digitalkin/utils/__init__.py,sha256=sJnY-ZUgsjMfojAjONC1VN14mhgIDnzyOlGkw21rRnM,28
 digitalkin/utils/arg_parser.py,sha256=3YyI6oZhhrlTmPTrzlwpQzbCNWDFAT3pggcLxNtJoc0,4388
-digitalkin-0.2.6.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
+digitalkin/utils/llm_ready_schema.py,sha256=JjMug_lrQllqFoanaC091VgOqwAd-_YzcpqFlS7p778,2375
+digitalkin-0.2.7.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
 modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 modules/minimal_llm_module.py,sha256=W-E3OrRbAsRJ6hvSeTU8pzmacdJC_PbcWfDapRv5A1A,5617
 modules/storage_module.py,sha256=bu52lW4RFcWB8VFDhrpBFfCaTSkVL6so3zrkfW4LO9E,6270
 modules/text_transform_module.py,sha256=1KaA7abwxltKKtbmiW1rkkIK3BTYFPegUq54px0LOQs,7277
-digitalkin-0.2.6.dist-info/METADATA,sha256=MSrxM2ATU0easkm4lkbVJdPbUzoH0xDgbmQ622ACwfM,29125
-digitalkin-0.2.6.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
-digitalkin-0.2.6.dist-info/top_level.txt,sha256=5_5e35inSM5YfWNZE21p5wGBojiVtQQML_WzbEk4BRU,31
-digitalkin-0.2.6.dist-info/RECORD,,
+digitalkin-0.2.7.dist-info/METADATA,sha256=1-9MQvIcY6e4S5Cfzea9Nk2lHFx7coGQZ935DzJ00zg,29125
+digitalkin-0.2.7.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+digitalkin-0.2.7.dist-info/top_level.txt,sha256=5_5e35inSM5YfWNZE21p5wGBojiVtQQML_WzbEk4BRU,31
+digitalkin-0.2.7.dist-info/RECORD,,

{digitalkin-0.2.6.dist-info → digitalkin-0.2.7.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (79.0.1)
+Generator: setuptools (80.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any