frogml-core 0.0.96__py3-none-any.whl → 0.0.98__py3-none-any.whl

frogml_core/__init__.py

@@ -1,7 +1,7 @@
 """Top-level package for frogml."""
 
 __author__ = "jfrog"
-__version__ = "0.0.96"
+__version__ = "0.0.98"
 
 from frogml_core.inner.di_configuration import wire_dependencies
 from frogml_core.model.model_version_tracking import (  # noqa: F401,E501

frogml_core/clients/batch_job_management/client.py

@@ -36,6 +36,10 @@ from frogml_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
     StartBatchJobResponse,
     StartWarmupJobRequest,
     StartWarmupJobResponse,
+    UpdateTasksDetailsResponse,
+    UpdateTasksDetailsRequest,
+    BatchTaskDetails,
+    InputFileDetails,
 )
 from frogml_proto.qwak.batch_job.v1.batch_job_service_pb2_grpc import (
     BatchJobManagementServiceStub,
@@ -67,6 +71,7 @@ from .results import (
     GetExecutionReportResult,
     StartExecutionResult,
 )
+from typing import List
 
 CLIENT_TIMEOUT = 180  # Seconds
 
@@ -542,3 +547,37 @@ class BatchJobManagerClient:
 
         except grpc.RpcError as e:
             raise FrogmlException(f"Failed to get batch job details, error is: {e}")
+
+    def update_task_details(self, task_id: str, input_file_paths: List[str]):
+        """
+        Update task details for a specific task using the provided list of input file paths.
+
+        Args:
+            task_id: The ID of the task to update.
+            input_file_paths: A list of paths for the input files associated with the task.
+        """
+        # Create InputFileDetails messages for each input file path
+        input_files_details = [
+            InputFileDetails(path=file_path) for file_path in input_file_paths
+        ]
+
+        # Create BatchTaskDetails message for the single task
+        batch_task_details = BatchTaskDetails(
+            task_id=task_id, input_files_details=input_files_details
+        )
+
+        update_request = UpdateTasksDetailsRequest(
+            tasks_details=[batch_task_details]  # Wrap in a list
+        )
+
+        try:
+            response: UpdateTasksDetailsResponse = (
+                self.batch_job_management.UpdateTasksDetails(
+                    update_request, timeout=CLIENT_TIMEOUT
+                )
+            )
+            return response
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to update task details for task '{task_id}', error is: {e}"
+            )

frogml_core/clients/instance_template/client.py

@@ -50,10 +50,12 @@ class InstanceTemplateManagementClient:
 
     def list_instance_templates(self) -> List[InstanceTemplateSpec]:
         try:
-            result: ListInstanceTemplatesResponse = self._instance_template_service.ListInstanceTemplates(
-                ListInstanceTemplatesRequest(
-                    optional_instance_filter=InstanceFilter(
-                        instance_type_filter=InstanceTypeFilter.INSTANCE_TYPE_FILTER_ALL
+            result: ListInstanceTemplatesResponse = (
+                self._instance_template_service.ListInstanceTemplates(
+                    ListInstanceTemplatesRequest(
+                        optional_instance_filter=InstanceFilter(
+                            instance_type_filter=InstanceTypeFilter.INSTANCE_TYPE_FILTER_ALL
+                        )
                     )
                 )
             )

frogml_core/clients/model_version_manager/client.py

@@ -2,7 +2,7 @@ from platform import python_version
 from typing import List, Optional, Dict, Union, Tuple
 
 from dependency_injector.wiring import Provide
-from frogml_storage.entity_manifest import Artifact
+from frogml_storage.models.entity_manifest import Artifact
 from grpc import RpcError
 
 from frogml_core.exceptions import FrogmlException
@@ -210,9 +210,9 @@ class ModelVersionManagerClient:
         Optional[List[ArtifactProto]],
         List[ArtifactProto],
     ]:
-        model_artifact_proto: List[
-            ArtifactProto
-        ] = ProtoUtils.convert_artifacts_to_artifacts_proto(model_artifact)
+        model_artifact_proto: List[ArtifactProto] = (
+            ProtoUtils.convert_artifacts_to_artifacts_proto(model_artifact)
+        )
         dependency_artifacts_proto = (
             ProtoUtils.convert_artifacts_to_artifacts_proto(dependency_artifacts)
             if dependency_artifacts

frogml_core/clients/prompt_manager/model_descriptor_mapper.py

@@ -57,10 +57,10 @@ class ModelDescriptorMapper:
         model_id: str, openai_chat_params: ProtoOpenAIChatModelParams
     ) -> OpenAIChat:
         p = openai_chat_params
-        _tool_choice: Union[
-            str, ChatCompletionNamedToolChoiceParam
-        ] = ModelDescriptorMapper._from_tool_choice(
-            openai_chat_params=openai_chat_params
+        _tool_choice: Union[str, ChatCompletionNamedToolChoiceParam] = (
+            ModelDescriptorMapper._from_tool_choice(
+                openai_chat_params=openai_chat_params
+            )
         )
         _tools: List[ChatCompletionToolParam] = []
 

frogml_core/feature_store/_common/artifact_utils.py

@@ -45,9 +45,9 @@ class ArtifactsUploader:
         featureset_name: str,
         __instance_module_path__: str,
     ) -> Optional[ArtifactSpec]:
-        transformation_functions: Optional[
-            List[Callable[..., Any]]
-        ] = transformation.get_functions()
+        transformation_functions: Optional[List[Callable[..., Any]]] = (
+            transformation.get_functions()
+        )
         if transformation_functions is not None and transformation_functions:
             return ArtifactSpec(
                 artifact_name=featureset_name,

frogml_core/feature_store/data_sources/batch/athena.py

@@ -130,9 +130,9 @@ class AthenaSource(JdbcSource):
             )
         )
 
-        time_partition_columns: Optional[
-            TimePartitionColumns
-        ] = AthenaSource._extract_partition_column(proto_athena_source)
+        time_partition_columns: Optional[TimePartitionColumns] = (
+            AthenaSource._extract_partition_column(proto_athena_source)
+        )
         workgroup: Optional[str] = (
             proto_athena_source.workgroup
             if proto_athena_source.HasField("workgroup")

frogml_core/feature_store/feature_sets/context.py

@@ -5,9 +5,5 @@ from typing import Union
 
 @dataclass
 class Context:
-    start_time: Union[
-        str, datetime
-    ] = "${qwak_ingestion_start_timestamp}"  # todo mlops-2312 - rename?
-    end_time: Union[
-        str, datetime
-    ] = "${qwak_ingestion_end_timestamp}"  # todo mlops-2312 - rename?
+    start_time: Union[str, datetime] = "${qwak_ingestion_start_timestamp}"
+    end_time: Union[str, datetime] = "${qwak_ingestion_end_timestamp}"

frogml_core/feature_store/feature_sets/streaming.py

@@ -131,9 +131,9 @@ def feature_set(
             offline_scheduling_policy=offline_scheduling_policy,
         )
 
-        streaming_backfill: Optional[
-            StreamingBackfill
-        ] = StreamingBackfill.get_streaming_backfill_from_function(function=function)
+        streaming_backfill: Optional[StreamingBackfill] = (
+            StreamingBackfill.get_streaming_backfill_from_function(function=function)
+        )
 
         fs_name = name or function.__name__
         streaming_feature_set = StreamingFeatureSet(

frogml_core/feature_store/feature_sets/streaming_backfill.py

@@ -230,7 +230,7 @@ class StreamingBackfill:
 
     @staticmethod
     def _get_normalized_backfill_sources_spec(
-        data_sources: Union[List[str], List[DataSourceBackfillSpec]]
+        data_sources: Union[List[str], List[DataSourceBackfillSpec]],
     ) -> List[DataSourceBackfillSpec]:
         # reformat all data source specs to 'DataSourceBackfillSpec'
         return [

frogml_core/feature_store/online/client.py

@@ -197,12 +197,12 @@ class OnlineClient:
         )
         ordered_entities = [entity[0] for entity in ordered_entities_tuple]
 
-        request_chunks: List[
-            Tuple[RequestedEntitiesMatrix, pd.DataFrame]
-        ] = OnlineClient._split_entities(
-            entity_names=ordered_entities,
-            population_df=df,
-            max_entities_per_split=max_keys_per_request,
+        request_chunks: List[Tuple[RequestedEntitiesMatrix, pd.DataFrame]] = (
+            OnlineClient._split_entities(
+                entity_names=ordered_entities,
+                population_df=df,
+                max_entities_per_split=max_keys_per_request,
+            )
         )
 
         results: List[pd.DataFrame] = []

frogml_core/feature_store/sinks/streaming/factory.py

@@ -25,7 +25,7 @@ class StreamingSinkFactory:
     def get_streaming_sink(proto_streaming_sink: ProtoStreamingSink) -> BaseSink:
         sink_type = proto_streaming_sink.WhichOneof("sink_type")
 
-        auth_conf: BaseAuthentication
+        auth_conf: BaseAuthentication  # noqa: F842
         if sink_type == "kafka_sink":
             proto_kafka_sink: ProtoKafkaSink = proto_streaming_sink.kafka_sink
             auth_configuration: BaseAuthentication = cast(

frogml_core/inner/build_logic/phases/phase_010_fetch_model/fetch_strategy_manager/strategy/git/git_strategy.py

@@ -161,9 +161,9 @@ def make_ssh_key_file(git_ssh_key: str) -> str:
 
 
 def add_ssh_file_to_env(ssh_key_file_path: str) -> None:
-    os.environ[
-        "GIT_SSH_COMMAND"
-    ] = f"ssh -i {ssh_key_file_path} -o StrictHostKeyChecking=no"
+    os.environ["GIT_SSH_COMMAND"] = (
+        f"ssh -i {ssh_key_file_path} -o StrictHostKeyChecking=no"
+    )
     os.environ["GIT_SSH"] = f"ssh -i {ssh_key_file_path} -o StrictHostKeyChecking=no"
 
 

frogml_core/inner/di_configuration/account.py

@@ -3,8 +3,8 @@ import os
 from dataclasses import dataclass
 from typing import Optional
 
-from frogml_storage.authentication._authentication_utils import get_credentials
-from frogml_storage.cli._login_cli import login as frogml_login
+from frogml_storage.authentication.utils import get_credentials
+from frogml_storage.authentication.login import frogml_login
 from frogml_storage.constants import CONFIG_FILE_PATH
 
 from frogml_core.exceptions import FrogmlLoginException

frogml_core/inner/tool/auth.py

@@ -1,9 +1,8 @@
 from typing import Optional, Union, cast
 
 import requests
-from frogml_storage._utils import BearerAuth
-from frogml_storage.authentication._authentication_utils import get_credentials
-from frogml_storage.authentication.models._auth_config import AuthConfig
+from frogml_storage.authentication.models import AuthConfig, BearerAuth
+from frogml_storage.authentication.utils import get_credentials
 from requests.auth import AuthBase
 from typing_extensions import Self
 

frogml_core/llmops/provider/openai/provider.py

@@ -30,9 +30,9 @@ class OpenAIProvider:
         self.client = OpenAIClient()
 
     def _get_random_openai_api_key(self) -> Optional[str]:
-        openai_api_keys: List[
-            OpenAIApiKeySystemSecret
-        ] = IntegrationUtils().get_openai_api_keys()
+        openai_api_keys: List[OpenAIApiKeySystemSecret] = (
+            IntegrationUtils().get_openai_api_keys()
+        )
         if len(openai_api_keys) == 0:
             return None
 

frogml_core/model/tools/adapters/output.py

@@ -46,7 +46,7 @@ def get_output_adapter(
 
             first_result = (
                 return_result[0]
-                if type(return_result) == list and len(return_result) > 0
+                if isinstance(return_result, list) and len(return_result) > 0
                 else return_result
             )
             if issubclass(type(first_result), Message):

frogml_core/tools/logger/logger.py

@@ -276,7 +276,7 @@ def set_file_handler_log_file(
     logger: logging.Logger, handler_name: str, log_file: Path
 ):
     existing_handler = get_handler_from_logger(logger, handler_name)
-    if type(existing_handler) != RotatingFileHandler:
+    if not isinstance(existing_handler, RotatingFileHandler):
         raise FrogmlException(
             f"Error in setting log file. Error message: handler of name {handler_name} is not a file logger handler"
         )

frogml_core/utils/proto_utils.py

@@ -1,6 +1,6 @@
 from typing import List
 
-from frogml_storage.entity_manifest import Artifact
+from frogml_storage.models.entity_manifest import Artifact
 
 from frogml_proto.jfml.model_version.v1.artifact_pb2 import (
     Artifact as ArtifactProto,