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

frogml_core/vector_store/filters.py

@@ -0,0 +1,359 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+from google.protobuf.json_format import MessageToDict, ParseDict
+
+from frogml_proto.qwak.vectors.v1.filters_pb2 import And as ProtoAnd
+from frogml_proto.qwak.vectors.v1.filters_pb2 import (
+    AtomicLiteral as ProtoAtomicLiteral,
+)
+from frogml_proto.qwak.vectors.v1.filters_pb2 import Equal as ProtoEqual
+from frogml_proto.qwak.vectors.v1.filters_pb2 import Filter as ProtoFilter
+from frogml_proto.qwak.vectors.v1.filters_pb2 import GreaterThan as ProtoGreaterThan
+from frogml_proto.qwak.vectors.v1.filters_pb2 import (
+    GreaterThanEqual as ProtoGreaterThanEqual,
+)
+from frogml_proto.qwak.vectors.v1.filters_pb2 import IsNotNull as ProtoIsNotNull
+from frogml_proto.qwak.vectors.v1.filters_pb2 import IsNull as ProtoIsNull
+from frogml_proto.qwak.vectors.v1.filters_pb2 import LessThan as ProtoLessThan
+from frogml_proto.qwak.vectors.v1.filters_pb2 import (
+    LessThanEqual as ProtoLessThanEqual,
+)
+from frogml_proto.qwak.vectors.v1.filters_pb2 import Like as ProtoLike
+from frogml_proto.qwak.vectors.v1.filters_pb2 import NotEqual as ProtoNotEqual
+from frogml_proto.qwak.vectors.v1.filters_pb2 import Or as ProtoOr
+from frogml_core.vector_store.utils.filter_utils import transform
+
+
+class Filter(ABC):
+    """Abstract base class for filter objects."""
+
+    def And(self, other):
+        """Logical AND operation with another filter.
+
+        Usage Example:
+        ```
+        filter1 = Equal("age", 30)
+        filter2 = GreaterThan("score", 90)
+        combined_filter = filter1.And(filter2)
+        ```
+        """
+        return And(self, other)
+
+    def Or(self, other):
+        """Logical OR operation with another filter.
+
+        Usage Example:
+        ```
+        filter1 = Equal("age", 30)
+        filter2 = GreaterThan("score", 90)
+        combined_filter = filter1.Or(filter2)
+        ```
+        """
+        return Or(self, other)
+
+    @abstractmethod
+    def _to_proto(self):
+        """Convert the filter to a protobuf representation."""
+        pass
+
+
+@dataclass
+class And(Filter):
+    """Logical And operation with another filter.
+
+    Usage Example:
+    ```
+    filter1 = Equal("age", 30)
+    filter2 = GreaterThan("score", 90)
+    combined_filter = And(left=filter1, right=filter2)
+    ```
+    """
+
+    left: Filter
+    right: Filter
+
+    def _to_proto(self):
+        proto_filter = ProtoFilter()
+        proto_filter_dict = MessageToDict(proto_filter)
+        proto_and_dict = MessageToDict(
+            ProtoAnd(left=self.left._to_proto(), right=self.right._to_proto())
+        )
+        proto_filter_dict["and"] = proto_and_dict
+        return ParseDict(proto_filter_dict, proto_filter, ignore_unknown_fields=True)
+
+
+@dataclass
+class Or(Filter):
+    """Logical OR operation with another filter.
+
+    Usage Example:
+    ```
+    filter1 = Equal("age", 30)
+    filter2 = GreaterThan("score", 90)
+    combined_filter = Or(left=filter1, right=filter2)
+    ```
+    """
+
+    left: Filter
+    right: Filter
+
+    def _to_proto(self):
+        proto_filter = ProtoFilter()
+        proto_filter_dict = MessageToDict(proto_filter)
+        proto_or_dict = MessageToDict(
+            ProtoOr(left=self.left._to_proto(), right=self.right._to_proto())
+        )
+        proto_filter_dict["or"] = proto_or_dict
+        return ParseDict(proto_filter_dict, proto_filter, ignore_unknown_fields=True)
+
+
+@dataclass
+class _UnaryFilter(Filter):
+    property: str
+    value: Any
+
+    def _to_proto(self):
+        # Each UnaryFilter implements its own _to_proto
+        pass
+
+
+class Equal(_UnaryFilter):
+    """Equal operation.
+
+    Usage Example:
+    ```
+    filter = Equal(property="age", value=30)
+    ```
+    """
+
+    def _to_proto(self):
+        atomic_literal: ProtoAtomicLiteral = transform(value=self.value)
+        return ProtoFilter(eq=ProtoEqual(property=self.property, value=atomic_literal))
+
+
+class NotEqual(_UnaryFilter):
+    """NotEqual operation.
+
+    Usage Example:
+    ```
+    filter = NotEqual(property="age", value=30)
+    ```
+    """
+
+    def _to_proto(self):
+        atomic_literal: ProtoAtomicLiteral = transform(value=self.value)
+        return ProtoFilter(
+            ne=ProtoNotEqual(property=self.property, value=atomic_literal)
+        )
+
+
+class LessThanEqual(_UnaryFilter):
+    """LessThanEqual operation.
+
+    Usage Example:
+    ```
+    filter = LessThanEqual(property="age", value=30)
+    ```
+    """
+
+    def _to_proto(self):
+        atomic_literal: ProtoAtomicLiteral = transform(value=self.value)
+        return ProtoFilter(
+            lte=ProtoLessThanEqual(property=self.property, value=atomic_literal)
+        )
+
+
+class LessThan(_UnaryFilter):
+    """LessThan operation.
+
+    Usage Example:
+    ```
+    filter = LessThan(property="age", value=30)
+    ```
+    """
+
+    def _to_proto(self):
+        atomic_literal: ProtoAtomicLiteral = transform(value=self.value)
+        return ProtoFilter(
+            lt=ProtoLessThan(property=self.property, value=atomic_literal)
+        )
+
+
+class GreaterThanEqual(_UnaryFilter):
+    """GreaterThanEqual operation.
+
+    Usage Example:
+    ```
+    filter = GreaterThanEqual(property="age", value=30)
+    ```
+    """
+
+    def _to_proto(self):
+        atomic_literal: ProtoAtomicLiteral = transform(value=self.value)
+        return ProtoFilter(
+            gte=ProtoGreaterThanEqual(property=self.property, value=atomic_literal)
+        )
+
+
+class GreaterThan(_UnaryFilter):
+    """GreaterThan operation.
+
+    Usage Example:
+    ```
+    filter = GreaterThan(property="age", value=30)
+    ```
+    """
+
+    def _to_proto(self):
+        atomic_literal: ProtoAtomicLiteral = transform(value=self.value)
+        return ProtoFilter(
+            gt=ProtoGreaterThan(property=self.property, value=atomic_literal)
+        )
+
+
+@dataclass
+class Like(Filter):
+    """Like operation.
+
+    Usage Example:
+    ```
+    filter = Like(property="name", pattern="Tal")
+    ```
+    """
+
+    property: str
+    pattern: str
+
+    def _to_proto(self):
+        return ProtoFilter(like=ProtoLike(property=self.property, pattern=self.pattern))
+
+
+@dataclass
+class IsNull(Filter):
+    """IsNull operation.
+
+    Usage Example:
+    ```
+    filter = IsNull(property="zipcode)
+    ```
+    """
+
+    property: str
+
+    def _to_proto(self):
+        return ProtoFilter(is_null=ProtoIsNull(property=self.property))
+
+
+@dataclass
+class IsNotNull(Filter):
+    """IsNotNull operation.
+
+    Usage Example:
+    ```
+    filter = IsNotNull(property="zipcode)
+    ```
+    """
+
+    property: str
+
+    def _to_proto(self):
+        return ProtoFilter(is_not_null=ProtoIsNotNull(property=self.property))
+
+
+@dataclass
+class Property:
+    """Represents a property for building filter conditions."""
+
+    name: str
+
+    def gt(self, value: Any):
+        """Create a GreaterThan filter for this property.
+
+        Usage Example:
+        ```
+        filter = Property("age").gt(30)
+        ```
+        """
+        return GreaterThan(self.name, value)
+
+    def gte(self, value: Any):
+        """Create a GreaterThanEqual filter for this property.
+
+        Usage Example:
+        ```
+        filter = Property("score").gte(90)
+        ```
+        """
+        return GreaterThanEqual(self.name, value)
+
+    def lt(self, value: Any):
+        """Create a LessThan filter for this property.
+
+        Usage Example:
+        ```
+        filter = Property("age").lt(30)
+        ```
+        """
+        return LessThan(self.name, value)
+
+    def lte(self, value: Any):
+        """Create a LessThanEqual filter for this property.
+
+        Usage Example:
+        ```
+        filter = Property("score").lte(90)
+        ```
+        """
+        return LessThanEqual(self.name, value)
+
+    def eq(self, value: Any):
+        """Create an Equal filter for this property.
+
+        Usage Example:
+        ```
+        filter = Property("score").eq(90)
+        ```
+        """
+        return Equal(self.name, value)
+
+    def ne(self, value: Any):
+        """Create an NotEqual filter for this property.
+
+        Usage Example:
+        ```
+        filter = Property("age").ne(30)
+        ```
+        """
+        return NotEqual(self.name, value)
+
+    def is_null(self):
+        """Create an IsNull filter for this property.
+
+        Usage Example:
+        ```
+        filter = Property("zipcode").is_null()
+        ```
+        """
+        return IsNull(self.name)
+
+    def is_not_null(self):
+        """Create an IsNotNull filter for this property.
+
+        Usage Example:
+        ```
+        filter = Property("zipcode").is_not_null()
+        ```
+        """
+        return IsNotNull(self.name)
+
+    def like(self, pattern: str):
+        """Create a Like filter for this property.
+
+        Usage Example:
+        ```
+        filter = Property("name").like("Tal")
+        ```
+        """
+        return Like(self.name, pattern)

frogml_core/vector_store/inference_client.py

@@ -0,0 +1,105 @@
+import json
+
+from frogml_core.clients.administration.eco_system.client import EcosystemClient
+from frogml_core.exceptions import FrogmlException, FrogmlHTTPException
+from frogml_core.vector_store.rest_helpers import RestSession
+
+
+class VectorStoreInferenceClient:
+    def __init__(
+        self,
+        model_id: str,
+    ):
+        """
+        :param model_id: The model id to invoke against.
+        """
+
+        self.r_session = RestSession()
+        self.ecosystem_client = EcosystemClient()
+
+        authenticated_user_context = (
+            self.ecosystem_client.get_authenticated_user_context()
+        )
+        account_details = authenticated_user_context.user.account_details
+        default_environment_id = account_details.default_environment_id
+        model_url_prefix = account_details.environment_by_id[
+            default_environment_id
+        ].configuration.model_api_url
+
+        self.model_id = model_id
+        self.content_type = "application/json; format=pandas-split"
+        self.model_api = _get_model_url(
+            model_id=model_id, model_url_prefix=model_url_prefix
+        )
+
+    def predict(
+        self,
+        feature_vectors,
+    ):
+        """
+        Perform a prediction request against a Qwak based model
+
+        :param feature_vectors: A list of feature vectors to predict against. Each feature vector is modeled as a python
+         dictionary
+        :return: Prediction response from the model
+        """
+
+        if feature_vectors.__class__.__name__ == "DataFrame":
+            feature_vectors = feature_vectors.to_json(orient="split")
+
+        if isinstance(feature_vectors, dict) or isinstance(feature_vectors, list):
+            feature_vectors = json.dumps(feature_vectors)
+
+        try:
+            response = self.r_session.post(
+                self.model_api, data=feature_vectors, headers={}
+            )
+
+            if response.status_code >= 400:
+                exception_class_name = response.headers.get("X-Exception-Class")
+                msg = f"{response.status_code}: {response.text}"
+                raise FrogmlHTTPException(
+                    response.status_code, msg, exception_class_name
+                )
+
+            elif response.status_code != 200:
+                raise FrogmlHTTPException(response.status_code, response.content)
+
+            dict_response = json.loads(response.content)
+            return dict_response
+        except FrogmlHTTPException as e:
+            raise e
+        except Exception as e:
+            raise RuntimeError(f"Failed to make a prediction request. Error is: {e}")
+
+    def get_embedding(self, natural_input: str):
+        feature_vector = [
+            {
+                "input": natural_input,
+            }
+        ]
+        try:
+            result_list = self.predict(feature_vector)
+        except Exception as e:
+            raise FrogmlException(
+                f"Vectorizer {self.model_id} failed to transform input {feature_vector} to vectors. Error is: {str(e)}"
+            )
+        try:
+            vector = result_list[0]["embeddings"]
+        except Exception:
+            raise FrogmlException(
+                f"Vectorizer {self.model_id} must return a dataframe containing an 'embeddings' column"
+            )
+        if not vector:
+            raise FrogmlException(
+                f"Vectorizer {self.model_id} did not return embeddings for the given natural input. Unable to continue with the query."
+            )
+        return vector
+
+
+def _get_model_url(model_id: str, model_url_prefix: str) -> str:
+    scheme = "http" if model_url_prefix.startswith("localhost") else "https"
+    effective_model_id = model_id.replace("-", "_")
+    return (
+        f"{scheme}://{model_url_prefix}/v1/{effective_model_id}/predict"  # noqa: E231
+    )

frogml_core/vector_store/rest_helpers.py

@@ -0,0 +1,81 @@
+import base64
+import json
+import socket
+from datetime import datetime
+
+import requests
+
+from frogml_core.clients.administration import AuthenticationClient
+from frogml_core.inner.di_configuration.account import UserAccountConfiguration
+
+
+def _get_authorization():
+    user_account_configuration = UserAccountConfiguration()
+    token = (
+        AuthenticationClient()
+        .authenticate(user_account_configuration.get_user_token())
+        .access_token
+    )
+    token_split = token.split(".")
+    decoded_token = json.loads(_base64url_decode(token_split[1]).decode("utf-8"))
+    token_expiration = datetime.fromtimestamp(decoded_token["exp"])
+
+    return f"Bearer {token}", token_expiration
+
+
+def _base64url_decode(input):
+    rem = len(input) % 4
+    if rem > 0:
+        input += "=" * (4 - rem)
+
+    return base64.urlsafe_b64decode(input)
+
+
+class SocketAdapter(requests.adapters.HTTPAdapter):
+    def __init__(self, *args, socket_params=None, **kwargs):
+        self.socket_params = socket_params
+        super().__init__(*args, **kwargs)
+
+    def init_poolmanager(self, *args, **kwargs):
+        if self.socket_params:
+            kwargs["socket_options"] = self.socket_params
+        super().init_poolmanager(*args, **kwargs)
+
+
+# Configure only TCP attributes that are available in the OS
+def validate_socket_config(socket_options):
+    config = []
+    for line in socket_options:
+        if hasattr(socket, line[0]) and hasattr(socket, line[1]):
+            config.append((getattr(socket, line[0]), getattr(socket, line[1]), line[2]))
+    return config
+
+
+class RestSession(requests.Session):
+    def __init__(self):
+        super().__init__()
+        self.headers.update({"Content-Type": "application/json"})
+        socket_options = [
+            ("SOL_SOCKET", "SO_KEEPALIVE", 1),
+            ("SOL_TCP", "TCP_KEEPIDLE", 120),
+            ("SOL_TCP", "TCP_KEEPINTVL", 75),
+            ("SOL_TCP", "TCP_KEEPCNT", 9),
+        ]
+        socket_options = validate_socket_config(socket_options)
+        adapter = SocketAdapter(socket_params=socket_options)
+        self.mount("https://", adapter)
+
+    def prepare_request(self, request):
+        if "Authorization" not in self.headers:
+            self.prepare_request_token()
+        else:
+            if self.jwt_expiration <= datetime.utcnow():
+                self.prepare_request_token()
+
+        return super().prepare_request(request)
+
+    def prepare_request_token(self):
+        (
+            self.headers["Authorization"],
+            self.jwt_expiration,
+        ) = _get_authorization()

frogml_core/vector_store/utils/filter_utils.py

@@ -0,0 +1,23 @@
+from datetime import datetime
+from typing import Any
+
+from frogml_proto.qwak.vectors.v1.filters_pb2 import (
+    AtomicLiteral as ProtoAtomicLiteral,
+)
+from frogml_core.utils.datetime_utils import datetime_to_pts
+
+
+def transform(value: Any) -> ProtoAtomicLiteral:
+    if isinstance(value, bool):
+        return ProtoAtomicLiteral(bool_literal=value)
+    elif isinstance(value, str):
+        return ProtoAtomicLiteral(string_literal=value)
+    elif isinstance(value, int):
+        return ProtoAtomicLiteral(int_literal=value)
+    elif isinstance(value, float):
+        return ProtoAtomicLiteral(double_literal=value)
+    elif isinstance(value, datetime):
+        # Assuming that timestamp is a datetime
+        return ProtoAtomicLiteral(timestamp_literal=datetime_to_pts(value))
+    else:
+        raise ValueError(f"Unsupported data type: {type(value)}")