diff --git a/semantic_router/layer.py b/semantic_router/layer.py index b3173728d0c66e08288e79e93a1fbb57389d7940..72f8a8f0652d3fef69c31693505f574fd7646358 100644 --- a/semantic_router/layer.py +++ b/semantic_router/layer.py @@ -7,11 +7,11 @@ import yaml 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 c7d6962b5a2e59635a51497d3284faec053232a3..e5aedc85fd30cc0b576fc2170c1b7ca694bdf200 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 dd8a0afada5f74e64277f260a8487483c5927d3c..51db1fd0e5f317ba7a436cecaac891e8628a161e 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 be99bbc4b1583d14d6b917d615776ccdac9158fe..775817001f22a4474f8705b2335c697730035d47 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 5ee5639822e7212dd2e16c7d8190c1fb0ccc0b1a..43ddd642bd42702461da56fd5de87dea01635dfb 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 5c3b317f84e65065281d52c3af3942f76ff2af7f..587eeb121e60c4e6f63dac289e390dd5c9086a1a 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 f4e4e8b3e6998757b82f20ba6d864a9c19fb705d..5e94c23b13f8c9d9359609e46c790e18dd860ab4 100644 --- a/semantic_router/schema.py +++ b/semantic_router/schema.py @@ -6,8 +6,8 @@ from pydantic.dataclasses import dataclass 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 19afcc4732fd75790fd76225e912c801a3dcd206..cedd9b6ecd86131b630cf6d4921848604dc88fa0 100644 --- a/semantic_router/utils/function_call.py +++ b/semantic_router/utils/function_call.py @@ -5,8 +5,7 @@ from typing import Any, Callable, Union 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 32443f0485a4148748504b93437b2c39fc34f0b4..aaf8a7e5de4ea23c683d8e9e2393937579bf61cc 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 @@ class TestCohereLLM: 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 @@ class TestCohereLLM: "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 4b2b2f545efdce8ee664a46d62cc971b0faab393..2f1171db84177d65b1bd2f438d531cadc8522edf 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 @@ class TestOpenAILLM: 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 @@ class TestOpenAILLM: 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 3009e29367f5da3394f7acc383ba58c059f69980..9b1ee150f2b301984c24eeb144453cd6a5ea0973 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 @@ class TestOpenRouterLLM: 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 @@ class TestOpenRouterLLM: 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 e7842d39a1c9cc8fb79dba6777e1e404492f5e0a..0f7c4d8de7b60da14e41ecd7757a8d77f9fc8018 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 a9e794cbe0225234163f26d170729f86e2bbeb50..a41d5fa732eb585dd9cd5081de6bebd3538cbfcb 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,