Merge pull request #1 from aurelio-labs/simple_classification

Add simple_classify Method to DecisionLayer Class

.env.example

@@ -0,0 +1 @@
+COHERE_API_KEY=

.gitignore

@@ -1,4 +1,13 @@
-.env
-mac.env
-
+__pycache__
+*.pyc
+.venv
+.DS_Store
+venv/
+/.vscode
 **/__pycache__
+**/*.py[cod]
+
+# local env files
+.env*.local
+.env
+mac.env

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Aurelio AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1 +1 @@
-# Decision Layer
+# Semantic Router

pyproject.toml

@@ -1,8 +1,11 @@
 [tool.poetry]
-name = "decision-layer"
+name = "semantic-router"
 version = "0.0.1"
-description = "Super fast decision layer for AI"
-authors = ["James Briggs <[email protected]>"]
+description = "Super fast semantic router for AI decision making"
+authors = [
+    "James Briggs <[email protected]>",
+    "Siraj Aizlewood <[email protected]>"
+]
 readme = "README.md"
 
 [tool.poetry.dependencies]
@@ -15,6 +18,12 @@ openai = "^0.28.1"
 transformers = "^4.34.1"
 cohere = "^4.32"
 
+
+[tool.poetry.group.dev.dependencies]
+ipykernel = "^6.26.0"
+ruff = "^0.1.5"
+black = "^23.11.0"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

semantic_router/encoders/__init__.py

@@ -0,0 +1,6 @@
+from .base import BaseEncoder
+from .cohere import CohereEncoder
+from .huggingface import HuggingFaceEncoder
+from .openai import OpenAIEncoder
+
+__all__ = ["BaseEncoder", "CohereEncoder", "HuggingFaceEncoder", "OpenAIEncoder"]

decision_layer/encoders/base.py → semantic_router/encoders/base.py

@@ -8,4 +8,4 @@ class Config:
         arbitrary_types_allowed = True
 
     def __call__(self, texts: list[str]) -> list[float]:
-        pass
+        raise NotImplementedError("Subclasses must implement this method")

semantic_router/encoders/cohere.py

@@ -0,0 +1,28 @@
+import os
+
+import cohere
+
+from semantic_router.encoders import BaseEncoder
+
+
+class CohereEncoder(BaseEncoder):
+    client: cohere.Client | None
+
+    def __init__(
+        self, name: str = "embed-english-v3.0", cohere_api_key: str | None = None
+    ):
+        super().__init__(name=name)
+        cohere_api_key = cohere_api_key or os.getenv("COHERE_API_KEY")
+        if cohere_api_key is None:
+            raise ValueError("Cohere API key cannot be 'None'.")
+        self.client = cohere.Client(cohere_api_key)
+
+    def __call__(self, texts: list[str]) -> list[list[float]]:
+        if self.client is None:
+            raise ValueError("Cohere client is not initialized.")
+        if len(texts) == 1:
+            input_type = "search_query"
+        else:
+            input_type = "search_document"
+        embeds = self.client.embed(texts, input_type=input_type, model=self.name)
+        return embeds.embeddings

decision_layer/encoders/huggingface.py → semantic_router/encoders/huggingface.py

@@ -1,9 +1,9 @@
-from decision_layer.encoders import BaseEncoder
+from semantic_router.encoders import BaseEncoder
 
 
 class HuggingFaceEncoder(BaseEncoder):
     def __init__(self, name: str):
-        super().__init__(name)
+        self.name = name
 
     def __call__(self, texts: list[str]) -> list[float]:
-        raise NotImplementedError
+        raise NotImplementedError

decision_layer/encoders/openai.py → semantic_router/encoders/openai.py

@@ -1,8 +1,10 @@
 import os
+from time import sleep
 
-from decision_layer.encoders import BaseEncoder
 import openai
-from time import time
+from openai.error import RateLimitError
+
+from semantic_router.encoders import BaseEncoder
 
 
 class OpenAIEncoder(BaseEncoder):
@@ -16,20 +18,18 @@ def __call__(self, texts: list[str]) -> list[list[float]]:
         """Encode a list of texts using the OpenAI API. Returns a list of
         vector embeddings.
         """
-        passed = False
+        res = None
         # exponential backoff in case of RateLimitError
         for j in range(5):
             try:
-                # create embeddings
-                res = openai.Embedding.create(
-                    input=texts, engine=self.name
-                )
-                passed = True
-            except openai.error.RateLimitError:
-                time.sleep(2 ** j)
-        if not passed:
-            raise openai.error.RateLimitError
+                res = openai.Embedding.create(input=texts, engine=self.name)
+                if isinstance(res, dict) and "data" in res:
+                    break
+            except RateLimitError:
+                sleep(2**j)
+        if not res or not isinstance(res, dict) or "data" not in res:
+            raise ValueError("Failed to create embeddings.")
+
         # get embeddings
         embeds = [r["embedding"] for r in res["data"]]
         return embeds
-

semantic_router/layer.py

@@ -0,0 +1,104 @@
+import numpy as np
+from numpy.linalg import norm
+
+from semantic_router.encoders import BaseEncoder, CohereEncoder, OpenAIEncoder
+from semantic_router.schema import Decision
+
+
+class DecisionLayer:
+    index = None
+    categories = None
+    similarity_threshold = 0.82
+
+    def __init__(self, encoder: BaseEncoder, decisions: list[Decision] = []):
+        self.encoder = encoder
+        # decide on default threshold based on encoder
+        if isinstance(encoder, OpenAIEncoder):
+            self.similarity_threshold = 0.82
+        elif isinstance(encoder, CohereEncoder):
+            self.similarity_threshold = 0.3
+        else:
+            self.similarity_threshold = 0.82
+        # if decisions list has been passed, we initialize index now
+        if decisions:
+            # initialize index now
+            for decision in decisions:
+                self._add_decision(decision=decision)
+
+    def __call__(self, text: str) -> str | None:
+        results = self._query(text)
+        top_class, top_class_scores = self._semantic_classify(results)
+        passed = self._pass_threshold(top_class_scores, self.similarity_threshold)
+        if passed:
+            return top_class
+        else:
+            return None
+
+    def add(self, decision: Decision):
+        self._add_decision(decision=decision)
+
+    def _add_decision(self, decision: Decision):
+        # create embeddings
+        embeds = self.encoder(decision.utterances)
+
+        # create decision array
+        if self.categories is None:
+            self.categories = np.array([decision.name] * len(embeds))
+        else:
+            str_arr = np.array([decision.name] * len(embeds))
+            self.categories = np.concatenate([self.categories, str_arr])
+        # create utterance array (the index)
+        if self.index is None:
+            self.index = np.array(embeds)
+        else:
+            embed_arr = np.array(embeds)
+            self.index = np.concatenate([self.index, embed_arr])
+
+    def _query(self, text: str, top_k: int = 5):
+        """Given some text, encodes and searches the index vector space to
+        retrieve the top_k most similar records.
+        """
+        # create query vector
+        xq = np.array(self.encoder([text]))
+        xq = np.squeeze(xq)  # Reduce to 1d array.
+
+        if self.index is not None:
+            index_norm = norm(self.index, axis=1)
+            xq_norm = norm(xq.T)
+            sim = np.dot(self.index, xq.T) / (index_norm * xq_norm)
+            # get indices of top_k records
+            top_k = min(top_k, sim.shape[0])
+            idx = np.argpartition(sim, -top_k)[-top_k:]
+            scores = sim[idx]
+            # get the utterance categories (decision names)
+            decisions = self.categories[idx] if self.categories is not None else []
+            return [
+                {"decision": d, "score": s.item()} for d, s in zip(decisions, scores)
+            ]
+        else:
+            return []
+
+    def _semantic_classify(self, query_results: list[dict]) -> tuple[str, list[float]]:
+        scores_by_class = {}
+        for result in query_results:
+            score = result["score"]
+            decision = result["decision"]
+            if decision in scores_by_class:
+                scores_by_class[decision].append(score)
+            else:
+                scores_by_class[decision] = [score]
+
+        # Calculate total score for each class
+        total_scores = {
+            decision: sum(scores) for decision, scores in scores_by_class.items()
+        }
+        top_class = max(total_scores, key=lambda x: total_scores[x], default=None)
+
+        # Return the top class and its associated scores
+        return str(top_class), scores_by_class.get(top_class, [])
+
+    def _pass_threshold(self, scores: list[float], threshold: float) -> bool:
+        if scores:
+            return max(scores) > threshold
+        else:
+            return False