Merge pull request #1 from aurelio-labs/simple_classification

Add simple_classify Method to DecisionLayer Class

.env.example

+COHERE_API_KEY=
\ No newline at end of file

.gitignore

-.env
-mac.env
-
+__pycache__
+*.pyc
+.venv
+.DS_Store
+venv/
+/.vscode
 **/__pycache__
+**/*.py[cod]
+
+# local env files
+.env*.local
+.env
+mac.env
\ No newline at end of file

LICENSE

+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.
\ No newline at end of file

README.md

-# Decision Layer
+# Semantic Router
\ No newline at end of file

decision_layer/__init__.py

-from decision_layer.decision_layer import DecisionLayer
\ No newline at end of file

decision_layer/encoders/__init__.py

-from decision_layer.encoders.base import BaseEncoder
-from decision_layer.encoders.cohere import CohereEncoder
-from decision_layer.encoders.huggingface import HuggingFaceEncoder
-from decision_layer.encoders.openai import OpenAIEncoder
\ No newline at end of file

decision_layer/encoders/cohere.py

-from decision_layer.encoders import BaseEncoder
-
-class CohereEncoder(BaseEncoder):
-    def __init__(self, name: str):
-        super().__init__(name)
-
-    def __call__(self, texts: list[str]) -> list[float]:
-        raise NotImplementedError
\ No newline at end of file

pyproject.toml

 [tool.poetry]
-name = "decision-layer"
+name = "semantic-router"
 version = "0.0.1"
-description = "Super fast decision layer for AI"
-authors = ["James Briggs <james@aurelio.ai>"]
+description = "Super fast semantic router for AI decision making"
+authors = [
+    "James Briggs <james@aurelio.ai>",
+    "Siraj Aizlewood <siraj@aurelio.ai>"
+]
 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

+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 BaseEncoder(BaseModel):
         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

+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

-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
\ No newline at end of file

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

 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 @@ class OpenAIEncoder(BaseEncoder):
         """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
-        
\ No newline at end of file

decision_layer/decision_layer.py → semantic_router/layer.py

-from decision_layer.encoders import BaseEncoder
-from decision_layer.schema import Decision
 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):
+    def __call__(self, text: str) -> str | None:
         results = self._query(text)
-        raise NotImplementedError("To implement decision logic based on scores")
-
+        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(devision=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))
+            self.categories = np.array([decision.name] * len(embeds))
         else:
-            str_arr = np.array([decision.name]*len(embeds))
+            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:
@@ -40,19 +54,51 @@ class DecisionLayer:
             embed_arr = np.array(embeds)
             self.index = np.concatenate([self.index, embed_arr])
 
-    def _query(self, text: str, top_k: int=5):
+    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]))
-        # calculate cosine similarities
-        sim = np.dot(self.index, xq.T) / (norm(self.index)*norm(xq.T))
-        # get indices of top_k records
-        idx = np.argpartition(sim.T[0], -top_k)[-top_k:]
-        scores = sim[idx]
-        # get the utterance categories (decision names)
-        decisions = self.categories[idx]
-        return [
-            {"decision": d, "score": s.item()} for d, s in zip(decisions, scores)
-        ]
+        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

decision_layer/schema.py → semantic_router/schema.py

 from enum import Enum
+
 from pydantic import BaseModel
 from pydantic.dataclasses import dataclass
-from decision_layer.encoders import (
+
+from semantic_router.encoders import (
     BaseEncoder,
+    CohereEncoder,
     HuggingFaceEncoder,
     OpenAIEncoder,
-    CohereEncoder,
 )
 
 
@@ -14,11 +16,13 @@ class Decision(BaseModel):
     utterances: list[str]
     description: str | None = None
 
+
 class EncoderType(Enum):
     HUGGINGFACE = "huggingface"
     OPENAI = "openai"
     COHERE = "cohere"
 
+
 @dataclass
 class Encoder:
     type: EncoderType
@@ -38,6 +42,7 @@ class Encoder:
     def __call__(self, texts: list[str]) -> list[float]:
         return self.model(texts)
 
+
 @dataclass
 class SemanticSpace:
     id: str
@@ -49,4 +54,4 @@ class SemanticSpace:
         self.decisions = decisions
 
     def add(self, decision: Decision):
-        self.decisions.append(decision)
+        self.decisions.append(decision)
\ No newline at end of file

walkthrough.ipynb

 {
  "cells": [
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Decision Layer Walkthrough"
+    "# Semantic Router Walkthrough"
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The decision layer library can be used as a super fast decision making layer on top of LLMs. That means that rather than waiting on a slow agent to decide what to do, we can use the magic of semantic vector space to make decisions. Cutting decision making time down from seconds to milliseconds."
+    "The Semantic Router library can be used as a super fast decision making layer on top of LLMs. That means rather than waiting on a slow agent to decide what to do, we can use the magic of semantic vector space to make decisions. Cutting decision making time down from seconds to milliseconds."
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Getting Started"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We start by installing the library:"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "!pip install -qU \\\n",
-    "    decision-layer"
+    "!pip install -qU semantic-router==0.0.1"
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -44,11 +46,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from decision_layer.schema import Decision\n",
+    "from semantic_router.schema import Decision\n",
     "\n",
     "politics = Decision(\n",
     "    name=\"politics\",\n",
@@ -72,7 +74,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -94,71 +96,89 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now we initialize our embedding model (we will add support for Hugging Face):"
+    "Now we initialize our embedding model:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from decision_layer.encoders import OpenAIEncoder\n",
+    "from semantic_router.encoders import CohereEncoder\n",
+    "from getpass import getpass\n",
+    "\n",
     "import os\n",
     "\n",
-    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
-    "encoder = OpenAIEncoder(name=\"text-embedding-ada-002\")"
+    "# os.environ[\"COHERE_API_KEY\"] = getpass(\"Enter Cohere API Key: \")\n",
+    "os.environ[\"COHERE_API_KEY\"]\n",
+    "encoder = CohereEncoder()"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now we define the `DecisionLayer`. When called, the decision layer will consume text (a query) and output the category (`Decision`) it belongs to — for now we can only `_query` and get the most similar `Decision` `utterances`."
+    "Now we define the `DecisionLayer`. When called, the decision layer will consume text (a query) and output the category (`Decision`) it belongs to — to initialize a `DecisionLayer` we need our `encoder` model and a list of `decisions`."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from decision_layer import DecisionLayer\n",
+    "from semantic_router import DecisionLayer\n",
     "\n",
     "dl = DecisionLayer(encoder=encoder, decisions=decisions)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we can test it:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dl(\"don't you love politics?\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dl(\"how's the weather today?\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Both are classified accurately, what if we send a query that is unrelated to our existing `Decision` objects?"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "[{'decision': 'politics', 'score': 0.24968127755063652},\n",
-       " {'decision': 'politics', 'score': 0.2536216026530966},\n",
-       " {'decision': 'politics', 'score': 0.27568433588684954},\n",
-       " {'decision': 'politics', 'score': 0.27732789989574913},\n",
-       " {'decision': 'politics', 'score': 0.28110307885950714}]"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "out = dl._query(\"don't you love politics?\")\n",
-    "out"
-   ]
-  },
-  {
-   "attachments": {},
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dl(\"I'm interested in learning about llama 2\")"
+   ]
+  },
+  {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "---"
+    "In this case, we return `None` because no matches were identified."
    ]
   }
  ],
@@ -178,9 +198,8 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.5"
-  },
-  "orig_nbformat": 4
+   "version": "3.11.3"
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2