From 6815425b584782f90d7971c413e1c531939119f7 Mon Sep 17 00:00:00 2001
From: Ismail Ashraq
Date: Sat, 6 Jan 2024 19:21:04 +0500
Subject: [PATCH 1/8] Add llms
---
semantic_router/layer.py | 14 ++++++-
semantic_router/llms/__init__.py | 7 ++++
semantic_router/llms/base.py | 11 +++++
semantic_router/llms/cohere.py | 43 ++++++++++++++++++++
semantic_router/llms/openai.py | 51 +++++++++++++++++++++++
semantic_router/llms/openrouter.py | 56 ++++++++++++++++++++++++++
semantic_router/route.py | 26 ++++++++----
semantic_router/schema.py | 8 ++++
semantic_router/utils/function_call.py | 20 +++++----
9 files changed, 220 insertions(+), 16 deletions(-)
create mode 100644 semantic_router/llms/__init__.py
create mode 100644 semantic_router/llms/base.py
create mode 100644 semantic_router/llms/cohere.py
create mode 100644 semantic_router/llms/openai.py
create mode 100644 semantic_router/llms/openrouter.py
diff --git a/semantic_router/layer.py b/semantic_router/layer.py
index 5b2aad84..46eb2eb7 100644
--- a/semantic_router/layer.py
+++ b/semantic_router/layer.py
@@ -9,6 +9,7 @@
CohereEncoder,
OpenAIEncoder,
)
+from semantic_router.llms import BaseLLM
from semantic_router.linear import similarity_matrix, top_scores
from semantic_router.route import Route
from semantic_router.schema import Encoder, EncoderType, RouteChoice
@@ -142,12 +143,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
if isinstance(encoder, OpenAIEncoder):
@@ -168,6 +173,13 @@ 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:
+ raise ValueError(
+ "LLM is required for dynamic routes. Please ensure the 'llm' is set."
+ )
+ else:
+ route.llm = self.llm
return route(text)
else:
# if no route passes threshold, return empty route choice
diff --git a/semantic_router/llms/__init__.py b/semantic_router/llms/__init__.py
new file mode 100644
index 00000000..446f7c42
--- /dev/null
+++ b/semantic_router/llms/__init__.py
@@ -0,0 +1,7 @@
+from semantic_router.llms.base import BaseLLM
+from semantic_router.llms.openai import OpenAI
+from semantic_router.llms.openrouter import OpenRouter
+from semantic_router.llms.cohere import Cohere
+
+
+__all__ = ["BaseLLM", "OpenAI", "OpenRouter", "Cohere"]
diff --git a/semantic_router/llms/base.py b/semantic_router/llms/base.py
new file mode 100644
index 00000000..2a1a038e
--- /dev/null
+++ b/semantic_router/llms/base.py
@@ -0,0 +1,11 @@
+from pydantic import BaseModel
+
+
+class BaseLLM(BaseModel):
+ name: str
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def __call__(self, prompt) -> str | None:
+ raise NotImplementedError("Subclasses must implement this method")
diff --git a/semantic_router/llms/cohere.py b/semantic_router/llms/cohere.py
new file mode 100644
index 00000000..80512d5c
--- /dev/null
+++ b/semantic_router/llms/cohere.py
@@ -0,0 +1,43 @@
+import os
+import cohere
+from semantic_router.llms import BaseLLM
+from semantic_router.schema import Message
+
+
+class Cohere(BaseLLM):
+ client: cohere.Client | None = None
+
+ def __init__(
+ self,
+ name: str | None = None,
+ cohere_api_key: str | None = None,
+ ):
+ if name is None:
+ name = os.getenv("COHERE_CHAT_MODEL_NAME", "command")
+ 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'.")
+ try:
+ self.client = cohere.Client(cohere_api_key)
+ except Exception as e:
+ raise ValueError(f"Cohere API client failed to initialize. Error: {e}")
+
+ def __call__(self, messages: list[Message]) -> str:
+ if self.client is None:
+ raise ValueError("Cohere client is not initialized.")
+ try:
+ completion = self.client.chat(
+ model=self.name,
+ chat_history=[m.to_cohere() for m in messages[:-1]],
+ message=messages[-1].content,
+ )
+
+ output = completion.text
+
+ if not output:
+ raise Exception("No output generated")
+ return output
+
+ except Exception as e:
+ raise ValueError(f"Cohere API call failed. Error: {e}")
diff --git a/semantic_router/llms/openai.py b/semantic_router/llms/openai.py
new file mode 100644
index 00000000..18b6e706
--- /dev/null
+++ b/semantic_router/llms/openai.py
@@ -0,0 +1,51 @@
+import os
+import openai
+from semantic_router.utils.logger import logger
+from semantic_router.llms import BaseLLM
+from semantic_router.schema import Message
+
+
+class OpenAI(BaseLLM):
+ client: openai.OpenAI | None
+ temperature: float | None
+ max_tokens: int | None
+
+ def __init__(
+ self,
+ name: str | None = None,
+ openai_api_key: str | None = None,
+ temperature: float = 0.01,
+ max_tokens: int = 200,
+ ):
+ if name is None:
+ name = os.getenv("OPENAI_CHAT_MODEL_NAME", "gpt-3.5-turbo")
+ super().__init__(name=name)
+ api_key = openai_api_key or os.getenv("OPENAI_API_KEY")
+ if api_key is None:
+ raise ValueError("OpenAI API key cannot be 'None'.")
+ try:
+ self.client = openai.OpenAI(api_key=api_key)
+ except Exception as e:
+ raise ValueError(f"OpenAI API client failed to initialize. Error: {e}")
+ self.temperature = temperature
+ self.max_tokens = max_tokens
+
+ def __call__(self, messages: list[Message]) -> str:
+ if self.client is None:
+ raise ValueError("OpenAI client is not initialized.")
+ try:
+ completion = self.client.chat.completions.create(
+ model=self.name,
+ messages=[m.to_openai() for m in messages],
+ temperature=self.temperature,
+ max_tokens=self.max_tokens,
+ )
+
+ output = completion.choices[0].message.content
+
+ if not output:
+ raise Exception("No output generated")
+ return output
+ except Exception as e:
+ logger.error(f"LLM error: {e}")
+ raise Exception(f"LLM error: {e}")
diff --git a/semantic_router/llms/openrouter.py b/semantic_router/llms/openrouter.py
new file mode 100644
index 00000000..3b7a9b49
--- /dev/null
+++ b/semantic_router/llms/openrouter.py
@@ -0,0 +1,56 @@
+import os
+import openai
+from semantic_router.utils.logger import logger
+from semantic_router.llms import BaseLLM
+from semantic_router.schema import Message
+
+
+class OpenRouter(BaseLLM):
+ client: openai.OpenAI | None
+ base_url: str | None
+ temperature: float | None
+ max_tokens: int | None
+
+ def __init__(
+ self,
+ name: str | None = None,
+ openrouter_api_key: str | None = None,
+ base_url: str = "https://openrouter.ai/api/v1",
+ temperature: float = 0.01,
+ max_tokens: int = 200,
+ ):
+ if name is None:
+ name = os.getenv(
+ "OPENROUTER_CHAT_MODEL_NAME", "mistralai/mistral-7b-instruct"
+ )
+ super().__init__(name=name)
+ self.base_url = base_url
+ api_key = openrouter_api_key or os.getenv("OPENROUTER_API_KEY")
+ if api_key is None:
+ raise ValueError("OpenRouter API key cannot be 'None'.")
+ try:
+ self.client = openai.OpenAI(api_key=api_key, base_url=self.base_url)
+ except Exception as e:
+ raise ValueError(f"OpenRouter API client failed to initialize. Error: {e}")
+ self.temperature = temperature
+ self.max_tokens = max_tokens
+
+ def __call__(self, messages: list[Message]) -> str:
+ if self.client is None:
+ raise ValueError("OpenRouter client is not initialized.")
+ try:
+ completion = self.client.chat.completions.create(
+ model=self.name,
+ messages=[m.to_openai() for m in messages],
+ temperature=self.temperature,
+ max_tokens=self.max_tokens,
+ )
+
+ output = completion.choices[0].message.content
+
+ if not output:
+ raise Exception("No output generated")
+ return output
+ except Exception as e:
+ logger.error(f"LLM error: {e}")
+ raise Exception(f"LLM error: {e}")
diff --git a/semantic_router/route.py b/semantic_router/route.py
index 12afa7fe..7a8803d7 100644
--- a/semantic_router/route.py
+++ b/semantic_router/route.py
@@ -4,11 +4,13 @@
from pydantic import BaseModel
+from semantic_router.llms import BaseLLM
from semantic_router.schema import RouteChoice
from semantic_router.utils import function_call
-from semantic_router.utils.llm import llm
from semantic_router.utils.logger import logger
+from semantic_router.schema import Message
+
def is_valid(route_config: str) -> bool:
try:
@@ -43,12 +45,17 @@ class Route(BaseModel):
utterances: list[str]
description: str | None = None
function_schema: dict[str, Any] | None = None
+ llm: BaseLLM | None = None
def __call__(self, query: str) -> RouteChoice:
if self.function_schema:
+ if not self.llm:
+ raise ValueError(
+ "LLM is required for dynamic routes. Please ensure the 'llm' is set."
+ )
# if a function schema is provided we generate the inputs
extracted_inputs = function_call.extract_function_inputs(
- query=query, function_schema=self.function_schema
+ query=query, llm=self.llm, function_schema=self.function_schema
)
func_call = extracted_inputs
else:
@@ -60,16 +67,16 @@ def to_dict(self):
return self.dict()
@classmethod
- def from_dict(cls, data: dict):
+ def from_dict(cls, data: dict[str, Any]):
return cls(**data)
@classmethod
- def from_dynamic_route(cls, entity: Union[BaseModel, Callable]):
+ def from_dynamic_route(cls, llm: BaseLLM, entity: Union[BaseModel, Callable]):
"""
Generate a dynamic Route object from a function or Pydantic model using LLM
"""
schema = function_call.get_schema(item=entity)
- dynamic_route = cls._generate_dynamic_route(function_schema=schema)
+ dynamic_route = cls._generate_dynamic_route(llm=llm, function_schema=schema)
return dynamic_route
@classmethod
@@ -85,7 +92,7 @@ def _parse_route_config(cls, config: str) -> str:
raise ValueError("No tags found in the output.")
@classmethod
- def _generate_dynamic_route(cls, function_schema: dict[str, Any]):
+ def _generate_dynamic_route(cls, llm: BaseLLM, function_schema: dict[str, Any]):
logger.info("Generating dynamic route...")
prompt = f"""
@@ -113,7 +120,8 @@ def _generate_dynamic_route(cls, function_schema: dict[str, Any]):
{function_schema}
"""
- output = llm(prompt)
+ llm_input = [Message(role="user", content=prompt)]
+ output = llm(llm_input)
if not output:
raise Exception("No output generated for dynamic route")
@@ -122,5 +130,7 @@ def _generate_dynamic_route(cls, function_schema: dict[str, Any]):
logger.info(f"Generated route config:\n{route_config}")
if is_valid(route_config):
- return Route.from_dict(json.loads(route_config))
+ route_config_dict = json.loads(route_config)
+ route_config_dict["llm"] = llm
+ return Route.from_dict(route_config_dict)
raise Exception("No config generated")
diff --git a/semantic_router/schema.py b/semantic_router/schema.py
index 465cfaac..62eecc7d 100644
--- a/semantic_router/schema.py
+++ b/semantic_router/schema.py
@@ -49,6 +49,14 @@ 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}
+
class Conversation(BaseModel):
messages: list[Message]
diff --git a/semantic_router/utils/function_call.py b/semantic_router/utils/function_call.py
index 2ead3ab5..d93d027c 100644
--- a/semantic_router/utils/function_call.py
+++ b/semantic_router/utils/function_call.py
@@ -4,7 +4,8 @@
from pydantic import BaseModel
-from semantic_router.utils.llm import llm
+from semantic_router.llms import BaseLLM
+from semantic_router.schema import Message
from semantic_router.utils.logger import logger
@@ -40,7 +41,9 @@ def get_schema(item: Union[BaseModel, Callable]) -> dict[str, Any]:
return schema
-def extract_function_inputs(query: str, function_schema: dict[str, Any]) -> dict:
+def extract_function_inputs(
+ query: str, llm: BaseLLM, function_schema: dict[str, Any]
+) -> dict:
logger.info("Extracting function input...")
prompt = f"""
@@ -71,8 +74,8 @@ def extract_function_inputs(query: str, function_schema: dict[str, Any]) -> dict
schema: {function_schema}
Result:
"""
-
- output = llm(prompt)
+ llm_input = [Message(role="user", content=prompt)]
+ output = llm(llm_input)
if not output:
raise Exception("No output generated for extract function input")
@@ -113,15 +116,18 @@ def call_function(function: Callable, inputs: dict[str, str]):
# TODO: Add route layer object to the input, solve circular import issue
-async def route_and_execute(query: str, functions: list[Callable], route_layer):
+async def route_and_execute(
+ query: str, llm: BaseLLM, functions: list[Callable], route_layer
+):
function_name = route_layer(query)
if not function_name:
logger.warning("No function found, calling LLM...")
- return llm(query)
+ llm_input = [Message(role="user", content=query)]
+ return llm(llm_input)
for function in functions:
if function.__name__ == function_name:
print(f"Calling function: {function.__name__}")
schema = get_schema(function)
- inputs = extract_function_inputs(query, schema)
+ inputs = extract_function_inputs(query, llm, schema)
call_function(function, inputs)
From 5141474a1670f09e020fa8defa7d04502fb5fe58 Mon Sep 17 00:00:00 2001
From: Ismail Ashraq
Date: Sun, 7 Jan 2024 01:21:39 +0500
Subject: [PATCH 2/8] fix test issues
---
tests/unit/test_route.py | 41 ++++++++++++++++++----------------------
1 file changed, 18 insertions(+), 23 deletions(-)
diff --git a/tests/unit/test_route.py b/tests/unit/test_route.py
index 09a5d235..2eb784d4 100644
--- a/tests/unit/test_route.py
+++ b/tests/unit/test_route.py
@@ -1,6 +1,7 @@
-from unittest.mock import Mock, patch # , AsyncMock
+from unittest.mock import patch # , AsyncMock
# import pytest
+from semantic_router.llms import BaseLLM
from semantic_router.route import Route, is_valid
@@ -41,11 +42,9 @@ def test_is_valid_with_invalid_json():
mock_logger.error.assert_called_once()
-class TestRoute:
- @patch("semantic_router.route.llm", new_callable=Mock)
- def test_generate_dynamic_route(self, mock_llm):
- print(f"mock_llm: {mock_llm}")
- mock_llm.return_value = """
+class MockLLM(BaseLLM):
+ def __call__(self, prompt):
+ llm_output = """
{
"name": "test_function",
@@ -58,8 +57,16 @@ def test_generate_dynamic_route(self, mock_llm):
}
"""
+ return llm_output
+
+
+class TestRoute:
+ def test_generate_dynamic_route(self):
+ mock_llm = MockLLM(name="test")
function_schema = {"name": "test_function", "type": "function"}
- route = Route._generate_dynamic_route(function_schema)
+ route = Route._generate_dynamic_route(
+ llm=mock_llm, function_schema=function_schema
+ )
assert route.name == "test_function"
assert route.utterances == [
"example_utterance_1",
@@ -105,6 +112,7 @@ def test_to_dict(self):
"utterances": ["utterance"],
"description": None,
"function_schema": None,
+ "llm": None,
}
assert route.to_dict() == expected_dict
@@ -114,28 +122,15 @@ def test_from_dict(self):
assert route.name == "test"
assert route.utterances == ["utterance"]
- @patch("semantic_router.route.llm", new_callable=Mock)
- def test_from_dynamic_route(self, mock_llm):
+ def test_from_dynamic_route(self):
# Mock the llm function
- mock_llm.return_value = """
-
- {
- "name": "test_function",
- "utterances": [
- "example_utterance_1",
- "example_utterance_2",
- "example_utterance_3",
- "example_utterance_4",
- "example_utterance_5"]
- }
-
- """
+ mock_llm = MockLLM(name="test")
def test_function(input: str):
"""Test function docstring"""
pass
- dynamic_route = Route.from_dynamic_route(test_function)
+ dynamic_route = Route.from_dynamic_route(llm=mock_llm, entity=test_function)
assert dynamic_route.name == "test_function"
assert dynamic_route.utterances == [
From 548bd40393b8e4483d9092359051c699558c709f Mon Sep 17 00:00:00 2001
From: Ismail Ashraq
Date: Sun, 7 Jan 2024 03:51:54 +0500
Subject: [PATCH 3/8] add tests for llms
---
semantic_router/llms/base.py | 3 +-
tests/unit/llms/test_llm_base.py | 16 +++++++
tests/unit/llms/test_llm_cohere.py | 52 +++++++++++++++++++++++
tests/unit/llms/test_llm_openai.py | 55 ++++++++++++++++++++++++
tests/unit/llms/test_llm_openrouter.py | 59 ++++++++++++++++++++++++++
tests/unit/test_route.py | 17 +++++++-
tests/unit/test_schema.py | 27 +++++++++++-
7 files changed, 226 insertions(+), 3 deletions(-)
create mode 100644 tests/unit/llms/test_llm_base.py
create mode 100644 tests/unit/llms/test_llm_cohere.py
create mode 100644 tests/unit/llms/test_llm_openai.py
create mode 100644 tests/unit/llms/test_llm_openrouter.py
diff --git a/semantic_router/llms/base.py b/semantic_router/llms/base.py
index 2a1a038e..dd8a0afa 100644
--- a/semantic_router/llms/base.py
+++ b/semantic_router/llms/base.py
@@ -1,4 +1,5 @@
from pydantic import BaseModel
+from semantic_router.schema import Message
class BaseLLM(BaseModel):
@@ -7,5 +8,5 @@ class BaseLLM(BaseModel):
class Config:
arbitrary_types_allowed = True
- def __call__(self, prompt) -> str | None:
+ def __call__(self, messages: list[Message]) -> str | None:
raise NotImplementedError("Subclasses must implement this method")
diff --git a/tests/unit/llms/test_llm_base.py b/tests/unit/llms/test_llm_base.py
new file mode 100644
index 00000000..df78d8f5
--- /dev/null
+++ b/tests/unit/llms/test_llm_base.py
@@ -0,0 +1,16 @@
+import pytest
+
+from semantic_router.llms import BaseLLM
+
+
+class TestBaseLLM:
+ @pytest.fixture
+ def base_llm(self):
+ return BaseLLM(name="TestLLM")
+
+ def test_base_llm_initialization(self, base_llm):
+ assert base_llm.name == "TestLLM", "Initialization of name failed"
+
+ def test_base_llm_call_method_not_implemented(self, base_llm):
+ with pytest.raises(NotImplementedError):
+ base_llm("test")
diff --git a/tests/unit/llms/test_llm_cohere.py b/tests/unit/llms/test_llm_cohere.py
new file mode 100644
index 00000000..32443f04
--- /dev/null
+++ b/tests/unit/llms/test_llm_cohere.py
@@ -0,0 +1,52 @@
+import pytest
+
+from semantic_router.llms import Cohere
+from semantic_router.schema import Message
+
+
+@pytest.fixture
+def cohere_llm(mocker):
+ mocker.patch("cohere.Client")
+ return Cohere(cohere_api_key="test_api_key")
+
+
+class TestCohereLLM:
+ def test_initialization_with_api_key(self, cohere_llm):
+ assert cohere_llm.client is not None, "Client should be initialized"
+ assert cohere_llm.name == "command", "Default name not set correctly"
+
+ def test_initialization_without_api_key(self, mocker, monkeypatch):
+ monkeypatch.delenv("COHERE_API_KEY", raising=False)
+ mocker.patch("cohere.Client")
+ with pytest.raises(ValueError):
+ Cohere()
+
+ def test_call_method(self, cohere_llm, mocker):
+ mock_llm = mocker.MagicMock()
+ mock_llm.text = "test"
+ cohere_llm.client.chat.return_value = mock_llm
+
+ llm_input = [Message(role="user", content="test")]
+ result = cohere_llm(llm_input)
+ assert isinstance(result, str), "Result should be a str"
+ cohere_llm.client.chat.assert_called_once()
+
+ def test_raises_value_error_if_cohere_client_fails_to_initialize(self, mocker):
+ mocker.patch(
+ "cohere.Client", side_effect=Exception("Failed to initialize client")
+ )
+ with pytest.raises(ValueError):
+ Cohere(cohere_api_key="test_api_key")
+
+ def test_raises_value_error_if_cohere_client_is_not_initialized(self, mocker):
+ mocker.patch("cohere.Client", return_value=None)
+ llm = Cohere(cohere_api_key="test_api_key")
+ with pytest.raises(ValueError):
+ llm("test")
+
+ def test_call_method_raises_error_on_api_failure(self, cohere_llm, mocker):
+ mocker.patch.object(
+ cohere_llm.client, "__call__", side_effect=Exception("API call failed")
+ )
+ with pytest.raises(ValueError):
+ cohere_llm("test")
diff --git a/tests/unit/llms/test_llm_openai.py b/tests/unit/llms/test_llm_openai.py
new file mode 100644
index 00000000..4b2b2f54
--- /dev/null
+++ b/tests/unit/llms/test_llm_openai.py
@@ -0,0 +1,55 @@
+import pytest
+from semantic_router.llms import OpenAI
+from semantic_router.schema import Message
+
+
+@pytest.fixture
+def openai_llm(mocker):
+ mocker.patch("openai.Client")
+ return OpenAI(openai_api_key="test_api_key")
+
+
+class TestOpenAILLM:
+ def test_openai_llm_init_with_api_key(self, openai_llm):
+ assert openai_llm.client is not None, "Client should be initialized"
+ assert openai_llm.name == "gpt-3.5-turbo", "Default name not set correctly"
+
+ def test_openai_llm_init_success(self, mocker):
+ mocker.patch("os.getenv", return_value="fake-api-key")
+ llm = OpenAI()
+ assert llm.client is not None
+
+ def test_openai_llm_init_without_api_key(self, mocker):
+ mocker.patch("os.getenv", return_value=None)
+ with pytest.raises(ValueError) as _:
+ OpenAI()
+
+ def test_openai_llm_call_uninitialized_client(self, openai_llm):
+ # Set the client to None to simulate an uninitialized client
+ openai_llm.client = None
+ with pytest.raises(ValueError) as e:
+ llm_input = [Message(role="user", content="test")]
+ openai_llm(llm_input)
+ assert "OpenAI client is not initialized." in str(e.value)
+
+ def test_openai_llm_init_exception(self, mocker):
+ mocker.patch("os.getenv", return_value="fake-api-key")
+ mocker.patch("openai.OpenAI", side_effect=Exception("Initialization error"))
+ with pytest.raises(ValueError) as e:
+ OpenAI()
+ assert (
+ "OpenAI API client failed to initialize. Error: Initialization error"
+ in str(e.value)
+ )
+
+ def test_openai_llm_call_success(self, openai_llm, mocker):
+ mock_completion = mocker.MagicMock()
+ mock_completion.choices[0].message.content = "test"
+
+ mocker.patch("os.getenv", return_value="fake-api-key")
+ mocker.patch.object(
+ openai_llm.client.chat.completions, "create", return_value=mock_completion
+ )
+ llm_input = [Message(role="user", content="test")]
+ output = openai_llm(llm_input)
+ assert output == "test"
diff --git a/tests/unit/llms/test_llm_openrouter.py b/tests/unit/llms/test_llm_openrouter.py
new file mode 100644
index 00000000..3009e293
--- /dev/null
+++ b/tests/unit/llms/test_llm_openrouter.py
@@ -0,0 +1,59 @@
+import pytest
+from semantic_router.llms import OpenRouter
+from semantic_router.schema import Message
+
+
+@pytest.fixture
+def openrouter_llm(mocker):
+ mocker.patch("openai.Client")
+ return OpenRouter(openrouter_api_key="test_api_key")
+
+
+class TestOpenRouterLLM:
+ def test_openrouter_llm_init_with_api_key(self, openrouter_llm):
+ assert openrouter_llm.client is not None, "Client should be initialized"
+ assert (
+ openrouter_llm.name == "mistralai/mistral-7b-instruct"
+ ), "Default name not set correctly"
+
+ def test_openrouter_llm_init_success(self, mocker):
+ mocker.patch("os.getenv", return_value="fake-api-key")
+ llm = OpenRouter()
+ assert llm.client is not None
+
+ def test_openrouter_llm_init_without_api_key(self, mocker):
+ mocker.patch("os.getenv", return_value=None)
+ with pytest.raises(ValueError) as _:
+ OpenRouter()
+
+ def test_openrouter_llm_call_uninitialized_client(self, openrouter_llm):
+ # Set the client to None to simulate an uninitialized client
+ openrouter_llm.client = None
+ with pytest.raises(ValueError) as e:
+ llm_input = [Message(role="user", content="test")]
+ openrouter_llm(llm_input)
+ assert "OpenRouter client is not initialized." in str(e.value)
+
+ def test_openrouter_llm_init_exception(self, mocker):
+ mocker.patch("os.getenv", return_value="fake-api-key")
+ mocker.patch("openai.OpenAI", side_effect=Exception("Initialization error"))
+ with pytest.raises(ValueError) as e:
+ OpenRouter()
+ assert (
+ "OpenRouter API client failed to initialize. Error: Initialization error"
+ in str(e.value)
+ )
+
+ def test_openrouter_llm_call_success(self, openrouter_llm, mocker):
+ mock_completion = mocker.MagicMock()
+ mock_completion.choices[0].message.content = "test"
+
+ mocker.patch("os.getenv", return_value="fake-api-key")
+ mocker.patch.object(
+ openrouter_llm.client.chat.completions,
+ "create",
+ return_value=mock_completion,
+ )
+ llm_input = [Message(role="user", content="test")]
+ output = openrouter_llm(llm_input)
+ assert output == "test"
diff --git a/tests/unit/test_route.py b/tests/unit/test_route.py
index 2eb784d4..e7842d39 100644
--- a/tests/unit/test_route.py
+++ b/tests/unit/test_route.py
@@ -1,6 +1,6 @@
from unittest.mock import patch # , AsyncMock
-# import pytest
+import pytest
from semantic_router.llms import BaseLLM
from semantic_router.route import Route, is_valid
@@ -61,6 +61,21 @@ def __call__(self, prompt):
class TestRoute:
+ def test_value_error_in_route_call(self):
+ function_schema = {"name": "test_function", "type": "function"}
+
+ route = Route(
+ name="test_function",
+ utterances=["utterance1", "utterance2"],
+ function_schema=function_schema,
+ )
+
+ with pytest.raises(
+ ValueError,
+ match="LLM is required for dynamic routes. Please ensure the 'llm' is set.",
+ ):
+ route("test_query")
+
def test_generate_dynamic_route(self):
mock_llm = MockLLM(name="test")
function_schema = {"name": "test_function", "type": "function"}
diff --git a/tests/unit/test_schema.py b/tests/unit/test_schema.py
index 97b5028e..a9e794cb 100644
--- a/tests/unit/test_schema.py
+++ b/tests/unit/test_schema.py
@@ -1,9 +1,10 @@
import pytest
-
+from pydantic import ValidationError
from semantic_router.schema import (
CohereEncoder,
Encoder,
EncoderType,
+ Message,
OpenAIEncoder,
)
@@ -38,3 +39,27 @@ def test_encoder_call_method(self, mocker):
encoder = Encoder(type="openai", name="test-engine")
result = encoder(["test"])
assert result == [0.1, 0.2, 0.3]
+
+
+class TestMessageDataclass:
+ def test_message_creation(self):
+ message = Message(role="user", content="Hello!")
+ assert message.role == "user"
+ assert message.content == "Hello!"
+
+ with pytest.raises(ValidationError):
+ Message(user_role="invalid_role", message="Hello!")
+
+ def test_message_to_openai(self):
+ message = Message(role="user", content="Hello!")
+ openai_format = message.to_openai()
+ assert openai_format == {"role": "user", "content": "Hello!"}
+
+ message = Message(role="invalid_role", content="Hello!")
+ with pytest.raises(ValueError):
+ message.to_openai()
+
+ def test_message_to_cohere(self):
+ message = Message(role="user", content="Hello!")
+ cohere_format = message.to_cohere()
+ assert cohere_format == {"role": "user", "message": "Hello!"}
From 45b2079d5772434e9f453f41db291cff02b6e75e Mon Sep 17 00:00:00 2001
From: James Briggs <35938317+jamescalam@users.noreply.github.com>
Date: Sun, 7 Jan 2024 15:24:33 +0100
Subject: [PATCH 4/8] added LLM to llm classes, update version and docs
---
README.md | 22 ++++----
docs/02-dynamic-routes.ipynb | 80 ++++++++++++++++--------------
pyproject.toml | 2 +-
semantic_router/layer.py | 16 +++---
semantic_router/llms/__init__.py | 8 +--
semantic_router/llms/cohere.py | 2 +-
semantic_router/llms/openai.py | 2 +-
semantic_router/llms/openrouter.py | 2 +-
semantic_router/route.py | 7 ++-
9 files changed, 73 insertions(+), 68 deletions(-)
diff --git a/README.md b/README.md
index da3fe685..fe5db343 100644
--- a/README.md
+++ b/README.md
@@ -12,7 +12,7 @@
-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"] = ""
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)
diff --git a/docs/02-dynamic-routes.ipynb b/docs/02-dynamic-routes.ipynb
index d8078cb2..c695838e 100644
--- a/docs/02-dynamic-routes.ipynb
+++ b/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,16 +92,23 @@
"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"
]
}
],
@@ -119,13 +116,21 @@
"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': \"\"}"
]
},
- "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?\")"
diff --git a/pyproject.toml b/pyproject.toml
index d3561c64..b24ed4f3 100644
--- a/pyproject.toml
+++ b/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 ",
diff --git a/semantic_router/layer.py b/semantic_router/layer.py
index eed04e36..b3173728 100644
--- a/semantic_router/layer.py
+++ b/semantic_router/layer.py
@@ -10,7 +10,7 @@
OpenAIEncoder,
FastEmbedEncoder,
)
-from semantic_router.llms import BaseLLM
+from semantic_router.llms import BaseLLM, OpenAILLM
from semantic_router.linear import similarity_matrix, top_scores
from semantic_router.route import Route
from semantic_router.schema import Encoder, EncoderType, RouteChoice
@@ -193,9 +193,13 @@ def __call__(self, text: str) -> RouteChoice:
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:
- raise ValueError(
- "LLM is required for dynamic routes. Please ensure the 'llm' is set."
+ 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)
@@ -228,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
diff --git a/semantic_router/llms/__init__.py b/semantic_router/llms/__init__.py
index 446f7c42..c7d6962b 100644
--- a/semantic_router/llms/__init__.py
+++ b/semantic_router/llms/__init__.py
@@ -1,7 +1,7 @@
from semantic_router.llms.base import BaseLLM
-from semantic_router.llms.openai import OpenAI
-from semantic_router.llms.openrouter import OpenRouter
-from semantic_router.llms.cohere import Cohere
+from semantic_router.llms.openai import OpenAILLM
+from semantic_router.llms.openrouter import OpenRouterLLM
+from semantic_router.llms.cohere import CohereLLM
-__all__ = ["BaseLLM", "OpenAI", "OpenRouter", "Cohere"]
+__all__ = ["BaseLLM", "OpenAILLM", "OpenRouterLLM", "CohereLLM"]
diff --git a/semantic_router/llms/cohere.py b/semantic_router/llms/cohere.py
index 80512d5c..be99bbc4 100644
--- a/semantic_router/llms/cohere.py
+++ b/semantic_router/llms/cohere.py
@@ -4,7 +4,7 @@
from semantic_router.schema import Message
-class Cohere(BaseLLM):
+class CohereLLM(BaseLLM):
client: cohere.Client | None = None
def __init__(
diff --git a/semantic_router/llms/openai.py b/semantic_router/llms/openai.py
index 18b6e706..5ee56398 100644
--- a/semantic_router/llms/openai.py
+++ b/semantic_router/llms/openai.py
@@ -5,7 +5,7 @@
from semantic_router.schema import Message
-class OpenAI(BaseLLM):
+class OpenAILLM(BaseLLM):
client: openai.OpenAI | None
temperature: float | None
max_tokens: int | None
diff --git a/semantic_router/llms/openrouter.py b/semantic_router/llms/openrouter.py
index 3b7a9b49..5c3b317f 100644
--- a/semantic_router/llms/openrouter.py
+++ b/semantic_router/llms/openrouter.py
@@ -5,7 +5,7 @@
from semantic_router.schema import Message
-class OpenRouter(BaseLLM):
+class OpenRouterLLM(BaseLLM):
client: openai.OpenAI | None
base_url: str | None
temperature: float | None
diff --git a/semantic_router/route.py b/semantic_router/route.py
index 454cfe79..0d8269f0 100644
--- a/semantic_router/route.py
+++ b/semantic_router/route.py
@@ -5,12 +5,10 @@
from pydantic import BaseModel
from semantic_router.llms import BaseLLM
-from semantic_router.schema import RouteChoice
+from semantic_router.schema import Message, RouteChoice
from semantic_router.utils import function_call
from semantic_router.utils.logger import logger
-from semantic_router.schema import Message
-
def is_valid(route_config: str) -> bool:
try:
@@ -51,7 +49,8 @@ def __call__(self, query: str) -> RouteChoice:
if self.function_schema:
if not self.llm:
raise ValueError(
- "LLM is required for dynamic routes. Please ensure the 'llm' is set."
+ "LLM is required for dynamic routes. Please ensure the `llm` "
+ "attribute is set."
)
# if a function schema is provided we generate the inputs
extracted_inputs = function_call.extract_function_inputs(
From 0696edcfc1073dd25315f6ee4f4c20f37b48a4b9 Mon Sep 17 00:00:00 2001
From: James Briggs <35938317+jamescalam@users.noreply.github.com>
Date: Sun, 7 Jan 2024 15:32:04 +0100
Subject: [PATCH 5/8] lint and test fix
---
semantic_router/layer.py | 4 ++--
semantic_router/llms/__init__.py | 3 +--
semantic_router/llms/base.py | 1 +
semantic_router/llms/cohere.py | 2 ++
semantic_router/llms/openai.py | 4 +++-
semantic_router/llms/openrouter.py | 4 +++-
semantic_router/schema.py | 2 +-
semantic_router/utils/function_call.py | 3 +--
tests/unit/llms/test_llm_cohere.py | 10 +++++-----
tests/unit/llms/test_llm_openai.py | 11 ++++++-----
tests/unit/llms/test_llm_openrouter.py | 11 ++++++-----
tests/unit/test_route.py | 1 +
tests/unit/test_schema.py | 1 +
13 files changed, 33 insertions(+), 24 deletions(-)
diff --git a/semantic_router/layer.py b/semantic_router/layer.py
index b3173728..72f8a8f0 100644
--- a/semantic_router/layer.py
+++ b/semantic_router/layer.py
@@ -7,11 +7,11 @@
from semantic_router.encoders import (
BaseEncoder,
CohereEncoder,
- OpenAIEncoder,
FastEmbedEncoder,
+ OpenAIEncoder,
)
-from semantic_router.llms import BaseLLM, OpenAILLM
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
diff --git a/semantic_router/llms/__init__.py b/semantic_router/llms/__init__.py
index c7d6962b..e5aedc85 100644
--- a/semantic_router/llms/__init__.py
+++ b/semantic_router/llms/__init__.py
@@ -1,7 +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
-from semantic_router.llms.cohere import CohereLLM
-
__all__ = ["BaseLLM", "OpenAILLM", "OpenRouterLLM", "CohereLLM"]
diff --git a/semantic_router/llms/base.py b/semantic_router/llms/base.py
index dd8a0afa..51db1fd0 100644
--- a/semantic_router/llms/base.py
+++ b/semantic_router/llms/base.py
@@ -1,4 +1,5 @@
from pydantic import BaseModel
+
from semantic_router.schema import Message
diff --git a/semantic_router/llms/cohere.py b/semantic_router/llms/cohere.py
index be99bbc4..77581700 100644
--- a/semantic_router/llms/cohere.py
+++ b/semantic_router/llms/cohere.py
@@ -1,5 +1,7 @@
import os
+
import cohere
+
from semantic_router.llms import BaseLLM
from semantic_router.schema import Message
diff --git a/semantic_router/llms/openai.py b/semantic_router/llms/openai.py
index 5ee56398..43ddd642 100644
--- a/semantic_router/llms/openai.py
+++ b/semantic_router/llms/openai.py
@@ -1,8 +1,10 @@
import os
+
import openai
-from semantic_router.utils.logger import logger
+
from semantic_router.llms import BaseLLM
from semantic_router.schema import Message
+from semantic_router.utils.logger import logger
class OpenAILLM(BaseLLM):
diff --git a/semantic_router/llms/openrouter.py b/semantic_router/llms/openrouter.py
index 5c3b317f..587eeb12 100644
--- a/semantic_router/llms/openrouter.py
+++ b/semantic_router/llms/openrouter.py
@@ -1,8 +1,10 @@
import os
+
import openai
-from semantic_router.utils.logger import logger
+
from semantic_router.llms import BaseLLM
from semantic_router.schema import Message
+from semantic_router.utils.logger import logger
class OpenRouterLLM(BaseLLM):
diff --git a/semantic_router/schema.py b/semantic_router/schema.py
index f4e4e8b3..5e94c23b 100644
--- a/semantic_router/schema.py
+++ b/semantic_router/schema.py
@@ -6,8 +6,8 @@
from semantic_router.encoders import (
BaseEncoder,
CohereEncoder,
- OpenAIEncoder,
FastEmbedEncoder,
+ OpenAIEncoder,
)
from semantic_router.utils.splitters import semantic_splitter
diff --git a/semantic_router/utils/function_call.py b/semantic_router/utils/function_call.py
index 19afcc47..cedd9b6e 100644
--- a/semantic_router/utils/function_call.py
+++ b/semantic_router/utils/function_call.py
@@ -5,8 +5,7 @@
from pydantic import BaseModel
from semantic_router.llms import BaseLLM
-from semantic_router.schema import Message
-from semantic_router.schema import RouteChoice
+from semantic_router.schema import Message, RouteChoice
from semantic_router.utils.logger import logger
diff --git a/tests/unit/llms/test_llm_cohere.py b/tests/unit/llms/test_llm_cohere.py
index 32443f04..aaf8a7e5 100644
--- a/tests/unit/llms/test_llm_cohere.py
+++ b/tests/unit/llms/test_llm_cohere.py
@@ -1,13 +1,13 @@
import pytest
-from semantic_router.llms import Cohere
+from semantic_router.llms import CohereLLM
from semantic_router.schema import Message
@pytest.fixture
def cohere_llm(mocker):
mocker.patch("cohere.Client")
- return Cohere(cohere_api_key="test_api_key")
+ return CohereLLM(cohere_api_key="test_api_key")
class TestCohereLLM:
@@ -19,7 +19,7 @@ def test_initialization_without_api_key(self, mocker, monkeypatch):
monkeypatch.delenv("COHERE_API_KEY", raising=False)
mocker.patch("cohere.Client")
with pytest.raises(ValueError):
- Cohere()
+ CohereLLM()
def test_call_method(self, cohere_llm, mocker):
mock_llm = mocker.MagicMock()
@@ -36,11 +36,11 @@ def test_raises_value_error_if_cohere_client_fails_to_initialize(self, mocker):
"cohere.Client", side_effect=Exception("Failed to initialize client")
)
with pytest.raises(ValueError):
- Cohere(cohere_api_key="test_api_key")
+ CohereLLM(cohere_api_key="test_api_key")
def test_raises_value_error_if_cohere_client_is_not_initialized(self, mocker):
mocker.patch("cohere.Client", return_value=None)
- llm = Cohere(cohere_api_key="test_api_key")
+ llm = CohereLLM(cohere_api_key="test_api_key")
with pytest.raises(ValueError):
llm("test")
diff --git a/tests/unit/llms/test_llm_openai.py b/tests/unit/llms/test_llm_openai.py
index 4b2b2f54..2f1171db 100644
--- a/tests/unit/llms/test_llm_openai.py
+++ b/tests/unit/llms/test_llm_openai.py
@@ -1,12 +1,13 @@
import pytest
-from semantic_router.llms import OpenAI
+
+from semantic_router.llms import OpenAILLM
from semantic_router.schema import Message
@pytest.fixture
def openai_llm(mocker):
mocker.patch("openai.Client")
- return OpenAI(openai_api_key="test_api_key")
+ return OpenAILLM(openai_api_key="test_api_key")
class TestOpenAILLM:
@@ -16,13 +17,13 @@ def test_openai_llm_init_with_api_key(self, openai_llm):
def test_openai_llm_init_success(self, mocker):
mocker.patch("os.getenv", return_value="fake-api-key")
- llm = OpenAI()
+ llm = OpenAILLM()
assert llm.client is not None
def test_openai_llm_init_without_api_key(self, mocker):
mocker.patch("os.getenv", return_value=None)
with pytest.raises(ValueError) as _:
- OpenAI()
+ OpenAILLM()
def test_openai_llm_call_uninitialized_client(self, openai_llm):
# Set the client to None to simulate an uninitialized client
@@ -36,7 +37,7 @@ def test_openai_llm_init_exception(self, mocker):
mocker.patch("os.getenv", return_value="fake-api-key")
mocker.patch("openai.OpenAI", side_effect=Exception("Initialization error"))
with pytest.raises(ValueError) as e:
- OpenAI()
+ OpenAILLM()
assert (
"OpenAI API client failed to initialize. Error: Initialization error"
in str(e.value)
diff --git a/tests/unit/llms/test_llm_openrouter.py b/tests/unit/llms/test_llm_openrouter.py
index 3009e293..9b1ee150 100644
--- a/tests/unit/llms/test_llm_openrouter.py
+++ b/tests/unit/llms/test_llm_openrouter.py
@@ -1,12 +1,13 @@
import pytest
-from semantic_router.llms import OpenRouter
+
+from semantic_router.llms import OpenRouterLLM
from semantic_router.schema import Message
@pytest.fixture
def openrouter_llm(mocker):
mocker.patch("openai.Client")
- return OpenRouter(openrouter_api_key="test_api_key")
+ return OpenRouterLLM(openrouter_api_key="test_api_key")
class TestOpenRouterLLM:
@@ -18,13 +19,13 @@ def test_openrouter_llm_init_with_api_key(self, openrouter_llm):
def test_openrouter_llm_init_success(self, mocker):
mocker.patch("os.getenv", return_value="fake-api-key")
- llm = OpenRouter()
+ llm = OpenRouterLLM()
assert llm.client is not None
def test_openrouter_llm_init_without_api_key(self, mocker):
mocker.patch("os.getenv", return_value=None)
with pytest.raises(ValueError) as _:
- OpenRouter()
+ OpenRouterLLM()
def test_openrouter_llm_call_uninitialized_client(self, openrouter_llm):
# Set the client to None to simulate an uninitialized client
@@ -38,7 +39,7 @@ def test_openrouter_llm_init_exception(self, mocker):
mocker.patch("os.getenv", return_value="fake-api-key")
mocker.patch("openai.OpenAI", side_effect=Exception("Initialization error"))
with pytest.raises(ValueError) as e:
- OpenRouter()
+ OpenRouterLLM()
assert (
"OpenRouter API client failed to initialize. Error: Initialization error"
in str(e.value)
diff --git a/tests/unit/test_route.py b/tests/unit/test_route.py
index e7842d39..0f7c4d8d 100644
--- a/tests/unit/test_route.py
+++ b/tests/unit/test_route.py
@@ -1,6 +1,7 @@
from unittest.mock import patch # , AsyncMock
import pytest
+
from semantic_router.llms import BaseLLM
from semantic_router.route import Route, is_valid
diff --git a/tests/unit/test_schema.py b/tests/unit/test_schema.py
index a9e794cb..a41d5fa7 100644
--- a/tests/unit/test_schema.py
+++ b/tests/unit/test_schema.py
@@ -1,5 +1,6 @@
import pytest
from pydantic import ValidationError
+
from semantic_router.schema import (
CohereEncoder,
Encoder,
From 3c3dce4cb16d1ad7e85a52ef7f641a3672b8d75b Mon Sep 17 00:00:00 2001
From: James Briggs <35938317+jamescalam@users.noreply.github.com>
Date: Sun, 7 Jan 2024 15:37:25 +0100
Subject: [PATCH 6/8] update test for ValueError in missing LLM
---
tests/unit/test_route.py | 5 +----
1 file changed, 1 insertion(+), 4 deletions(-)
diff --git a/tests/unit/test_route.py b/tests/unit/test_route.py
index 0f7c4d8d..33a9ac13 100644
--- a/tests/unit/test_route.py
+++ b/tests/unit/test_route.py
@@ -71,10 +71,7 @@ def test_value_error_in_route_call(self):
function_schema=function_schema,
)
- with pytest.raises(
- ValueError,
- match="LLM is required for dynamic routes. Please ensure the 'llm' is set.",
- ):
+ with pytest.raises(ValueError):
route("test_query")
def test_generate_dynamic_route(self):
From e9660ee0388f0aa2798786c41dfb0c54d0c6f659 Mon Sep 17 00:00:00 2001
From: James Briggs <35938317+jamescalam@users.noreply.github.com>
Date: Sun, 7 Jan 2024 15:48:50 +0100
Subject: [PATCH 7/8] added dynamic route init test
---
tests/unit/test_layer.py | 12 ++++++++++++
1 file changed, 12 insertions(+)
diff --git a/tests/unit/test_layer.py b/tests/unit/test_layer.py
index 32754997..1511a983 100644
--- a/tests/unit/test_layer.py
+++ b/tests/unit/test_layer.py
@@ -91,6 +91,13 @@ def routes():
Route(name="Route 2", utterances=["Goodbye", "Bye", "Au revoir"]),
]
+@pytest.fixture
+def dynamic_routes():
+ return [
+ Route(name="Route 1", utterances=["Hello", "Hi"], function_schema="test"),
+ Route(name="Route 2", utterances=["Goodbye", "Bye", "Au revoir"]),
+ ]
+
class TestRouteLayer:
def test_initialization(self, openai_encoder, routes):
@@ -106,7 +113,12 @@ def test_initialization(self, openai_encoder, routes):
def test_initialization_different_encoders(self, cohere_encoder, openai_encoder):
route_layer_cohere = RouteLayer(encoder=cohere_encoder)
assert route_layer_cohere.score_threshold == 0.3
+ route_layer_openai = RouteLayer(encoder=openai_encoder)
+ assert route_layer_openai.score_threshold == 0.82
+ def test_initialization_dynamic_route(self, cohere_encoder, openai_encoder):
+ route_layer_cohere = RouteLayer(encoder=cohere_encoder)
+ assert route_layer_cohere.score_threshold == 0.3
route_layer_openai = RouteLayer(encoder=openai_encoder)
assert route_layer_openai.score_threshold == 0.82
From 774b8eda4500ea3b32cc0c49c09bcd8ce2e3b55d Mon Sep 17 00:00:00 2001
From: James Briggs <35938317+jamescalam@users.noreply.github.com>
Date: Sun, 7 Jan 2024 15:49:51 +0100
Subject: [PATCH 8/8] lint
---
tests/unit/test_layer.py | 1 +
1 file changed, 1 insertion(+)
diff --git a/tests/unit/test_layer.py b/tests/unit/test_layer.py
index 1511a983..495d1bdc 100644
--- a/tests/unit/test_layer.py
+++ b/tests/unit/test_layer.py
@@ -91,6 +91,7 @@ def routes():
Route(name="Route 2", utterances=["Goodbye", "Bye", "Au revoir"]),
]
+
@pytest.fixture
def dynamic_routes():
return [