diff --git a/semantic_router/schema.py b/semantic_router/schema.py new file mode 100644 index 0000000000000000000000000000000000000000..b7a3c9faf7817e48215bcebdab8f3e7140aae970 --- /dev/null +++ b/semantic_router/schema.py @@ -0,0 +1,43 @@ +from enum import Enum + +from pydantic import BaseModel +from pydantic.dataclasses import dataclass + +from semantic_router.encoders import ( + BaseEncoder, + CohereEncoder, + OpenAIEncoder, +) + + +class EncoderType(Enum): + HUGGINGFACE = "huggingface" + OPENAI = "openai" + COHERE = "cohere" + + +class RouteChoice(BaseModel): + name: str | None = None + function_call: dict | None = None + + +@dataclass +class Encoder: + type: EncoderType + name: str | None + model: BaseEncoder + + def __init__(self, type: str, name: str | None): + self.type = EncoderType(type) + self.name = name + if self.type == EncoderType.HUGGINGFACE: + raise NotImplementedError + elif self.type == EncoderType.OPENAI: + self.model = OpenAIEncoder(name) + elif self.type == EncoderType.COHERE: + self.model = CohereEncoder(name) + else: + raise ValueError + + def __call__(self, texts: list[str]) -> list[list[float]]: + return self.model(texts)