schema.py

from datetime import datetime, timezone
from difflib import Differ
from enum import Enum
import json
import numpy as np
from typing import List, Optional, Union, Any, Dict, Tuple
from pydantic import BaseModel, Field
from semantic_router.utils.logger import logger
from aurelio_sdk.schema import BM25Embedding


class EncoderType(Enum):
    AURELIO = "aurelio"
    AZURE = "azure"
    COHERE = "cohere"
    OPENAI = "openai"
    BM25 = "bm25"
    TFIDF = "tfidf"
    FASTEMBED = "fastembed"
    HUGGINGFACE = "huggingface"
    MISTRAL = "mistral"
    VIT = "vit"
    CLIP = "clip"
    GOOGLE = "google"
    BEDROCK = "bedrock"


class EncoderInfo(BaseModel):
    name: str
    token_limit: int
    threshold: Optional[float] = None


class RouteChoice(BaseModel):
    name: Optional[str] = None
    function_call: Optional[List[Dict]] = None
    similarity_score: Optional[float] = None


class Message(BaseModel):
    role: str
    content: str

    def to_openai(self):
        if self.role.lower() not in ["user", "assistant", "system"]:
            raise ValueError("Role must be either 'user', 'assistant' or 'system'")
        return {"role": self.role, "content": self.content}

    def to_cohere(self):
        return {"role": self.role, "message": self.content}

    def to_llamacpp(self):
        return {"role": self.role, "content": self.content}

    def to_mistral(self):
        return {"role": self.role, "content": self.content}

    def __str__(self):
        return f"{self.role}: {self.content}"


class ConfigParameter(BaseModel):
    field: str
    value: str
    scope: Optional[str] = None
    created_at: str = Field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )

    def to_pinecone(self, dimensions: int):
        namespace = self.scope or ""
        return {
            "id": f"{self.field}#{namespace}",
            "values": [0.1] * dimensions,
            "metadata": {
                "value": self.value,
                "created_at": self.created_at,
                "namespace": namespace,
                "field": self.field,
            },
        }


class Utterance(BaseModel):
    route: str
    utterance: Union[str, Any]
    function_schemas: Optional[List[Dict]] = None
    metadata: dict = {}
    diff_tag: str = " "

    @classmethod
    def from_tuple(cls, tuple_obj: Tuple):
        """Create an Utterance object from a tuple. The tuple must contain
        route and utterance as the first two elements. Then optionally
        function schemas and metadata as the third and fourth elements
        respectively. If this order is not followed an invalid Utterance
        object will be returned.

        :param tuple_obj: A tuple containing route, utterance, function schemas and metadata.
        :type tuple_obj: Tuple
        :return: An Utterance object.
        :rtype: Utterance
        """
        route, utterance = tuple_obj[0], tuple_obj[1]
        function_schemas = tuple_obj[2] if len(tuple_obj) > 2 else None
        if isinstance(function_schemas, dict):
            function_schemas = [function_schemas]
        metadata = tuple_obj[3] if len(tuple_obj) > 3 else {}
        return cls(
            route=route,
            utterance=utterance,
            function_schemas=function_schemas,
            metadata=metadata,
        )

    def to_tuple(self):
        """Convert an Utterance object to a tuple.

        :return: A tuple containing (route, utterance, function schemas, metadata).
        :rtype: Tuple
        """
        return (
            self.route,
            self.utterance,
            self.function_schemas,
            self.metadata,
        )

    def to_str(self, include_metadata: bool = False):
        if include_metadata:
            # we sort the dicts to ensure consistent order as we need this to compare
            # stringified function schemas accurately
            if self.function_schemas is not None:
                function_schemas_sorted: List[str] | None = [
                    json.dumps(schema, sort_keys=True)
                    for schema in self.function_schemas
                ]
            else:
                function_schemas_sorted = None
            # we must do the same for metadata
            metadata_sorted = json.dumps(self.metadata, sort_keys=True)
            return f"{self.route}: {self.utterance} | {function_schemas_sorted} | {metadata_sorted}"
        return f"{self.route}: {self.utterance}"

    def to_diff_str(self, include_metadata: bool = False):
        return f"{self.diff_tag} {self.to_str(include_metadata=include_metadata)}"


class SyncMode(Enum):
    """Synchronization modes for local (route layer) and remote (index)
    instances.
    """

    ERROR = "error"
    REMOTE = "remote"
    LOCAL = "local"
    MERGE_FORCE_REMOTE = "merge-force-remote"
    MERGE_FORCE_LOCAL = "merge-force-local"
    MERGE = "merge"


SYNC_MODES = [x.value for x in SyncMode]


class UtteranceDiff(BaseModel):
    diff: List[Utterance]

    @classmethod
    def from_utterances(
        cls, local_utterances: List[Utterance], remote_utterances: List[Utterance]
    ):
        local_utterances_map = {
            x.to_str(include_metadata=True): x for x in local_utterances
        }
        remote_utterances_map = {
            x.to_str(include_metadata=True): x for x in remote_utterances
        }
        # sort local and remote utterances
        local_utterances_str = list(local_utterances_map.keys())
        local_utterances_str.sort()
        remote_utterances_str = list(remote_utterances_map.keys())
        remote_utterances_str.sort()
        # get diff
        differ = Differ()
        diff_obj = list(differ.compare(local_utterances_str, remote_utterances_str))
        # create UtteranceDiff list
        utterance_diffs = []
        for line in diff_obj:
            utterance_str = line[2:]
            utterance_diff_tag = line[0]
            if utterance_diff_tag == "?":
                # this is a new line from diff string, we can ignore
                continue
            utterance = (
                remote_utterances_map[utterance_str]
                if utterance_diff_tag == "+"
                else local_utterances_map[utterance_str]
            )
            utterance.diff_tag = utterance_diff_tag
            utterance_diffs.append(utterance)
        return UtteranceDiff(diff=utterance_diffs)

    def to_utterance_str(self, include_metadata: bool = False) -> List[str]:
        """Outputs the utterance diff as a list of diff strings. Returns a list
        of strings showing what is different in the remote when compared to the
        local. For example:

        ["  route1: utterance1",
         "  route1: utterance2",
         "- route2: utterance3",
         "- route2: utterance4"]

        Tells us that the remote is missing "route2: utterance3" and "route2:
        utterance4", which do exist locally. If we see:

        ["  route1: utterance1",
         "  route1: utterance2",
         "+ route2: utterance3",
         "+ route2: utterance4"]

        This diff tells us that the remote has "route2: utterance3" and
        "route2: utterance4", which do not exist locally.
        """
        return [x.to_diff_str(include_metadata=include_metadata) for x in self.diff]

    def get_tag(self, diff_tag: str) -> List[Utterance]:
        """Get all utterances with a given diff tag.

        :param diff_tag: The diff tag to filter by. Must be one of "+", "-", or
        " ".
        :type diff_tag: str
        :return: A list of Utterance objects.
        :rtype: List[Utterance]
        """
        if diff_tag not in ["+", "-", " "]:
            raise ValueError("diff_tag must be one of '+', '-', or ' '")
        return [x for x in self.diff if x.diff_tag == diff_tag]

    def get_sync_strategy(self, sync_mode: str) -> dict:
        """Generates the optimal synchronization plan for local and remote
        instances.

        :param sync_mode: The mode to sync the routes with the remote index.
        :type sync_mode: str
        :return: A dictionary describing the synchronization strategy.
        :rtype: dict
        """
        if sync_mode not in SYNC_MODES:
            raise ValueError(f"sync_mode must be one of {SYNC_MODES}")
        local_only = self.get_tag("-")
        local_only_mapper = {
            utt.route: (utt.function_schemas, utt.metadata) for utt in local_only
        }
        remote_only = self.get_tag("+")
        remote_only_mapper = {
            utt.route: (utt.function_schemas, utt.metadata) for utt in remote_only
        }
        local_and_remote = self.get_tag(" ")
        if sync_mode == "error":
            if len(local_only) > 0 or len(remote_only) > 0:
                raise ValueError(
                    "There are utterances that exist in the local or remote "
                    "instance that do not exist in the other instance. Please "
                    "sync the routes before running this command."
                )
            else:
                return {
                    "remote": {"upsert": [], "delete": []},
                    "local": {"upsert": [], "delete": []},
                }
        elif sync_mode == "local":
            return {
                "remote": {
                    "upsert": local_only,  # + remote_updates,
                    "delete": remote_only,
                },
                "local": {"upsert": [], "delete": []},
            }
        elif sync_mode == "remote":
            return {
                "remote": {"upsert": [], "delete": []},
                "local": {"upsert": remote_only, "delete": local_only},
            }
        elif sync_mode == "merge-force-local":  # merge-to-local merge-join-local
            # PRIORITIZE LOCAL
            # get set of route names that exist in local (we keep these if
            # they are in remote)
            local_route_names = set([utt.route for utt in local_only])
            # if we see route: utterance exists in local, we do not pull it in
            # from remote
            local_route_utt_strs = set([utt.to_str() for utt in local_only])
            # get remote utterances that are in local
            remote_to_keep = [
                utt
                for utt in remote_only
                if (
                    utt.route in local_route_names
                    and utt.to_str() not in local_route_utt_strs
                )
            ]
            # overwrite remote routes with local metadata and function schemas
            logger.info(f"local_only_mapper: {local_only_mapper}")
            remote_to_update = [
                Utterance(
                    route=utt.route,
                    utterance=utt.utterance,
                    metadata=local_only_mapper[utt.route][1],
                    function_schemas=local_only_mapper[utt.route][0],
                )
                for utt in remote_only
                if (
                    utt.route in local_only_mapper
                    and (
                        utt.metadata != local_only_mapper[utt.route][1]
                        or utt.function_schemas != local_only_mapper[utt.route][0]
                    )
                )
            ]
            remote_to_keep = [
                Utterance(
                    route=utt.route,
                    utterance=utt.utterance,
                    metadata=local_only_mapper[utt.route][1],
                    function_schemas=local_only_mapper[utt.route][0],
                )
                for utt in remote_to_keep
                if utt.to_str() not in [x.to_str() for x in remote_to_update]
            ]
            # get remote utterances that are NOT in local
            remote_to_delete = [
                utt for utt in remote_only if utt.route not in local_route_names
            ]
            return {
                "remote": {
                    "upsert": local_only + remote_to_update,
                    "delete": remote_to_delete,
                },
                "local": {"upsert": remote_to_keep, "delete": []},
            }
        elif sync_mode == "merge-force-remote":  # merge-to-remote merge-join-remote
            # get set of route names that exist in remote (we keep these if
            # they are in local)
            remote_route_names = set([utt.route for utt in remote_only])
            # if we see route: utterance exists in remote, we do not pull it in
            # from local
            remote_route_utt_strs = set([utt.to_str() for utt in remote_only])
            # get local utterances that are in remote
            local_to_keep = [
                utt
                for utt in local_only
                if (
                    utt.route in remote_route_names
                    and utt.to_str() not in remote_route_utt_strs
                )
            ]
            # overwrite remote routes with local metadata and function schemas
            local_to_keep = [
                Utterance(
                    route=utt.route,
                    utterance=utt.utterance,
                    metadata=remote_only_mapper[utt.route][1],
                    function_schemas=remote_only_mapper[utt.route][0],
                )
                for utt in local_to_keep
            ]
            # get local utterances that are NOT in remote
            local_to_delete = [
                utt for utt in local_only if utt.route not in remote_route_names
            ]
            return {
                "remote": {"upsert": local_to_keep, "delete": []},
                "local": {"upsert": remote_only, "delete": local_to_delete},
            }
        elif sync_mode == "merge":
            # overwrite remote routes with local metadata and function schemas
            remote_only_updated = [
                (
                    Utterance(
                        route=utt.route,
                        utterance=utt.utterance,
                        metadata=local_only_mapper[utt.route][1],
                        function_schemas=local_only_mapper[utt.route][0],
                    )
                    if utt.route in local_only_mapper
                    else utt
                )
                for utt in remote_only
            ]
            # propogate same to shared routes
            shared_updated = [
                Utterance(
                    route=utt.route,
                    utterance=utt.utterance,
                    metadata=local_only_mapper[utt.route][1],
                    function_schemas=local_only_mapper[utt.route][0],
                )
                for utt in local_and_remote
                if (
                    utt.route in local_only_mapper
                    and (
                        utt.metadata != local_only_mapper[utt.route][1]
                        or utt.function_schemas != local_only_mapper[utt.route][0]
                    )
                )
            ]
            return {
                "remote": {
                    "upsert": local_only + shared_updated + remote_only_updated,
                    "delete": [],
                },
                "local": {"upsert": remote_only_updated + shared_updated, "delete": []},
            }
        else:
            raise ValueError(f"sync_mode must be one of {SYNC_MODES}")


class Metric(Enum):
    COSINE = "cosine"
    DOTPRODUCT = "dotproduct"
    EUCLIDEAN = "euclidean"
    MANHATTAN = "manhattan"


class SparseEmbedding(BaseModel):
    """Sparse embedding interface. Primarily uses numpy operations for faster
    operations.
    """

    embedding: np.ndarray

    class Config:
        arbitrary_types_allowed = True

    @classmethod
    def from_compact_array(cls, array: np.ndarray):
        if array.ndim != 2 or array.shape[1] != 2:
            raise ValueError(
                f"Expected a 2D array with 2 columns, got a {array.ndim}D array with {array.shape[1]} columns. "
                "Column 0 should contain index positions, and column 1 should contain respective values."
            )
        return cls(embedding=array)

    @classmethod
    def from_vector(cls, vector: np.ndarray):
        """Consumes an array of sparse vectors containing zero-values."""
        if vector.ndim != 1:
            raise ValueError(f"Expected a 1D array, got a {vector.ndim}D array.")
        return cls.from_compact_array(np.array([np.arange(len(vector)), vector]).T)

    @classmethod
    def from_aurelio(cls, embedding: BM25Embedding):
        arr = np.array([embedding.indices, embedding.values]).T
        return cls.from_compact_array(arr)

    @classmethod
    def from_dict(cls, sparse_dict: dict):
        arr = np.array([list(sparse_dict.keys()), list(sparse_dict.values())]).T
        return cls.from_compact_array(arr)

    @classmethod
    def from_pinecone_dict(cls, sparse_dict: dict):
        arr = np.array([sparse_dict["indices"], sparse_dict["values"]]).T
        return cls.from_compact_array(arr)

    def to_dict(self):
        return {
            i: v for i, v in zip(self.embedding[:, 0].astype(int), self.embedding[:, 1])
        }

    def to_pinecone(self):
        return {
            "indices": self.embedding[:, 0].astype(int).tolist(),
            "values": self.embedding[:, 1].tolist(),
        }

    # dictionary interface
    def items(self):
        return [
            (i, v)
            for i, v in zip(self.embedding[:, 0].astype(int), self.embedding[:, 1])
        ]