qwak-core 0.1.37__py3-none-any.whl → 0.5.19__py3-none-any.whl

qwak/llmops/prompt/manager.py

@@ -0,0 +1,134 @@
+from typing import Optional
+
+from _qwak_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
+    Prompt as ProtoPrompt,
+    PromptVersion as ProtoPromptVersion,
+    PromptVersionDefinition as ProtoPromptVersionDefinition,
+    PromptVersionSpec as ProtoPromptVersionSpec,
+)
+from qwak.clients.prompt_manager.prompt_manager_client import PromptManagerClient
+from qwak.clients.prompt_manager.prompt_proto_mapper import PromptProtoMapper
+from qwak.exceptions import QwakException
+from qwak.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 QwakException(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 QwakException(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,
+        )

qwak/llmops/prompt/template.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Dict
+
+from qwak.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)
+        )

qwak/llmops/prompt/value.py

@@ -0,0 +1,14 @@
+from abc import ABC
+from dataclasses import dataclass
+
+
+class PromptValue(ABC):
+    pass
+
+
+@dataclass
+class StringPromptValue(PromptValue):
+    text: str
+
+    def to_string(self) -> str:
+        return self.text

qwak/llmops/provider/chat.py

@@ -0,0 +1,42 @@
+from functools import lru_cache
+from typing import Union
+
+from qwak.llmops.generation.chat.openai.types.chat.chat_completion import ChatCompletion
+from qwak.llmops.generation.streaming import ChatCompletionStream
+from qwak.llmops.model.descriptor import ChatModelDescriptor, OpenAIChat
+from qwak.llmops.prompt.chat.value import ChatPromptValue
+from qwak.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,
+        )

qwak/llmops/provider/openai/client.py

@@ -0,0 +1,123 @@
+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 qwak.exceptions.qwak_external_exception import QwakExternalException
+from qwak.llmops.generation._steaming import BaseSSEDecoder  # noqa
+from qwak.llmops.generation.chat.openai.types.chat.chat_completion import ChatCompletion
+from qwak.llmops.generation.chat.openai.types.chat.chat_completion_chunk import (
+    ChatCompletionChunk,
+)
+from qwak.llmops.generation.chat.openai.types.chat.chat_completion_tool_choice_option_param import (
+    ChatCompletionToolChoiceOptionParam,
+)
+from qwak.llmops.generation.chat.openai.types.chat.chat_completion_tool_param import (
+    ChatCompletionToolParam,
+)
+from qwak.llmops.generation.streaming import ChatCompletionStream
+from qwak.utils.dict_utils import remove_none_value_keys
+from requests import Response
+from typing_extensions import Literal
+
+
+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 QwakExternalException(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),
+            )

qwak/llmops/provider/openai/provider.py

@@ -0,0 +1,89 @@
+import random
+from typing import Dict, List, Optional, Union
+
+from qwak.clients.integration_management.integration_utils import IntegrationUtils
+from qwak.clients.integration_management.openai.openai_system_secret import (
+    OpenAIApiKeySystemSecret,
+)
+from qwak.exceptions import QwakException
+from qwak.llmops.generation.chat.openai.types.chat.chat_completion import ChatCompletion
+from qwak.llmops.generation.streaming import ChatCompletionStream
+from qwak.llmops.model.descriptor import OpenAIChat
+from qwak.llmops.prompt.chat.message import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+)
+from qwak.llmops.prompt.chat.value import ChatPromptValue
+from qwak.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 QwakException(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 QwakException(
+                "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,
+        )

qwak/model/_entity_extraction.py

@@ -1,37 +1,50 @@
-from typing import Dict, List
+import dataclasses
+from typing import Dict, List, cast
 
-from qwak.clients.feature_store import FeatureRegistryClient
 from qwak.exceptions import QwakException
 from qwak.feature_store._common.feature_set_utils import (
     FeatureSetInfo,
     get_entity_type,
-    get_feature_set_info,
     get_typed_feature,
 )
+from qwak.model.fs_info_mapping_retriever import retrieve_fs_mapping
 from qwak.model.schema_entities import BaseFeature, Entity, FeatureStoreInput
-
-
-def get_feature_set_name(feature: BaseFeature) -> str:
-    feature_full_name: List[str] = feature.name.split(".")
-    return feature_full_name[0].lower()
+from qwak.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], feature_set_to_entities: Dict[str, FeatureSetInfo]
+    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
-        feature_set_to_entities: dict of feature set name and it's entity
-
+        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_name = get_feature_set_name(feature)
-            entity_spec = feature_set_to_entities.get(feature_set_name).entity_spec
+            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 QwakException(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 QwakException(
+                    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
@@ -48,16 +61,16 @@ def normalize_features(
                 )
             feature.entity = retrieved_entity
             feature.entity.name = feature.entity.name.lower()
-            if isinstance(feature, FeatureStoreInput):
-                feature_type = feature_set_to_entities.get(
-                    feature_set_name
-                ).feature_set_type
-                feature_version = feature_set_to_entities.get(
-                    feature_set_name
-                ).feature_version
-                feature = get_typed_feature(feature, feature_type, feature_version)
+            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
 
 
@@ -82,37 +95,18 @@ def adding_entities_to_schema(entities: List[Entity], features: List[BaseFeature
 
 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 = features
-    normalized_entities = entities
-    if features and hasattr(features, "__iter__"):
-        # featureset_name -> FeatureSetInfo
-        fs_info_cache: Dict[str, FeatureSetInfo] = dict()
-        features_manager_client = FeatureRegistryClient()
-        validation_errors = []
-        for feature in features:
-            try:
-                if isinstance(feature, FeatureStoreInput):
-                    feature_set_name = get_feature_set_name(feature)
-                    if feature_set_name not in fs_info_cache:
-                        fs_info_cache[feature_set_name] = get_feature_set_info(
-                            features_manager_client, feature_set_name
-                        )
-                else:
-                    fs_info_cache[feature.name] = feature
-            except QwakException as e:
-                validation_errors.append(e.message)
-
-        if validation_errors:
-            raise QwakException(f"Schema validation error: {validation_errors}")
-        normalized_features = normalize_features(features, fs_info_cache)
-        normalized_entities = adding_entities_to_schema(entities, features)
+    normalized_features = normalize_features(sanitized_features, fs_info_cache)
+    normalized_entities = adding_entities_to_schema(entities, sanitized_features)
 
     return normalized_features, normalized_entities
 

qwak/model/adapters/input_adapters/multi_input_adapter.py

@@ -16,8 +16,24 @@ class MultiInputAdapter(BaseInputAdapter):
     def __init__(self, adapters: List[BaseInputAdapter]):
         self.adapters = adapters
 
-    def extract_user_func_arg(self, data: bytes) -> Sequence[Any]:
+    def extract_user_func_arg(self, data: bytes, headers=None) -> Sequence[Any]:
         task = InferenceTask(data=data)
+        if headers:
+            adapter_name = headers.get("Adapter-Name")
+            if adapter_name:
+                for inner_adapter in self.adapters:
+                    if isinstance(inner_adapter, ProtoInputAdapter):
+                        if (
+                            inner_adapter.message.DESCRIPTOR.name == adapter_name
+                            or inner_adapter.message.DESCRIPTOR.full_name
+                            == adapter_name
+                        ):
+                            extracted_data = inner_adapter.extract_user_func_arg(data)
+                            if isinstance(extracted_data, tuple):
+                                extracted_data = extracted_data[0]
+                            if extracted_data is None:
+                                raise ValueError("Parser returned None")
+                            return extracted_data
         for inner_adapter in self.adapters:
             try:
                 if isinstance(inner_adapter, ProtoInputAdapter):

qwak/model/base.py

@@ -5,7 +5,7 @@ from qwak.model.schema import ModelSchema
 
 class QwakModel:
     """
-    Base class for all Qwak based models.
+    Base class for all Qwak based models
     """
 
     @abstractmethod

qwak/model/decorators/api.py

@@ -16,7 +16,7 @@ from qwak.model.adapters.output_adapters.dataframe_output_adapter import (
 API_NOT_CONFIGURED_ERROR_MESSAGE = (
     "There has been an error configuring the Qwak model. When testing locally, "
     "please use 'run_local' as described in our docs: "
-    "https://docs-saas.qwak.com/docs/testing-locally#running-local-inference"
+    "https://docs.qwak.com/docs/testing-locally#/running-models-locally"
 )
 
 
@@ -24,6 +24,7 @@ API_NOT_CONFIGURED_ERROR_MESSAGE = (
 def api_decorator(
     analytics: bool = True,
     analytics_sample_ratio: float = 1.0,
+    analytics_exclude_columns: list = [],
     feature_extraction: bool = False,
     input_adapter: BaseInputAdapter = DataFrameInputAdapter(),
     output_adapter: BaseOutputAdapter = DataFrameOutputAdapter(),
@@ -54,6 +55,7 @@ def api_decorator(
             input_adapter,
             output_adapter,
             analytics_sample_ratio,
+            analytics_exclude_columns,
         )
     except TypeError as e:
         if "__call__() takes 1" in str(e):

qwak/model/decorators/impl/api_implementation.py

@@ -11,6 +11,7 @@ def create_decorator_function(
     input_adapter: BaseInputAdapter,
     output_adapter: BaseOutputAdapter,
     analytics_sample_ratio: float,
+    analytics_exclude_columns: list,
 ) -> Callable:
     def api_decorator_function_load_logic(func: Callable[..., Any]):
         func = extract_wrapped(func)