Merge pull request #79 from aurelio-labs/ashraq/llm-module

feat: Add OpenAI, Cohere and OpenRouter as separate LLMs

README.md

@@ -12,7 +12,7 @@
 <img alt="Github License" src="https://img.shields.io/badge/License-MIT-yellow.svg" />
 </p>
 
-Semantic Router is a superfast decision layer for your LLMs and agents. Rather than waiting for slow LLM generations to make tool-use decisions, we use the magic of semantic vector space to make those decisions — _routing_ our requests using _semantic_ meaning.
+Semantic Router is a superfast decision-making layer for your LLMs and agents. Rather than waiting for slow LLM generations to make tool-use decisions, we use the magic of semantic vector space to make those decisions — _routing_ our requests using _semantic_ meaning.
 
 ## Quickstart
 
@@ -22,7 +22,9 @@ To get started with _semantic-router_ we install it like so:
 pip install -qU semantic-router
 ```
 
-We begin by defining a set of `Decision` objects. These are the decision paths that the semantic router can decide to use, let's try two simple decisions for now — one for talk on _politics_ and another for _chitchat_:
+❗️ _If wanting to use local embeddings you can use `FastEmbedEncoder` (`pip install -qU semantic-router[fastembed]`). To use the `HybridRouteLayer` you must `pip install -qU semantic-router[hybrid]`._
+
+We begin by defining a set of `Route` objects. These are the decision paths that the semantic router can decide to use, let's try two simple routes for now — one for talk on _politics_ and another for _chitchat_:
 
 ```python
 from semantic_router import Route
@@ -56,7 +58,7 @@ chitchat = Route(
 routes = [politics, chitchat]
 ```
 
-We have our decisions ready, now we initialize an embedding / encoder model. We currently support a `CohereEncoder` and `OpenAIEncoder` — more encoders will be added soon. To initialize them we do:
+We have our routes ready, now we initialize an embedding / encoder model. We currently support a `CohereEncoder` and `OpenAIEncoder` — more encoders will be added soon. To initialize them we do:
 
 ```python
 import os
@@ -71,18 +73,18 @@ os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>"
 encoder = OpenAIEncoder()
 ```
 
-With our `decisions` and `encoder` defined we now create a `DecisionLayer`. The decision layer handles our semantic decision making.
+With our `routes` and `encoder` defined we now create a `RouteLayer`. The route layer handles our semantic decision making.
 
 ```python
 from semantic_router.layer import RouteLayer
 
-dl = RouteLayer(encoder=encoder, routes=routes)
+rl = RouteLayer(encoder=encoder, routes=routes)
 ```
 
-We can now use our decision layer to make super fast decisions based on user queries. Let's try with two queries that should trigger our decisions:
+We can now use our route layer to make super fast decisions based on user queries. Let's try with two queries that should trigger our route decisions:
 
 ```python
-dl("don't you love politics?").name
+rl("don't you love politics?").name
 ```
 
 ```
@@ -92,7 +94,7 @@ dl("don't you love politics?").name
 Correct decision, let's try another:
 
 ```python
-dl("how's the weather today?").name
+rl("how's the weather today?").name
 ```
 
 ```
@@ -102,14 +104,14 @@ dl("how's the weather today?").name
 We get both decisions correct! Now lets try sending an unrelated query:
 
 ```python
-dl("I'm interested in learning about llama 2").name
+rl("I'm interested in learning about llama 2").name
 ```
 
 ```
 [Out]:
 ```
 
-In this case, no decision could be made as we had no matches — so our decision layer returned `None`!
+In this case, no decision could be made as we had no matches — so our route layer returned `None`!
 
 ## 📚 [Resources](https://github.com/aurelio-labs/semantic-router/tree/main/docs)
 

docs/02-dynamic-routes.ipynb

@@ -36,7 +36,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!pip install -qU semantic-router==0.0.14"
+    "!pip install -qU semantic-router==0.0.15"
    ]
   },
   {
@@ -64,17 +64,7 @@
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/Users/jamesbriggs/opt/anaconda3/envs/decision-layer/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n",
-      "None of PyTorch, TensorFlow >= 2.0, or Flax have been found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from semantic_router import Route\n",
     "\n",
@@ -102,30 +92,45 @@
     "routes = [politics, chitchat]"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We initialize our `RouteLayer` with our `encoder` and `routes`. We can use popular encoder APIs like `CohereEncoder` and `OpenAIEncoder`, or local alternatives like `FastEmbedEncoder`."
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[32m2023-12-28 19:19:39 INFO semantic_router.utils.logger Initializing RouteLayer\u001b[0m\n"
+      "\u001b[32m2024-01-07 15:23:12 INFO semantic_router.utils.logger Initializing RouteLayer\u001b[0m\n"
      ]
     }
    ],
    "source": [
     "import os\n",
     "from getpass import getpass\n",
     "from semantic_router import RouteLayer\n",
+    "from semantic_router.encoders import CohereEncoder, OpenAIEncoder\n",
     "\n",
     "# dashboard.cohere.ai\n",
-    "os.environ[\"COHERE_API_KEY\"] = os.getenv(\"COHERE_API_KEY\") or getpass(\n",
-    "    \"Enter Cohere API Key: \"\n",
+    "# os.environ[\"COHERE_API_KEY\"] = os.getenv(\"COHERE_API_KEY\") or getpass(\n",
+    "#     \"Enter Cohere API Key: \"\n",
+    "# )\n",
+    "# platform.openai.com\n",
+    "os.environ[\"OPENAI_API_KEY\"] = os.getenv(\"OPENAI_API_KEY\") or getpass(\n",
+    "    \"Enter OpenAI API Key: \"\n",
     ")\n",
     "\n",
-    "rl = RouteLayer(routes=routes)"
+    "# encoder = CohereEncoder()\n",
+    "encoder = OpenAIEncoder()\n",
+    "\n",
+    "rl = RouteLayer(encoder=encoder, routes=routes)"
    ]
   },
   {
@@ -137,7 +142,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -146,7 +151,7 @@
        "RouteChoice(name='chitchat', function_call=None)"
       ]
      },
-     "execution_count": 5,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -171,7 +176,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -193,16 +198,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "'13:19'"
+       "'09:23'"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -220,7 +225,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -232,7 +237,7 @@
        " 'output': \"<class 'str'>\"}"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -253,7 +258,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -277,16 +282,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
+     "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Adding route `get_time`\n",
-      "Adding route to categories\n",
-      "Adding route to index\n"
+      "\u001b[32m2024-01-07 15:23:16 INFO semantic_router.utils.logger Adding `get_time` route\u001b[0m\n"
      ]
     }
    ],
@@ -303,31 +306,32 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[32m2023-12-28 19:21:58 INFO semantic_router.utils.logger Extracting function input...\u001b[0m\n"
+      "\u001b[33m2024-01-07 15:23:17 WARNING semantic_router.utils.logger No LLM provided for dynamic route, will use OpenAI LLM default. Ensure API key is set in OPENAI_API_KEY environment variable.\u001b[0m\n",
+      "\u001b[32m2024-01-07 15:23:17 INFO semantic_router.utils.logger Extracting function input...\u001b[0m\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "RouteChoice(name='get_time', function_call={'timezone': 'America/New_York'})"
+       "RouteChoice(name='get_time', function_call={'timezone': 'new york city'})"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "# https://openrouter.ai/keys\n",
-    "os.environ[\"OPENROUTER_API_KEY\"] = os.getenv(\"OPENROUTER_API_KEY\") or getpass(\n",
-    "    \"Enter OpenRouter API Key: \"\n",
+    "# https://platform.openai.com/\n",
+    "os.environ[\"OPENAI_API_KEY\"] = os.getenv(\"OPENAI_API_KEY\") or getpass(\n",
+    "    \"Enter OpenAI API Key: \"\n",
     ")\n",
     "\n",
     "rl(\"what is the time in new york city?\")"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "semantic-router"
-version = "0.0.14"
+version = "0.0.15"
 description = "Super fast semantic router for AI decision making"
 authors = [
     "James Briggs <james@aurelio.ai>",

semantic_router/layer.py

@@ -7,10 +7,11 @@
 from semantic_router.encoders import (
     BaseEncoder,
     CohereEncoder,
-    OpenAIEncoder,
     FastEmbedEncoder,
+    OpenAIEncoder,
 )
 from semantic_router.linear import similarity_matrix, top_scores
+from semantic_router.llms import BaseLLM, OpenAILLM
 from semantic_router.route import Route
 from semantic_router.schema import Encoder, EncoderType, RouteChoice
 from semantic_router.utils.logger import logger
@@ -156,12 +157,16 @@ class RouteLayer:
     score_threshold: float = 0.82
 
     def __init__(
-        self, encoder: BaseEncoder | None = None, routes: list[Route] | None = None
+        self,
+        encoder: BaseEncoder | None = None,
+        llm: BaseLLM | None = None,
+        routes: list[Route] | None = None,
     ):
         logger.info("Initializing RouteLayer")
         self.index = None
         self.categories = None
         self.encoder = encoder if encoder is not None else CohereEncoder()
+        self.llm = llm
         self.routes: list[Route] = routes if routes is not None else []
         # decide on default threshold based on encoder
         # TODO move defaults to the encoder objects and extract from there
@@ -186,6 +191,17 @@ def __call__(self, text: str) -> RouteChoice:
         if passed:
             # get chosen route object
             route = [route for route in self.routes if route.name == top_class][0]
+            if route.function_schema and not isinstance(route.llm, BaseLLM):
+                if not self.llm:
+                    logger.warning(
+                        "No LLM provided for dynamic route, will use OpenAI LLM "
+                        "default. Ensure API key is set in OPENAI_API_KEY environment "
+                        "variable."
+                    )
+                    self.llm = OpenAILLM()
+                    route.llm = self.llm
+                else:
+                    route.llm = self.llm
             return route(text)
         else:
             # if no route passes threshold, return empty route choice
@@ -216,24 +232,20 @@ def from_config(cls, config: LayerConfig):
         return cls(encoder=encoder, routes=config.routes)
 
     def add(self, route: Route):
-        print(f"Adding route `{route.name}`")
+        logger.info(f"Adding `{route.name}` route")
         # create embeddings
         embeds = self.encoder(route.utterances)
 
         # create route array
         if self.categories is None:
-            print("Initializing categories array")
             self.categories = np.array([route.name] * len(embeds))
         else:
-            print("Adding route to categories")
             str_arr = np.array([route.name] * len(embeds))
             self.categories = np.concatenate([self.categories, str_arr])
         # create utterance array (the index)
         if self.index is None:
-            print("Initializing index array")
             self.index = np.array(embeds)
         else:
-            print("Adding route to index")
             embed_arr = np.array(embeds)
             self.index = np.concatenate([self.index, embed_arr])
         # add route to routes list

semantic_router/llms/__init__.py

@@ -0,0 +1,6 @@
+from semantic_router.llms.base import BaseLLM
+from semantic_router.llms.cohere import CohereLLM
+from semantic_router.llms.openai import OpenAILLM
+from semantic_router.llms.openrouter import OpenRouterLLM
+
+__all__ = ["BaseLLM", "OpenAILLM", "OpenRouterLLM", "CohereLLM"]

semantic_router/llms/base.py

@@ -0,0 +1,13 @@
+from pydantic import BaseModel
+
+from semantic_router.schema import Message
+
+
+class BaseLLM(BaseModel):
+    name: str
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def __call__(self, messages: list[Message]) -> str | None:
+        raise NotImplementedError("Subclasses must implement this method")