frogml-core 0.0.1__py3-none-any.whl

frogml/clients/model_version_manager/client.py

@@ -0,0 +1,108 @@
+from dependency_injector.wiring import Provide
+from grpc import RpcError
+
+from _frogml_proto.jfml.model_version.v1.model_repository_spec_pb2 import (
+    ModelRepositorySpec,
+)
+from _frogml_proto.jfml.model_version.v1.model_version_framework_pb2 import (
+    ModelVersionFramework,
+)
+from _frogml_proto.jfml.model_version.v1.model_version_manager_service_pb2 import (
+    CreateModelVersionRequest,
+    CreateModelVersionResponse,
+)
+from _frogml_proto.jfml.model_version.v1.model_version_manager_service_pb2_grpc import (
+    ModelVersionManagerServiceStub,
+)
+from _frogml_proto.jfml.model_version.v1.model_version_pb2 import ModelVersionSpec
+from frogml.exceptions import FrogmlException
+from frogml.inner.di_configuration import FrogmlContainer
+
+
+class ModelVersionManagerClient:
+    """
+    Used for interacting with the model version manager's endpoints
+    """
+
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self.__model_version_manager_stub = ModelVersionManagerServiceStub(grpc_channel)
+
+    @staticmethod
+    def __build_create_model_request(
+        project_key: str,
+        repository_key: str,
+        model_id: str,
+        model_version_name: str,
+        model_version_framework: ModelVersionFramework,
+        dry_run: bool,
+    ) -> CreateModelVersionRequest:
+        return CreateModelVersionRequest(
+            dry_run=dry_run,
+            model_version=ModelVersionSpec(
+                repository_spec=ModelRepositorySpec(
+                    project_key=project_key,
+                    repository_key=repository_key,
+                    model_id=model_id,
+                ),
+                name=model_version_name,
+                framework=model_version_framework,
+            ),
+        )
+
+    def validate_create_model_version(
+        self,
+        project_key: str,
+        repository_key: str,
+        model_name: str,
+        model_version_name: str,
+        model_version_framework: ModelVersionFramework,
+    ):
+        try:
+            create_model_request = self.__build_create_model_request(
+                project_key=project_key,
+                repository_key=repository_key,
+                model_id=model_name,
+                model_version_name=model_version_name,
+                model_version_framework=model_version_framework,
+                dry_run=True,
+            )
+            self.__model_version_manager_stub.CreateModelVersion(create_model_request)
+        except RpcError as e:
+            message = f"Failed to validate model version, details [{e.details()}]"
+            raise FrogmlException(message)
+
+        except Exception as e:
+            message = f"Failed to validate model version, details [{e}]"
+            raise FrogmlException(message)
+
+    def create_model_version(
+        self,
+        project_key: str,
+        repository_key: str,
+        model_name: str,
+        model_version_name: str,
+        model_version_framework: ModelVersionFramework,
+    ) -> CreateModelVersionResponse:
+        try:
+            create_model_request = self.__build_create_model_request(
+                project_key=project_key,
+                repository_key=repository_key,
+                model_id=model_name,
+                model_version_name=model_version_name,
+                model_version_framework=model_version_framework,
+                dry_run=False,
+            )
+            create_model_version_response: CreateModelVersionResponse = (
+                self.__model_version_manager_stub.CreateModelVersion(
+                    create_model_request
+                )
+            )
+
+            return create_model_version_response
+        except RpcError as e:
+            message = f"Failed to validate model version, details [{e.details()}]"
+            raise FrogmlException(message)
+
+        except Exception as e:
+            message = f"Failed to validate model version, details [{e}]"
+            raise FrogmlException(message)

frogml/clients/project/client.py

@@ -0,0 +1,60 @@
+import grpc
+from dependency_injector.wiring import Provide, inject
+
+from _frogml_proto.qwak.projects.projects_pb2 import (
+    CreateProjectRequest,
+    DeleteProjectRequest,
+    GetProjectRequest,
+    ListProjectsRequest,
+)
+from _frogml_proto.qwak.projects.projects_pb2_grpc import ProjectsManagementServiceStub
+from frogml.exceptions import FrogmlException
+from frogml.inner.di_configuration import FrogmlContainer
+
+
+class ProjectsManagementClient:
+    """
+    Used for interacting with Project Management endpoints
+    """
+
+    @inject
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._projects_management_service = ProjectsManagementServiceStub(grpc_channel)
+
+    def list_projects(self):
+        try:
+            return self._projects_management_service.ListProjects(ListProjectsRequest())
+
+        except grpc.RpcError as e:
+            raise FrogmlException(f"Failed to list projects, error is {e.details()}")
+
+    def create_project(self, project_name, project_description):
+        try:
+            return self._projects_management_service.CreateProject(
+                CreateProjectRequest(
+                    project_name=project_name, project_description=project_description
+                )
+            )
+
+        except grpc.RpcError as e:
+            raise FrogmlException(f"Failed to create project, error is {e.details()}")
+
+    def delete_project(self, project_id):
+        try:
+            return self._projects_management_service.DeleteProject(
+                DeleteProjectRequest(project_id=project_id)
+            )
+
+        except grpc.RpcError as e:
+            raise FrogmlException(f"Failed to delete project, error is {e.details()}")
+
+    def get_project(self, project_id: str = "", project_name: str = ""):
+        try:
+            return self._projects_management_service.GetProject(
+                GetProjectRequest(project_id=project_id, project_name=project_name)
+            )
+
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to fetch models associated with project id {project_id}, error is {e.details()}"
+            )

frogml/clients/prompt_manager/model_descriptor_mapper.py

@@ -0,0 +1,196 @@
+from typing import List, Optional, Union
+
+from google.protobuf import json_format
+from google.protobuf.struct_pb2 import ListValue, Struct
+
+from _frogml_proto.qwak.model_descriptor.open_ai_descriptor_pb2 import (
+    OpenAICapabilities as ProtoOpenAICapabilities,
+)
+from _frogml_proto.qwak.model_descriptor.open_ai_descriptor_pb2 import (
+    OpenAIChatAPI as ProtoOpenAIChatAPI,
+)
+from _frogml_proto.qwak.model_descriptor.open_ai_descriptor_pb2 import (
+    OpenAIChatModelParams as ProtoOpenAIChatModelParams,
+)
+from _frogml_proto.qwak.model_descriptor.open_ai_descriptor_pb2 import (
+    OpenAIModelDescriptor as ProtoOpenAIModelDescriptor,
+)
+from _frogml_proto.qwak.model_descriptor.open_ai_descriptor_pb2 import (
+    ToolChoice as ProtoToolChoice,
+)
+from _frogml_proto.qwak.model_descriptor.open_ai_descriptor_pb2 import (
+    Tools as ProtoTools,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
+    PromptModelDescriptor as ProtoPromptModelDescriptor,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
+    PromptOpenAIProvider as ProtoPromptOpenAIProvider,
+)
+from frogml.exceptions import FrogmlException
+from frogml.llmops.generation.chat.openai.types.chat.chat_completion_named_tool_choice_param import (
+    ChatCompletionNamedToolChoiceParam,
+)
+from frogml.llmops.generation.chat.openai.types.chat.chat_completion_tool_param import (
+    ChatCompletionToolParam,
+)
+from frogml.llmops.model.descriptor import ModelDescriptor, OpenAIChat
+
+
+class ModelDescriptorMapper:
+    @staticmethod
+    def _from_tool_choice(
+        openai_chat_params: ProtoOpenAIChatModelParams,
+    ) -> Optional[Union[str, ChatCompletionNamedToolChoiceParam]]:
+        _tool_choice: Optional[Union[str, ChatCompletionNamedToolChoiceParam]] = None
+        if openai_chat_params.HasField("tool_choice"):
+            if openai_chat_params.tool_choice.WhichOneof("value_type") == "json":
+                _tool_choice = json_format.MessageToDict(
+                    openai_chat_params.tool_choice.json
+                )
+            elif openai_chat_params.tool_choice.WhichOneof("value_type") == "literal":
+                _tool_choice = openai_chat_params.tool_choice.literal
+        return _tool_choice
+
+    @staticmethod
+    def from_openai_chat_capability(
+        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
+        )
+        _tools: List[ChatCompletionToolParam] = []
+
+        if p.HasField("tools_spec"):
+            for tool in p.tools_spec.tools:
+                _tools.append(json_format.MessageToDict(tool))
+
+        return OpenAIChat(
+            model_id=model_id,
+            frequency_penalty=(
+                p.frequency_penalty if p.HasField("frequency_penalty") else None
+            ),
+            logit_bias=(
+                {k: int(v) for k, v in p.logit_bias.items()}
+                if p.HasField("logit_bias")
+                else None
+            ),
+            logprobs=p.logprobs if p.HasField("logprobs") else None,
+            max_tokens=p.max_tokens if p.HasField("max_tokens") else None,
+            n=p.n if p.HasField("n") else None,
+            presence_penalty=(
+                p.presence_penalty if p.HasField("presence_penalty") else None
+            ),
+            response_format=(
+                p.response_format if p.HasField("response_format") else None
+            ),  # noqa
+            seed=p.seed if p.HasField("seed") else None,
+            stop=[_ for _ in p.stop] if p.HasField("stop") else None,
+            temperature=p.temperature if p.HasField("temperature") else None,
+            top_p=p.top_p if p.HasField("top_p") else None,
+            top_logprobs=p.top_logprobs if p.HasField("top_logprobs") else None,
+            user=p.user if p.HasField("user") else None,
+            tool_choice=_tool_choice if p.HasField("tool_choice") else None,
+            tools=_tools if p.HasField("tools_spec") else None,
+        )
+
+    @staticmethod
+    def from_prompt_openai_provider(
+        open_ai_provider: ProtoPromptOpenAIProvider,
+    ) -> ModelDescriptor:
+        descriptor: ProtoOpenAIModelDescriptor = (
+            open_ai_provider.open_ai_model_descriptor
+        )
+        model_id: str = descriptor.model_id
+
+        if descriptor.capabilities.WhichOneof("optional_chat"):
+            return ModelDescriptorMapper.from_openai_chat_capability(
+                model_id=model_id,
+                openai_chat_params=descriptor.capabilities.chat_api.chat_params,
+            )
+        else:
+            raise FrogmlException(
+                f"Got unsupported openai capability: {repr(open_ai_provider)}"
+            )
+
+    @staticmethod
+    def from_prompt_model_descriptor(
+        model_descriptor: ProtoPromptModelDescriptor,
+    ) -> ModelDescriptor:
+        if model_descriptor.WhichOneof("model_provider") == "open_ai_provider":
+            return ModelDescriptorMapper.from_prompt_openai_provider(
+                model_descriptor.open_ai_provider
+            )
+        else:
+            raise FrogmlException(
+                f"Got unsupported model descriptor: {repr(model_descriptor)}"
+            )
+
+    @staticmethod
+    def to_openai_chat(model_descriptor: OpenAIChat) -> ProtoOpenAIModelDescriptor:
+        d: OpenAIChat = model_descriptor
+        logit_bias_struct = Struct()
+
+        logit_bias_struct.update(d.logit_bias) if d.logit_bias else None
+
+        stop_list_value = ListValue()
+        stop_list_value.extend(d.stop) if d.stop else None
+        tool_choice_proto: ProtoToolChoice
+
+        if isinstance(d.tool_choice, str):
+            tool_choice_proto = ProtoToolChoice(literal=d.tool_choice)
+        elif d.tool_choice is not None:
+            tool_choice_struct = Struct()
+            tool_choice_struct.update(d.tool_choice)
+            tool_choice_proto = ProtoToolChoice(json=tool_choice_struct)
+
+        tools_structs = []
+        if d.tools:
+            for tool in d.tools:
+                s = Struct()
+                s.update(tool)
+                tools_structs.append(s)
+
+        tools_proto = ProtoTools(tools=tools_structs)
+
+        model_capabilities = ProtoOpenAICapabilities(
+            chat_api=ProtoOpenAIChatAPI(
+                chat_params=ProtoOpenAIChatModelParams(
+                    frequency_penalty=d.frequency_penalty,
+                    logit_bias=logit_bias_struct if d.logit_bias else None,
+                    logprobs=d.logprobs,
+                    max_tokens=d.max_tokens,
+                    n=d.n,
+                    presence_penalty=d.presence_penalty,
+                    response_format=(
+                        d.response_format if d.response_format else None
+                    ),  # noqa
+                    seed=d.seed,
+                    stop=stop_list_value if d.stop else None,
+                    temperature=d.temperature,
+                    tool_choice=tool_choice_proto if d.tool_choice else None,
+                    tools_spec=tools_proto if d.tools else None,
+                    top_logprobs=d.top_logprobs,
+                    top_p=d.top_p,
+                    user=d.user,
+                )
+            )
+        )
+
+        return ProtoOpenAIModelDescriptor(
+            model_id=model_descriptor.model_id, capabilities=model_capabilities
+        )
+
+    @staticmethod
+    def to_model_descriptor(
+        model_descriptor: ModelDescriptor,
+    ) -> Union[ProtoOpenAIModelDescriptor]:
+        if isinstance(model_descriptor, OpenAIChat):
+            return ModelDescriptorMapper.to_openai_chat(model_descriptor)
+
+        raise FrogmlException(
+            f"Got unsupported model descriptor: {repr(model_descriptor)}"
+        )

frogml/clients/prompt_manager/prompt_manager_client.py

@@ -0,0 +1,190 @@
+import grpc
+from dependency_injector.wiring import Provide, inject
+
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    CreatePromptRequest as ProtoCreatePromptRequest,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    CreatePromptResponse as ProtoCreatePromptResponse,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    CreatePromptVersionRequest as ProtoCreatePromptVersionRequest,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    CreatePromptVersionResponse as ProtoCreatePromptVersionResponse,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    DeletePromptRequest as ProtoDeletePromptRequest,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    DeletePromptVersionRequest as ProtoDeletePromptVersionRequest,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    GetPromptByNameRequest as ProtoGetPromptByNameRequest,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    GetPromptByNameResponse as ProtoGetPromptByNameResponse,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    GetPromptVersionByPromptNameRequest as ProtoGetPromptVersionByPromptNameRequest,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    GetPromptVersionByPromptNameResponse as ProtoGetPromptVersionByPromptNameResponse,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    SetDefaultPromptVersionRequest as ProtoSetDefaultPromptVersionRequest,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2_grpc import (
+    PromptManagerServiceStub,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import Prompt as ProtoPrompt
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import PromptSpec as ProtoPromptSpec
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
+    PromptVersion as ProtoPromptVersion,
+)
+from _frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
+    PromptVersionSpec as ProtoPromptVersionSpec,
+)
+from frogml.exceptions import FrogmlException
+from frogml.inner.di_configuration import FrogmlContainer
+
+
+class PromptManagerClient:
+    @inject
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._grpc_client: PromptManagerServiceStub = PromptManagerServiceStub(
+            grpc_channel
+        )
+
+    def create_prompt(
+        self,
+        name: str,
+        prompt_description: str,
+        version_spec: ProtoPromptVersionSpec,
+    ) -> ProtoPrompt:
+        request = ProtoCreatePromptRequest(
+            prompt_name=name,
+            prompt_spec=ProtoPromptSpec(description=prompt_description),
+            prompt_version_spec=version_spec,
+        )
+        try:
+            response: ProtoCreatePromptResponse = self._grpc_client.CreatePrompt(
+                request
+            )
+            return response.prompt
+        except grpc.RpcError as error:
+            call: grpc.Call = error  # noqa
+            if call.code() == grpc.StatusCode.ALREADY_EXISTS:
+                raise FrogmlException(f"Prompt with name: {name} already exists")
+            elif call.code() == grpc.StatusCode.INVALID_ARGUMENT:
+                raise FrogmlException(
+                    f"Got an illegal prompt specification: {call.details()}"
+                )
+            else:
+                raise FrogmlException(f"Internal Error: {call.details()}")
+
+    def create_prompt_version(
+        self,
+        name: str,
+        version_spec: ProtoPromptVersionSpec,
+        set_default: bool,
+    ) -> ProtoPromptVersion:
+        request = ProtoCreatePromptVersionRequest(
+            prompt_name=name, prompt_version_spec=version_spec, set_default=set_default
+        )
+        try:
+            response: ProtoCreatePromptVersionResponse = (
+                self._grpc_client.CreatePromptVersion(request)
+            )
+            return response.prompt_version
+        except grpc.RpcError as error:
+            call: grpc.Call = error  # noqa
+            if call.code() == grpc.StatusCode.NOT_FOUND:
+                raise FrogmlException(
+                    f"Can not update prompt: '{name}', prompt was not found"
+                )
+            elif call.code() == grpc.StatusCode.INVALID_ARGUMENT:
+                raise FrogmlException(
+                    f"Got an illegal prompt specification: {call.details()}"
+                )
+            else:
+                raise FrogmlException(f"Internal Error: {call.details()}")
+
+    def delete_prompt(self, name: str):
+        try:
+            self._grpc_client.DeletePrompt(ProtoDeletePromptRequest(prompt_name=name))
+        except grpc.RpcError as error:
+            call: grpc.Call = error  # noqa
+            if call.code() == grpc.StatusCode.NOT_FOUND:
+                raise FrogmlException(f"Prompt named '{name}' was not found")
+            else:
+                raise FrogmlException(f"Internal Error: {call.details()}")
+
+    def delete_prompt_version(self, name: str, version: int):
+        try:
+            self._grpc_client.DeletePromptVersion(
+                ProtoDeletePromptVersionRequest(
+                    prompt_name=name, version_number=version
+                )
+            )
+        except grpc.RpcError as error:
+            call: grpc.Call = error  # noqa
+            if call.code() == grpc.StatusCode.NOT_FOUND:
+                raise FrogmlException(str(call.details()))
+            elif call.code() == grpc.StatusCode.FAILED_PRECONDITION:
+                raise FrogmlException(
+                    f"Cannot delete the default version '{version}' of a prompt '{name}',"
+                    f" please set another version as the default to delete this version."
+                )
+            else:
+                raise FrogmlException(f"Internal Error: {call.details()}")
+
+    def get_prompt_by_name(self, name: str) -> ProtoPrompt:
+        """
+        Get prompt's default version
+        """
+        try:
+            response: ProtoGetPromptByNameResponse = self._grpc_client.GetPromptByName(
+                ProtoGetPromptByNameRequest(prompt_name=name)
+            )
+            return response.prompt
+        except grpc.RpcError as error:
+            call: grpc.Call = error  # noqa
+            if call.code() == grpc.StatusCode.NOT_FOUND:
+                raise FrogmlException(str(call.details()))
+            else:
+                raise FrogmlException(f"Internal Error: {call.details()}")
+
+    def get_prompt_version_by_name(self, name: str, version: int) -> ProtoPromptVersion:
+        """
+        Get prompt specific version
+        """
+        try:
+            response: ProtoGetPromptVersionByPromptNameResponse = (
+                self._grpc_client.GetPromptVersionByPromptName(
+                    ProtoGetPromptVersionByPromptNameRequest(
+                        prompt_name=name, version_number=version
+                    )
+                )
+            )
+            return response.prompt_version
+        except grpc.RpcError as error:
+            call: grpc.Call = error  # noqa
+            if call.code() == grpc.StatusCode.NOT_FOUND:
+                raise FrogmlException(str(call.details()))
+            else:
+                raise FrogmlException(f"Internal Error: {call.details()}")
+
+    def set_default_prompt_version(self, name: str, version: int):
+        try:
+            self._grpc_client.SetDefaultPromptVersion(
+                ProtoSetDefaultPromptVersionRequest(
+                    prompt_name=name, version_number=version
+                )
+            )
+        except grpc.RpcError as error:
+            call: grpc.Call = error  # noqa
+            if call.code() == grpc.StatusCode.NOT_FOUND:
+                raise FrogmlException(str(call.details()))
+            else:
+                raise FrogmlException(f"Internal Error: {call.details()}")