frogml-core 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

frogml_core/llmops/prompt/chat/template.py

@@ -0,0 +1,113 @@
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, List, Tuple, Union
+
+from frogml_core.llmops.prompt.chat.message import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+)
+from frogml_core.llmops.prompt.chat.value import ChatPromptValue
+from frogml_core.llmops.prompt.template import BasePromptTemplate, StringPromptTemplate
+
+
+@dataclass
+class BaseMessagePromptTemplate(BasePromptTemplate):
+    @abstractmethod
+    def render(self, variables: Dict[str, any]) -> BaseMessage:
+        pass
+
+
+@dataclass
+class BaseStringMessagePromptTemplate(BaseMessagePromptTemplate, ABC):
+    template: StringPromptTemplate = field(init=False)
+    role_name: str = field(init=False)
+
+    def __init__(self, template: str):
+        self.template = StringPromptTemplate(template=template)
+
+
+class AIMessagePromptTemplate(BaseStringMessagePromptTemplate):
+    role_name: str = "ai"
+
+    def render(self, variables: Dict[str, any]) -> BaseMessage:
+        return AIMessage(content=self.template.render(variables=variables).to_string())
+
+
+class HumanMessagePromptTemplate(BaseStringMessagePromptTemplate):
+    role_name: str = "human"
+
+    def render(self, variables: Dict[str, any]) -> BaseMessage:
+        return HumanMessage(
+            content=self.template.render(variables=variables).to_string()
+        )
+
+
+class SystemMessagePromptTemplate(BaseStringMessagePromptTemplate):
+    role_name: str = "system"
+
+    def render(self, variables: Dict[str, any]) -> BaseMessage:
+        return SystemMessage(
+            content=self.template.render(variables=variables).to_string()
+        )
+
+
+@dataclass
+class ChatPromptTemplate(BasePromptTemplate):
+    messages: List[Union[BaseMessage, BaseStringMessagePromptTemplate]]
+
+    def render(self, variables: Dict[str, any]) -> ChatPromptValue:
+        resulting_messages: List[BaseMessage] = list()
+
+        for message in self.messages:
+            if isinstance(message, BaseMessage):
+                resulting_messages.append(message)
+            elif isinstance(message, BaseStringMessagePromptTemplate):
+                resulting_messages.append(message.render(variables=variables))
+            else:
+                raise ValueError(
+                    f"Got unsupported message type: {repr(message)}. \n"
+                    "Supported messages are: "
+                    "AIMessagePromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate, \n"
+                    "AIMessage, HumanMessage, SystemMessage."
+                )
+
+        return ChatPromptValue(messages=resulting_messages)
+
+    def to_messages(self) -> List[Tuple[str, str]]:
+        """
+        Useful for integration with other libraries such as Langchain.
+
+        ```
+        ChatPromptTemplate(
+                messages=[
+                    SystemMessage("you are an assistant"),
+                    HumanMessagePromptTemplate("{{question}}")
+                ]
+            ).to_messages()
+
+        resulting in:
+
+        [("system", "you are an assistant"),
+        ("human", "{question}")]
+        ```
+
+        """
+
+        def strip_curly(string: str) -> str:
+            return re.sub(r"\{\{\s*([\w\s]+)\s*\}\}", repl=r"{\g<1>}", string=string)
+
+        if not self.messages:
+            return []
+
+        result: List[Tuple[str, str]] = []
+
+        for msg in self.messages:
+            if isinstance(msg, BaseMessage):
+                result.append((msg.role_name, msg.content))
+            elif isinstance(msg, BaseStringMessagePromptTemplate):
+                result.append((msg.role_name, strip_curly(msg.template.template)))
+
+        return result

frogml_core/llmops/prompt/chat/value.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass
+from typing import List
+
+from frogml_core.llmops.prompt.chat.message import BaseMessage
+from frogml_core.llmops.prompt.value import PromptValue
+
+
+@dataclass
+class ChatPromptValue(PromptValue):
+    messages: List[BaseMessage]

frogml_core/llmops/prompt/manager.py

@@ -0,0 +1,138 @@
+from typing import Optional
+
+from frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import Prompt as ProtoPrompt
+from frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
+    PromptVersion as ProtoPromptVersion,
+)
+from frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
+    PromptVersionDefinition as ProtoPromptVersionDefinition,
+)
+from frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
+    PromptVersionSpec as ProtoPromptVersionSpec,
+)
+from frogml_core.clients.prompt_manager.prompt_manager_client import PromptManagerClient
+from frogml_core.clients.prompt_manager.prompt_proto_mapper import PromptProtoMapper
+from frogml_core.exceptions import FrogmlException
+from frogml_core.llmops.prompt.base import BasePrompt, ChatPrompt, RegisteredPrompt
+
+
+class PromptManager:
+    _prompt_manager_client: PromptManagerClient
+
+    def __init__(self):
+        self._prompt_manager_client = PromptManagerClient()
+
+    def register(
+        self,
+        name: str,
+        prompt: BasePrompt,
+        prompt_description: Optional[str] = None,
+        version_description: Optional[str] = None,
+    ) -> RegisteredPrompt:
+        """
+        Registers a new prompt in Qwak platform. Name must be unique
+         and conform to ^[a-z0-9](?:[-_]?[a-z0-9]+)+$
+        """
+        if not isinstance(prompt, ChatPrompt):
+            raise FrogmlException(f"Got unsupported prompt type: {prompt}")
+
+        version_spec: ProtoPromptVersionSpec = PromptProtoMapper.to_prompt_version_spec(
+            version_description=version_description,
+            prompt_template=prompt.template,
+            model_descriptor=prompt.model,
+        )
+
+        registered_prompt: ProtoPrompt = self._prompt_manager_client.create_prompt(
+            name=name, prompt_description=prompt_description, version_spec=version_spec
+        )
+
+        return PromptProtoMapper.from_prompt(
+            name=registered_prompt.name,
+            prompt_description=registered_prompt.prompt_spec.description,
+            version_description=registered_prompt.default_version_definition.version_spec.description,
+            version=registered_prompt.default_version_definition.version_number,
+            target_default_version=True,
+            prompt_version_definition=registered_prompt.default_version_definition,
+        )
+
+    def update(
+        self,
+        name: str,
+        prompt: BasePrompt,
+        version_description: Optional[str] = None,
+        set_default: bool = False,
+    ) -> RegisteredPrompt:
+        """
+        Creates a new version for an existing prompt, prompt name must already exist.
+        `set_default` set to True if this version is to become the default one immediately.
+        """
+        if not isinstance(prompt, ChatPrompt):
+            raise FrogmlException(f"Got unsupported prompt type: {prompt}")
+
+        version_spec: ProtoPromptVersionSpec = PromptProtoMapper.to_prompt_version_spec(
+            version_description=version_description,
+            prompt_template=prompt.template,
+            model_descriptor=prompt.model,
+        )
+
+        prompt_version: ProtoPromptVersion = (
+            self._prompt_manager_client.create_prompt_version(
+                name=name, version_spec=version_spec, set_default=set_default
+            )
+        )
+
+        version_for_get_request = (
+            None
+            if set_default
+            else prompt_version.prompt_version_definition.version_number
+        )
+        return self.get_prompt(name=name, version=version_for_get_request)
+
+    def set_default(self, name: str, version: int):
+        """
+        Set a version of a registered prompt named: `name`, as the default version
+        """
+        self._prompt_manager_client.set_default_prompt_version(
+            name=name, version=version
+        )
+
+    def delete_prompt(self, name: str):
+        """
+        Delete all version of a prompt, by name
+        """
+        self._prompt_manager_client.delete_prompt(name=name)
+
+    def delete_prompt_version(self, name: str, version: int):
+        """
+        Deletes a specific version of a registered prompt
+        """
+        self._prompt_manager_client.delete_prompt_version(name=name, version=version)
+
+    def get_prompt(self, name: str, version: Optional[int] = None) -> RegisteredPrompt:
+        """
+        Get a registered prompt by name. To get the default version omit the `version` param, else
+        fetch the specified version.
+        """
+
+        prompt_default_version: ProtoPrompt = (
+            self._prompt_manager_client.get_prompt_by_name(name=name)
+        )
+        prompt_version_definition: ProtoPromptVersionDefinition = (
+            prompt_default_version.default_version_definition
+        )
+        if version:
+            prompt_version: ProtoPromptVersion = (
+                self._prompt_manager_client.get_prompt_version_by_name(
+                    name=name, version=version
+                )
+            )
+            prompt_version_definition = prompt_version.prompt_version_definition
+
+        return PromptProtoMapper.from_prompt(
+            name=prompt_default_version.name,
+            prompt_description=prompt_default_version.prompt_spec.description,
+            version_description=prompt_version_definition.version_spec.description,
+            version=prompt_version_definition.version_number,
+            target_default_version=not bool(version),
+            prompt_version_definition=prompt_version_definition,
+        )

frogml_core/llmops/prompt/template.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Dict
+
+from frogml_core.llmops.prompt.value import PromptValue, StringPromptValue
+
+
+@dataclass
+class BasePromptTemplate(ABC):
+    @abstractmethod
+    def render(self, variables: Dict[str, any]) -> PromptValue:
+        pass
+
+
+@dataclass
+class StringPromptTemplate(BasePromptTemplate):
+    template: str
+
+    def render(self, variables: Dict[str, any]) -> StringPromptValue:
+        from chevron import renderer
+
+        return StringPromptValue(
+            text=renderer.render(template=self.template, data=variables, warn=True)
+        )

frogml_core/llmops/provider/chat.py

@@ -0,0 +1,44 @@
+from functools import lru_cache
+from typing import Union
+
+from frogml_core.llmops.generation.chat.openai.types.chat.chat_completion import (
+    ChatCompletion,
+)
+from frogml_core.llmops.generation.streaming import ChatCompletionStream
+from frogml_core.llmops.model.descriptor import ChatModelDescriptor, OpenAIChat
+from frogml_core.llmops.prompt.chat.value import ChatPromptValue
+from frogml_core.llmops.provider.openai.provider import OpenAIProvider
+
+
+class ChatCompletionProvider:
+    @staticmethod
+    @lru_cache(maxsize=None)
+    def _get_openai_provider():
+        return OpenAIProvider()
+
+    @staticmethod
+    def invoke(
+        chat_prompt_value: ChatPromptValue,
+        chat_model_descriptor: ChatModelDescriptor,
+        stream: bool = False,
+    ) -> Union[ChatCompletion, ChatCompletionStream]:
+        if isinstance(chat_model_descriptor, OpenAIChat):
+            return ChatCompletionProvider._invoke_openai_chat(
+                chat_prompt_value=chat_prompt_value,
+                chat_model_descriptor=chat_model_descriptor,
+                stream=stream,
+            )
+        else:
+            raise ValueError("Can't invoke prompt and model combination!")
+
+    @staticmethod
+    def _invoke_openai_chat(
+        chat_prompt_value: ChatPromptValue,
+        chat_model_descriptor: OpenAIChat,
+        stream: bool = False,
+    ) -> Union[ChatCompletion, ChatCompletionStream]:
+        return ChatCompletionProvider._get_openai_provider().create_chat_completion(
+            chat_prompt_value=chat_prompt_value,
+            chat_model_descriptor=chat_model_descriptor,
+            stream=stream,
+        )

frogml_core/llmops/provider/openai/client.py

@@ -0,0 +1,126 @@
+import json
+import os
+from typing import Dict, Iterable, List, Optional, Union
+from urllib.parse import urljoin
+
+import requests
+from dacite import Config, from_dict
+from requests import Response
+from typing_extensions import Literal
+
+from frogml_core.exceptions.frogml_external_exception import FrogmlExternalException
+from frogml_core.llmops.generation._steaming import BaseSSEDecoder  # noqa
+from frogml_core.llmops.generation.chat.openai.types.chat.chat_completion import (
+    ChatCompletion,
+)
+from frogml_core.llmops.generation.chat.openai.types.chat.chat_completion_chunk import (
+    ChatCompletionChunk,
+)
+from frogml_core.llmops.generation.chat.openai.types.chat.chat_completion_tool_choice_option_param import (
+    ChatCompletionToolChoiceOptionParam,
+)
+from frogml_core.llmops.generation.chat.openai.types.chat.chat_completion_tool_param import (
+    ChatCompletionToolParam,
+)
+from frogml_core.llmops.generation.streaming import ChatCompletionStream
+from frogml_core.utils.dict_utils import remove_none_value_keys
+
+
+class OpenAIChatCompletionStream(
+    BaseSSEDecoder[ChatCompletionChunk], ChatCompletionStream
+):
+    def __init__(self, response: requests.Response):
+        super().__init__(response=response, parse_to=ChatCompletionChunk)
+
+
+class OpenAIClient:
+    base_url: str
+
+    def __init__(self):
+        self.base_url: str = os.environ.get(
+            "_QWAK_OPEN_AI_BASE_URL", "https://api.openai.com"
+        )
+
+    def invoke_chat_completion(
+        self,
+        api_key: str,
+        model: str,
+        messages: List[Dict],
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, int]] = None,
+        logprobs: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Literal["text", "json_object"] = None,
+        seed: Optional[int] = None,
+        stop: Union[Optional[str], List[str]] = None,
+        stream: Optional[bool] = False,
+        temperature: Optional[float] = None,
+        top_logprobs: Optional[int] = None,
+        top_p: Optional[float] = None,
+        user: Optional[str] = None,
+        tool_choice: Optional[ChatCompletionToolChoiceOptionParam] = None,
+        tools: Iterable[ChatCompletionToolParam] = None,
+        extra_headers: Optional[Dict[str, str]] = None,
+        extra_body: Optional[Dict[str, str]] = None,
+        timeout_seconds: Optional[float] = None,
+    ) -> Union[ChatCompletion, ChatCompletionStream]:
+        url: str = urljoin(self.base_url, "v1/chat/completions")
+        headers: Dict[str, str] = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {api_key}",
+        }
+        body = {
+            "messages": messages,
+            "model": model,
+            "frequency_penalty": frequency_penalty,
+            "logit_bias": logit_bias,
+            "logprobs": logprobs,
+            "max_tokens": max_tokens,
+            "n": n,
+            "presence_penalty": presence_penalty,
+            "response_format": {"type": response_format} if response_format else None,
+            "seed": seed,
+            "stop": stop,
+            "temperature": temperature,
+            "tool_choice": tool_choice if tools else None,
+            "tools": tools if tools else None,
+            "top_logprobs": top_logprobs,
+            "top_p": top_p,
+            "user": user,
+            "stream": stream if stream else None,
+        }
+        body = remove_none_value_keys(body)
+
+        if extra_headers:
+            headers.update(extra_headers)
+
+        if extra_body:
+            body.update(extra_body)
+
+        http_request_timeout_seconds: float = (
+            timeout_seconds
+            if timeout_seconds
+            else float(os.environ.get("_QWAK_OPEN_AI_TIMEOUT_SECONDS", 60.0))
+        )
+        response: Response = requests.post(
+            url=url,
+            data=json.dumps(body),
+            headers=headers,
+            stream=stream,
+            timeout=http_request_timeout_seconds,
+        )
+
+        try:
+            response.raise_for_status()
+        except requests.exceptions.HTTPError as e:
+            raise FrogmlExternalException(message=e.response.content.decode())
+        if stream:
+            return OpenAIChatCompletionStream(response=response)
+        else:
+            return from_dict(
+                data_class=ChatCompletion,
+                data=response.json(),
+                config=Config(check_types=False),
+            )

frogml_core/llmops/provider/openai/provider.py

@@ -0,0 +1,93 @@
+import random
+from typing import Dict, List, Optional, Union
+
+from frogml_core.clients.integration_management.integration_utils import (
+    IntegrationUtils,
+)
+from frogml_core.clients.integration_management.openai.openai_system_secret import (
+    OpenAIApiKeySystemSecret,
+)
+from frogml_core.exceptions import FrogmlException
+from frogml_core.llmops.generation.chat.openai.types.chat.chat_completion import (
+    ChatCompletion,
+)
+from frogml_core.llmops.generation.streaming import ChatCompletionStream
+from frogml_core.llmops.model.descriptor import OpenAIChat
+from frogml_core.llmops.prompt.chat.message import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+)
+from frogml_core.llmops.prompt.chat.value import ChatPromptValue
+from frogml_core.llmops.provider.openai.client import OpenAIClient
+
+
+class OpenAIProvider:
+    client: OpenAIClient
+
+    def __init__(self):
+        self.client = OpenAIClient()
+
+    def _get_random_openai_api_key(self) -> Optional[str]:
+        openai_api_keys: List[
+            OpenAIApiKeySystemSecret
+        ] = IntegrationUtils().get_openai_api_keys()
+        if len(openai_api_keys) == 0:
+            return None
+
+        return random.choice(openai_api_keys).get_api_key()  # nosec
+
+    def _chat_value_to_json(self, chat_prompt_value: ChatPromptValue) -> List[Dict]:
+        return [self._map_message(m) for m in chat_prompt_value.messages]
+
+    def _map_message(self, message: BaseMessage) -> Dict[str, str]:
+        role: str
+        content: str = message.content
+
+        if isinstance(message, AIMessage):
+            role = "assistant"
+        elif isinstance(message, SystemMessage):
+            role = "system"
+        elif isinstance(message, HumanMessage):
+            role = "user"
+        else:
+            raise FrogmlException(f"Can't handle message of type: {repr(message)}")
+
+        return {"role": role, "content": content}
+
+    def create_chat_completion(
+        self,
+        chat_prompt_value: ChatPromptValue,
+        chat_model_descriptor: OpenAIChat,
+        stream: bool = False,
+    ) -> Union[ChatCompletion, ChatCompletionStream]:
+        openai_api_key: Optional[str] = self._get_random_openai_api_key()
+        if not openai_api_key:
+            raise FrogmlException(
+                "Could not find Open AI integration, Please create one."
+            )
+
+        d = chat_model_descriptor
+
+        return self.client.invoke_chat_completion(
+            stream=stream,
+            api_key=openai_api_key,
+            model=d.model_id,
+            messages=self._chat_value_to_json(chat_prompt_value),
+            frequency_penalty=d.frequency_penalty,
+            logit_bias=d.logit_bias,
+            logprobs=d.logprobs,
+            max_tokens=d.max_tokens,
+            n=d.n,
+            presence_penalty=d.presence_penalty,
+            response_format=d.response_format,
+            seed=d.seed,
+            stop=d.stop,
+            temperature=d.temperature,
+            top_logprobs=d.top_logprobs,
+            top_p=d.top_p,
+            user=d.user,
+            tool_choice=d.tool_choice,
+            tools=d.tools,
+        )

frogml_core/model/_entity_extraction.py

@@ -0,0 +1,116 @@
+import dataclasses
+from typing import Dict, List, cast
+
+from frogml_core.exceptions import FrogmlException
+from frogml_core.feature_store._common.feature_set_utils import (
+    FeatureSetInfo,
+    get_entity_type,
+    get_typed_feature,
+)
+from frogml_core.model.fs_info_mapping_retriever import retrieve_fs_mapping
+from frogml_core.model.schema_entities import BaseFeature, Entity, FeatureStoreInput
+from frogml_core.model.utils.feature_utils import (
+    discard_env_from_name,
+    extract_env_name,
+    extract_featureset_name,
+    validate_and_sanitize_features_name,
+)
+
+
+def normalize_features(
+    features: List[BaseFeature],
+    env_to_feature_sets_infos: Dict[str, Dict[str, FeatureSetInfo]],
+) -> List[BaseFeature]:
+    """
+    Adding the relevant entity to the feature set features
+    Args:
+        features: list of features
+        env_to_feature_sets_infos: dict of environment name to feature set name to it's entity
+    return normalize features - the features with the entities
+    :param
+    """
+    new_normalize_features = []
+    for feature in features:
+        if isinstance(feature, FeatureStoreInput):
+            feature_set_env = extract_env_name(feature.name)
+            feature_set_name = extract_featureset_name(feature.name)
+            feature_set_to_feature_set_info = env_to_feature_sets_infos.get(
+                feature_set_env
+            )
+            if feature_set_to_feature_set_info is None:
+                raise FrogmlException(f"Environment '{feature_set_env}' does not exist")
+            feature_set_info = feature_set_to_feature_set_info.get(feature_set_name)
+            if feature_set_info is None:
+                raise FrogmlException(
+                    f"Featureset '{feature_set_name}' does not exist in env '{feature_set_env}'"  # noqa
+                )
+            entity_spec = feature_set_info.entity_spec
+            retrieved_entity = Entity(
+                name=entity_spec.keys[
+                    0
+                ],  # currently support only one entity key per feature set
+                type=get_entity_type(entity_spec.value_type),
+            )
+            if (
+                feature.entity
+                and retrieved_entity.name.lower() != feature.entity.name.lower()
+            ):
+                raise FrogmlException(
+                    f"Explicitly supplied with an invalid entity: {feature.entity}, "
+                    f"actual: {retrieved_entity}"
+                )
+            feature.entity = retrieved_entity
+            feature.entity.name = feature.entity.name.lower()
+            feature_type = feature_set_info.feature_set_type
+            feature_version = feature_set_info.feature_version
+            feature_without_env = cast(
+                FeatureStoreInput,
+                dataclasses.replace(feature, name=discard_env_from_name(feature.name)),
+            )
+            feature = get_typed_feature(
+                feature_without_env, feature_type, feature_version
+            )
+        new_normalize_features.append(feature)
+    return new_normalize_features
+
+
+def adding_entities_to_schema(entities: List[Entity], features: List[BaseFeature]):
+    """
+    Adding the new entities for feature store feature to the entities list from schema
+    Args:
+        entities: list of entities
+        features: list of features
+    Returns: the entities with the new entities from the feature store features
+    """
+    entity_name_to_entity = {entity.name: entity for entity in entities}
+    for feature in features:
+        if isinstance(feature, FeatureStoreInput):
+            entity_name_to_entity[feature.entity.name] = feature.entity
+
+    for entity in entity_name_to_entity.values():
+        entity.name = entity.name.lower()
+
+    return list(set(entity_name_to_entity.values()))
+
+
+def normalize_schema(features: List[BaseFeature], entities: List[Entity]):
+    """
+    validate and sanitize the feature names and entities
+    Normalize schema - will add entity to each features in schema if not exists
+    Arg:
+        model: the model with the required schema
+    Return:
+        updated model
+    """
+    sanitized_features = validate_and_sanitize_features_name(features)
+    fs_info_cache = retrieve_fs_mapping(sanitized_features)
+
+    normalized_features = normalize_features(sanitized_features, fs_info_cache)
+    normalized_entities = adding_entities_to_schema(entities, sanitized_features)
+
+    return normalized_features, normalized_entities
+
+
+def enrich_schema(features: List[BaseFeature], entities: List[Entity]):
+    normalized_features, normalized_entities = normalize_schema(features, entities)
+    return normalized_features, normalized_entities