diff --git a/docs/source/conf.py b/docs/source/conf.py index 763e91ce5e75f1283ee7adfd56f688aa0c22e23d..2ac9efd80d907883faa3b322a2e04856bb984245 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -1,6 +1,6 @@ -from typing import List import os import sys +from typing import List sys.path.insert(0, os.path.abspath("../..")) # Source code dir relative to this file @@ -15,7 +15,7 @@ sys.path.insert(0, os.path.abspath("../..")) # Source code dir relative to this project = "Semantic Router" copyright = "2025, Aurelio AI" author = "Aurelio AI" -release = "0.1.0.dev7" +release = "0.1.0.dev8" # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration diff --git a/pyproject.toml b/pyproject.toml index ee754493648cfe75203416fc204d7ee71a4975be..ef967e30ecc43b396dcedf11fd4de5b3f1b2ea83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "semantic-router" -version = "0.1.0.dev7" +version = "0.1.0.dev8" description = "Super fast semantic router for AI decision making" authors = ["Aurelio AI <hello@aurelio.ai>"] readme = "README.md" diff --git a/semantic_router/__init__.py b/semantic_router/__init__.py index 2b446694b21367208814ee3824e4f35653cd96a4..41caa586bc8adb54b15b0dd15fe2148afcd266b1 100644 --- a/semantic_router/__init__.py +++ b/semantic_router/__init__.py @@ -1,6 +1,6 @@ -from semantic_router.routers import RouterConfig, SemanticRouter, HybridRouter from semantic_router.route import Route +from semantic_router.routers import HybridRouter, RouterConfig, SemanticRouter __all__ = ["SemanticRouter", "HybridRouter", "Route", "RouterConfig"] -__version__ = "0.1.0.dev7" +__version__ = "0.1.0.dev8" diff --git a/semantic_router/encoders/__init__.py b/semantic_router/encoders/__init__.py index 07e468d893afe8fb2b8ecf69d1dcd8bbc2f4d8ee..5fdeb789098aa009996df1c6b264f1db1bbeaf82 100644 --- a/semantic_router/encoders/__init__.py +++ b/semantic_router/encoders/__init__.py @@ -8,8 +8,7 @@ from semantic_router.encoders.clip import CLIPEncoder from semantic_router.encoders.cohere import CohereEncoder from semantic_router.encoders.fastembed import FastEmbedEncoder from semantic_router.encoders.google import GoogleEncoder -from semantic_router.encoders.huggingface import HuggingFaceEncoder -from semantic_router.encoders.huggingface import HFEndpointEncoder +from semantic_router.encoders.huggingface import HFEndpointEncoder, HuggingFaceEncoder from semantic_router.encoders.mistral import MistralEncoder from semantic_router.encoders.openai import OpenAIEncoder from semantic_router.encoders.tfidf import TfidfEncoder diff --git a/semantic_router/encoders/aurelio.py b/semantic_router/encoders/aurelio.py index c50c6c4e51d648221b2421cdf4dcfb3481ef6c35..d1d40aa757b89bd58bdb57b61fe23110a719a6e1 100644 --- a/semantic_router/encoders/aurelio.py +++ b/semantic_router/encoders/aurelio.py @@ -1,8 +1,8 @@ import os from typing import Any, List, Optional -from pydantic import Field -from aurelio_sdk import AurelioClient, AsyncAurelioClient, EmbeddingResponse +from aurelio_sdk import AsyncAurelioClient, AurelioClient, EmbeddingResponse +from pydantic import Field from semantic_router.encoders.base import SparseEncoder from semantic_router.schema import SparseEmbedding diff --git a/semantic_router/encoders/base.py b/semantic_router/encoders/base.py index 993093c1a2132c58e769b3035ed7588aefe7a689..3e8cba216511eceb535fcd28872cc0ca3d7ef739 100644 --- a/semantic_router/encoders/base.py +++ b/semantic_router/encoders/base.py @@ -1,7 +1,7 @@ from typing import Any, Coroutine, List, Optional -from pydantic import BaseModel, Field, field_validator import numpy as np +from pydantic import BaseModel, Field, field_validator from semantic_router.schema import SparseEmbedding diff --git a/semantic_router/encoders/bedrock.py b/semantic_router/encoders/bedrock.py index 7adf60c717252b71da0d9f604aec43d2dc5226b9..f200bea8548bcc60d735d0c0bcb4f6fab430f675 100644 --- a/semantic_router/encoders/bedrock.py +++ b/semantic_router/encoders/bedrock.py @@ -17,10 +17,12 @@ Classes: """ import json -from typing import Dict, List, Optional, Any, Union import os from time import sleep +from typing import Any, Dict, List, Optional, Union + import tiktoken + from semantic_router.encoders import DenseEncoder from semantic_router.utils.defaults import EncoderDefault from semantic_router.utils.logger import logger diff --git a/semantic_router/encoders/bm25.py b/semantic_router/encoders/bm25.py index f42bf9c28937419b29ae4429151a1fe3120c0487..6c180179f63640734d279b3e1c387362f879a04c 100644 --- a/semantic_router/encoders/bm25.py +++ b/semantic_router/encoders/bm25.py @@ -1,9 +1,9 @@ from typing import Any, Dict, List, Optional from semantic_router.encoders.tfidf import TfidfEncoder -from semantic_router.utils.logger import logger -from semantic_router.schema import SparseEmbedding from semantic_router.route import Route +from semantic_router.schema import SparseEmbedding +from semantic_router.utils.logger import logger class BM25Encoder(TfidfEncoder): diff --git a/semantic_router/encoders/clip.py b/semantic_router/encoders/clip.py index 065ff115bede8d4597b57a62857627ea8099525f..1ecfe41cc3102ab914ee05a7b8248553e92a64ce 100644 --- a/semantic_router/encoders/clip.py +++ b/semantic_router/encoders/clip.py @@ -1,8 +1,8 @@ -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional import numpy as np from pydantic import PrivateAttr -from typing import Dict + from semantic_router.encoders import DenseEncoder diff --git a/semantic_router/encoders/huggingface.py b/semantic_router/encoders/huggingface.py index f553a082f663bb6dc697c9ecd80eb581fb21be73..b659ebac842949f2cd2fca0227e9398c3c89de15 100644 --- a/semantic_router/encoders/huggingface.py +++ b/semantic_router/encoders/huggingface.py @@ -20,11 +20,11 @@ Classes: HFEndpointEncoder: A class for generating embeddings using a Huggingface endpoint. """ -import requests -import time import os -from typing import Any, List, Optional, Dict +import time +from typing import Any, Dict, List, Optional +import requests from pydantic import PrivateAttr from semantic_router.encoders import DenseEncoder diff --git a/semantic_router/encoders/openai.py b/semantic_router/encoders/openai.py index 1ac5e56e4e519ffb8f71ba18f7fdcb86d3168f23..4865cc298514601d798ba3d3c0cb842a08cac3a0 100644 --- a/semantic_router/encoders/openai.py +++ b/semantic_router/encoders/openai.py @@ -1,21 +1,20 @@ -from asyncio import sleep as asleep import os +from asyncio import sleep as asleep from time import sleep from typing import Any, List, Optional, Union -from pydantic import PrivateAttr import openai +import tiktoken from openai import OpenAIError from openai._types import NotGiven from openai.types import CreateEmbeddingResponse -import tiktoken +from pydantic import PrivateAttr from semantic_router.encoders import DenseEncoder from semantic_router.schema import EncoderInfo from semantic_router.utils.defaults import EncoderDefault from semantic_router.utils.logger import logger - model_configs = { "text-embedding-ada-002": EncoderInfo( name="text-embedding-ada-002", diff --git a/semantic_router/encoders/zure.py b/semantic_router/encoders/zure.py index fd8594a986ea5911cd78c4761b304f72fb6b4b2e..f16213dcb7d0c155dd380f467ec685bcdbb42b11 100644 --- a/semantic_router/encoders/zure.py +++ b/semantic_router/encoders/zure.py @@ -1,11 +1,11 @@ -from asyncio import sleep as asleep import os +from asyncio import sleep as asleep from time import sleep from typing import List, Optional, Union import openai -from openai._types import NotGiven from openai import OpenAIError +from openai._types import NotGiven from openai.types import CreateEmbeddingResponse from semantic_router.encoders import DenseEncoder diff --git a/semantic_router/index/base.py b/semantic_router/index/base.py index 0025a0bbc652b69bf2a13e3839703416569e6283..cad75daf5f0a010f384e7ee38575f25b7d1ceea1 100644 --- a/semantic_router/index/base.py +++ b/semantic_router/index/base.py @@ -1,17 +1,16 @@ import asyncio -from datetime import datetime -import time -from typing import Any, List, Optional, Tuple, Union, Dict import json +import time +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np from pydantic import BaseModel -from semantic_router.schema import ConfigParameter, SparseEmbedding, Utterance from semantic_router.route import Route +from semantic_router.schema import ConfigParameter, SparseEmbedding, Utterance from semantic_router.utils.logger import logger - RETRY_WAIT_TIME = 2.5 diff --git a/semantic_router/index/hybrid_local.py b/semantic_router/index/hybrid_local.py index 704df2575e114dfa3aabe0fb802860a3820bfdae..c473a56c35a3436279090879cec4cd508f230da1 100644 --- a/semantic_router/index/hybrid_local.py +++ b/semantic_router/index/hybrid_local.py @@ -1,12 +1,11 @@ -from typing import List, Optional, Tuple, Dict +from typing import Any, Dict, List, Optional, Tuple import numpy as np from numpy.linalg import norm -from semantic_router.schema import ConfigParameter, SparseEmbedding, Utterance from semantic_router.index.local import LocalIndex +from semantic_router.schema import ConfigParameter, SparseEmbedding, Utterance from semantic_router.utils.logger import logger -from typing import Any class HybridLocalIndex(LocalIndex): diff --git a/semantic_router/index/local.py b/semantic_router/index/local.py index 9cccfc4ade9a84e520f95eac42de90ffa87922c9..d66f6159e57f918c772f31fd9c8ae5f86311063b 100644 --- a/semantic_router/index/local.py +++ b/semantic_router/index/local.py @@ -1,12 +1,11 @@ -from typing import List, Optional, Tuple, Dict +from typing import Any, Dict, List, Optional, Tuple import numpy as np -from semantic_router.schema import ConfigParameter, SparseEmbedding, Utterance from semantic_router.index.base import BaseIndex, IndexConfig from semantic_router.linear import similarity_matrix, top_scores +from semantic_router.schema import ConfigParameter, SparseEmbedding, Utterance from semantic_router.utils.logger import logger -from typing import Any class LocalIndex(BaseIndex): diff --git a/semantic_router/index/pinecone.py b/semantic_router/index/pinecone.py index 15889013b300533e78223b0a2ae5d20d2c274965..75e0190291a0eaff347c635633d03b156a07fb6c 100644 --- a/semantic_router/index/pinecone.py +++ b/semantic_router/index/pinecone.py @@ -1,12 +1,11 @@ -import aiohttp import asyncio import hashlib +import json import os import time -import json - -from typing import Any, Dict, List, Optional, Union, Tuple +from typing import Any, Dict, List, Optional, Tuple, Union +import aiohttp import numpy as np from pydantic import BaseModel, Field diff --git a/semantic_router/index/postgres.py b/semantic_router/index/postgres.py index 6f4a9f2a1e56eaaeb890e5850c79acb0d61039ac..362741efc1bc2f48f0039f31b5f7b449bbf672c4 100644 --- a/semantic_router/index/postgres.py +++ b/semantic_router/index/postgres.py @@ -1,7 +1,7 @@ import os import uuid from enum import Enum -from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import numpy as np from pydantic import BaseModel, Field diff --git a/semantic_router/llms/base.py b/semantic_router/llms/base.py index 02df2764502d2d2bdd70135406c5e27c5a900fda..1ebc7aa50156e3c4830882637d49983a5b111f6d 100644 --- a/semantic_router/llms/base.py +++ b/semantic_router/llms/base.py @@ -1,5 +1,5 @@ import json -from typing import Any, List, Optional, Dict +from typing import Any, Dict, List, Optional from pydantic import BaseModel diff --git a/semantic_router/llms/llamacpp.py b/semantic_router/llms/llamacpp.py index 8257593c6f250a8bc1ded5f3af0dbc30a8b44141..dda11bea5532768daf21427abd027648b848caed 100644 --- a/semantic_router/llms/llamacpp.py +++ b/semantic_router/llms/llamacpp.py @@ -1,6 +1,6 @@ from contextlib import contextmanager from pathlib import Path -from typing import Any, Optional, List, Dict +from typing import Any, Dict, List, Optional from pydantic import PrivateAttr diff --git a/semantic_router/llms/openai.py b/semantic_router/llms/openai.py index 7dc565e8877a46d74e5fd86d4e3d05e04c52175c..3e991148db715c20eaef1b9729d9baeeb921e133 100644 --- a/semantic_router/llms/openai.py +++ b/semantic_router/llms/openai.py @@ -1,24 +1,24 @@ +import inspect +import json import os -from typing import List, Optional, Any, Callable, Dict, Union -from pydantic import PrivateAttr +import re +from typing import Any, Callable, Dict, List, Optional, Union import openai -from openai._types import NotGiven, NOT_GIVEN +from openai._types import NOT_GIVEN, NotGiven +from openai.types.chat.chat_completion_message_tool_call import ( + ChatCompletionMessageToolCall, +) +from pydantic import PrivateAttr from semantic_router.llms import BaseLLM from semantic_router.schema import Message from semantic_router.utils.defaults import EncoderDefault -from semantic_router.utils.logger import logger -import json from semantic_router.utils.function_call import ( - get_schema, convert_python_type_to_json_type, + get_schema, ) -import inspect -import re -from openai.types.chat.chat_completion_message_tool_call import ( - ChatCompletionMessageToolCall, -) +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 0376236486c06f32775730bb332f13c5f87b0ab8..34dc147dbc7da1ab8724f79e90844c0c1a9d3684 100644 --- a/semantic_router/llms/openrouter.py +++ b/semantic_router/llms/openrouter.py @@ -1,9 +1,8 @@ import os from typing import List, Optional -from pydantic import PrivateAttr - import openai +from pydantic import PrivateAttr from semantic_router.llms import BaseLLM from semantic_router.schema import Message diff --git a/semantic_router/llms/zure.py b/semantic_router/llms/zure.py index ba833d044785b775b0da123db554a79f55dd5be8..fae5149d0e1f8a566806b5eb0f475ef55b9ea3a5 100644 --- a/semantic_router/llms/zure.py +++ b/semantic_router/llms/zure.py @@ -1,8 +1,8 @@ import os from typing import List, Optional -from pydantic import PrivateAttr import openai +from pydantic import PrivateAttr from semantic_router.llms import BaseLLM from semantic_router.schema import Message diff --git a/semantic_router/routers/__init__.py b/semantic_router/routers/__init__.py index 79987406ed8c430993f714f433e6df892b53e299..754ea17d7c4751d52aa4ff6024e7856a67f31e9a 100644 --- a/semantic_router/routers/__init__.py +++ b/semantic_router/routers/__init__.py @@ -1,6 +1,6 @@ from semantic_router.routers.base import BaseRouter, RouterConfig -from semantic_router.routers.semantic import SemanticRouter from semantic_router.routers.hybrid import HybridRouter +from semantic_router.routers.semantic import SemanticRouter __all__ = [ "BaseRouter", diff --git a/semantic_router/routers/base.py b/semantic_router/routers/base.py index b5136f97750e55f25358707774cbd2abead5da30..a9ff6ebad0466876780c00c3054eabe1831e00b4 100644 --- a/semantic_router/routers/base.py +++ b/semantic_router/routers/base.py @@ -1,15 +1,15 @@ +import hashlib import importlib import json import os import random -import hashlib from typing import Any, Callable, Dict, List, Optional, Tuple, Union -from typing_extensions import deprecated -from pydantic import BaseModel, Field import numpy as np import yaml # type: ignore +from pydantic import BaseModel, Field from tqdm.auto import tqdm +from typing_extensions import deprecated from semantic_router.encoders import ( AutoEncoder, @@ -1343,11 +1343,15 @@ class BaseRouter(BaseModel): # Switch to a local index for fitting from semantic_router.index.local import LocalIndex - remote_routes = self.index.get_utterances(include_metadata=True) + remote_utterances = self.index.get_utterances(include_metadata=True) # TODO Enhance by retrieving directly the vectors instead of embedding all utterances again - routes, utterances, function_schemas, metadata = map( - list, zip(*remote_routes) - ) + routes = [] + utterances = [] + metadata = [] + for utterance in remote_utterances: + routes.append(utterance.route) + utterances.append(utterance.utterance) + metadata.append(utterance.metadata) embeddings = self.encoder(utterances) self.index = LocalIndex() self.index.add( diff --git a/semantic_router/routers/hybrid.py b/semantic_router/routers/hybrid.py index 94ef33a3476060eeb8f34dde32c6a130d839532c..b7857783cd3ddbdcb8da1374256b9fbd3da6b337 100644 --- a/semantic_router/routers/hybrid.py +++ b/semantic_router/routers/hybrid.py @@ -1,22 +1,22 @@ -from typing import Any, Dict, List, Optional, Union -from tqdm.auto import tqdm import asyncio -from pydantic import Field +from typing import Any, Dict, List, Optional, Union import numpy as np +from pydantic import Field +from tqdm.auto import tqdm from semantic_router.encoders import ( + BM25Encoder, DenseEncoder, SparseEncoder, - BM25Encoder, TfidfEncoder, ) -from semantic_router.route import Route from semantic_router.index import BaseIndex, HybridLocalIndex +from semantic_router.llms import BaseLLM +from semantic_router.route import Route +from semantic_router.routers.base import BaseRouter, threshold_random_search, xq_reshape from semantic_router.schema import RouteChoice, SparseEmbedding, Utterance from semantic_router.utils.logger import logger -from semantic_router.routers.base import BaseRouter, xq_reshape, threshold_random_search -from semantic_router.llms import BaseLLM class HybridRouter(BaseRouter): @@ -286,11 +286,15 @@ class HybridRouter(BaseRouter): # Switch to a local index for fitting from semantic_router.index.hybrid_local import HybridLocalIndex - remote_routes = self.index.get_utterances(include_metadata=True) + remote_utterances = self.index.get_utterances(include_metadata=True) # TODO Enhance by retrieving directly the vectors instead of embedding all utterances again - routes, utterances, function_schemas, metadata = map( - list, zip(*remote_routes) - ) + routes = [] + utterances = [] + metadata = [] + for utterance in remote_utterances: + routes.append(utterance.route) + utterances.append(utterance.utterance) + metadata.append(utterance.metadata) embeddings = self.encoder(utterances) sparse_embeddings = self.sparse_encoder(utterances) self.index = HybridLocalIndex() diff --git a/semantic_router/routers/semantic.py b/semantic_router/routers/semantic.py index 5efa2e48f0b5c2b622471a99a28848edbac0ee09..a468e0de7669ad7e6efdc44ad2fe4bf93608781f 100644 --- a/semantic_router/routers/semantic.py +++ b/semantic_router/routers/semantic.py @@ -5,9 +5,9 @@ import numpy as np from semantic_router.encoders import DenseEncoder from semantic_router.index.base import BaseIndex from semantic_router.llms import BaseLLM -from semantic_router.utils.logger import logger from semantic_router.route import Route from semantic_router.routers.base import BaseRouter +from semantic_router.utils.logger import logger class SemanticRouter(BaseRouter): diff --git a/semantic_router/schema.py b/semantic_router/schema.py index 273043f55ced8f5fdedb2b91e2ef9b0d35ee3d17..2ba85b498fbe8fe804ae6808ac6f2c7d694a71a9 100644 --- a/semantic_router/schema.py +++ b/semantic_router/schema.py @@ -1,12 +1,14 @@ +import json from datetime import datetime, timezone from difflib import Differ from enum import Enum -import json +from typing import Any, Dict, List, Optional, Tuple, Union + import numpy as np -from typing import List, Optional, Union, Any, Dict, Tuple +from aurelio_sdk.schema import BM25Embedding from pydantic import BaseModel, Field + from semantic_router.utils.logger import logger -from aurelio_sdk.schema import BM25Embedding class EncoderType(Enum): diff --git a/semantic_router/utils/function_call.py b/semantic_router/utils/function_call.py index f5fd5a87fd100482cd65f3502c3f8f67349e36c2..e9c0afd1c0b7cdf8faff6fc288a47d22e1d13465 100644 --- a/semantic_router/utils/function_call.py +++ b/semantic_router/utils/function_call.py @@ -1,12 +1,11 @@ import inspect from typing import Any, Callable, Dict, List, Optional, Union -from pydantic import BaseModel +from pydantic import BaseModel, Field from semantic_router.llms import BaseLLM from semantic_router.schema import Message, RouteChoice from semantic_router.utils.logger import logger -from pydantic import Field class Parameter(BaseModel): diff --git a/tests/integration/encoders/test_openai_integration.py b/tests/integration/encoders/test_openai_integration.py index 8c557d809d54413ae28a0676d5abe56026d29022..63c6968a101b15a3a187962bd0d0af713c915c89 100644 --- a/tests/integration/encoders/test_openai_integration.py +++ b/tests/integration/encoders/test_openai_integration.py @@ -1,6 +1,8 @@ import os + import pytest from openai import OpenAIError + from semantic_router.encoders.base import DenseEncoder from semantic_router.encoders.openai import OpenAIEncoder diff --git a/tests/unit/encoders/test_azure.py b/tests/unit/encoders/test_azure.py index 5a6841d7edf5038aae4a9a7124331a760279b626..01a495a44ec1f4e58f9ee0a0c45f211d62226d76 100644 --- a/tests/unit/encoders/test_azure.py +++ b/tests/unit/encoders/test_azure.py @@ -1,5 +1,6 @@ -import pytest from unittest.mock import AsyncMock, Mock, patch + +import pytest from openai import OpenAIError from openai.types import CreateEmbeddingResponse, Embedding from openai.types.create_embedding_response import Usage diff --git a/tests/unit/encoders/test_bedrock.py b/tests/unit/encoders/test_bedrock.py index 619325ed0ac43e7f54cf7eacd14a68a21a22ee8f..a1f40cc2121ba518ec2c07fe68997739dc7814de 100644 --- a/tests/unit/encoders/test_bedrock.py +++ b/tests/unit/encoders/test_bedrock.py @@ -1,7 +1,9 @@ -import os -import pytest import json +import os from io import BytesIO + +import pytest + from semantic_router.encoders import BedrockEncoder diff --git a/tests/unit/encoders/test_bm25.py b/tests/unit/encoders/test_bm25.py index a22b0cc2eb5c9e13c8a31038b76b18fe0b6e152c..6c2133482f8a54f543b199d5de58963439e61dca 100644 --- a/tests/unit/encoders/test_bm25.py +++ b/tests/unit/encoders/test_bm25.py @@ -1,9 +1,9 @@ -import pytest +import nltk import numpy as np +import pytest from semantic_router.encoders import BM25Encoder from semantic_router.route import Route -import nltk nltk.download("punkt_tab") diff --git a/tests/unit/encoders/test_clip.py b/tests/unit/encoders/test_clip.py index 297761519f29870723970da7f6ece10e98ed9ff8..ac87e6a09fa03b9c3c31e4a1c55c7576dc179168 100644 --- a/tests/unit/encoders/test_clip.py +++ b/tests/unit/encoders/test_clip.py @@ -1,12 +1,14 @@ import os + import numpy as np import pytest _ = pytest.importorskip("torch") +from unittest.mock import patch # noqa: E402 + import torch # noqa: E402 from PIL import Image # noqa: E402 -from unittest.mock import patch # noqa: E402 from semantic_router.encoders import CLIPEncoder # noqa: E402 diff --git a/tests/unit/encoders/test_fastembed.py b/tests/unit/encoders/test_fastembed.py index f46303c1bb069653c792fcf00c05ee9ced112923..a8ad3438e25eebd04edf62990a44bf1733b1abb2 100644 --- a/tests/unit/encoders/test_fastembed.py +++ b/tests/unit/encoders/test_fastembed.py @@ -1,7 +1,7 @@ -from semantic_router.encoders import FastEmbedEncoder - import pytest +from semantic_router.encoders import FastEmbedEncoder + _ = pytest.importorskip("fastembed") diff --git a/tests/unit/encoders/test_huggingface.py b/tests/unit/encoders/test_huggingface.py index 7f496e39f8117dbe4242070e59517acf751ca9a0..1677a450037285885034b73efbc9e26dd4d2b68f 100644 --- a/tests/unit/encoders/test_huggingface.py +++ b/tests/unit/encoders/test_huggingface.py @@ -1,6 +1,6 @@ +import os from unittest.mock import patch -import os import numpy as np import pytest diff --git a/tests/unit/encoders/test_mistral.py b/tests/unit/encoders/test_mistral.py index a236ad167139b5d5994ed8ad2d3610e4b49fc349..25dba6b759bc2a033c64ae0055e0b1b45d626f48 100644 --- a/tests/unit/encoders/test_mistral.py +++ b/tests/unit/encoders/test_mistral.py @@ -1,7 +1,6 @@ from unittest.mock import patch import pytest - from mistralai.exceptions import MistralException from mistralai.models.embeddings import EmbeddingObject, EmbeddingResponse, UsageInfo diff --git a/tests/unit/encoders/test_openai.py b/tests/unit/encoders/test_openai.py index 96171cce5b462f76dfbd88d7795bfc684ca100c8..a564f2e63a100469dbb78e26278cd9f4a9c4d750 100644 --- a/tests/unit/encoders/test_openai.py +++ b/tests/unit/encoders/test_openai.py @@ -1,5 +1,6 @@ -import pytest from unittest.mock import AsyncMock, Mock, patch + +import pytest from openai import OpenAIError from openai.types import CreateEmbeddingResponse, Embedding from openai.types.create_embedding_response import Usage diff --git a/tests/unit/encoders/test_vit.py b/tests/unit/encoders/test_vit.py index 622b22689f327054392a98cb336ec3f40adee6cd..c2488a1e43b1126a762f7ebcca0825f55a097ece 100644 --- a/tests/unit/encoders/test_vit.py +++ b/tests/unit/encoders/test_vit.py @@ -1,6 +1,6 @@ +import os from unittest.mock import patch -import os import numpy as np import pytest diff --git a/tests/unit/llms/test_llm_base.py b/tests/unit/llms/test_llm_base.py index 7fe8c0f06e3d153f17c589c086267690c703aa9a..c849901dc4346a3b665015cb9679c1f834096967 100644 --- a/tests/unit/llms/test_llm_base.py +++ b/tests/unit/llms/test_llm_base.py @@ -1,6 +1,8 @@ +from unittest.mock import patch + import pytest + from semantic_router.llms import BaseLLM -from unittest.mock import patch class TestBaseLLM: diff --git a/tests/unit/test_function_schema.py b/tests/unit/test_function_schema.py index 94a3f8e8b182dfd8064c0548ef3ba634440ed462..fbd2d7c6ca378683780bd08d86f94e1e8b119be9 100644 --- a/tests/unit/test_function_schema.py +++ b/tests/unit/test_function_schema.py @@ -1,4 +1,5 @@ import inspect + from semantic_router.utils.function_call import FunctionSchema diff --git a/tests/unit/test_router.py b/tests/unit/test_router.py index a210cc4cd32075dc8f43e522316d95d2a1cecfe1..5de1e98b28141449046a5f158d8a192c4516c163 100644 --- a/tests/unit/test_router.py +++ b/tests/unit/test_router.py @@ -1,22 +1,23 @@ import importlib -from functools import wraps import os import tempfile -from unittest.mock import mock_open, patch -from datetime import datetime -import pytest import time +from datetime import datetime +from functools import wraps +from platform import python_version from typing import Optional -from semantic_router.encoders import DenseEncoder, CohereEncoder, OpenAIEncoder +from unittest.mock import mock_open, patch + +import pytest + +from semantic_router.encoders import CohereEncoder, DenseEncoder, OpenAIEncoder from semantic_router.index.local import LocalIndex from semantic_router.index.pinecone import PineconeIndex from semantic_router.index.qdrant import QdrantIndex -from semantic_router.routers import RouterConfig, SemanticRouter, HybridRouter from semantic_router.llms import BaseLLM, OpenAILLM from semantic_router.route import Route +from semantic_router.routers import HybridRouter, RouterConfig, SemanticRouter from semantic_router.utils.logger import logger -from platform import python_version - PINECONE_SLEEP = 8 RETRY_COUNT = 10 @@ -1268,3 +1269,25 @@ class TestLayerFit: # unpack test data X, y = zip(*test_data) route_layer.fit(X=list(X), y=list(y), batch_size=int(len(X) / 5)) + + def test_fit_local(self, routes, test_data, index_cls, encoder_cls, router_cls): + # TODO: this is super slow for PineconeIndex, need to fix + encoder = encoder_cls() + index = init_index(index_cls, index_name=encoder.__class__.__name__) + route_layer = router_cls( + encoder=encoder, + routes=routes, + index=index, + auto_sync="local", + ) + + @retry(max_retries=RETRY_COUNT, delay=PINECONE_SLEEP) + def check_is_ready(): + assert route_layer.index.is_ready() + + check_is_ready() + # unpack test data + X, y = zip(*test_data) + route_layer.fit( + X=list(X), y=list(y), batch_size=int(len(X) / 5), local_execution=True + ) diff --git a/tests/unit/test_sync.py b/tests/unit/test_sync.py index d40a57fb0a1feb80d34aa2a8d44da21b44ba3db9..02f4aa891c9992b1c7e52a44d34ed6d0385ef5e7 100644 --- a/tests/unit/test_sync.py +++ b/tests/unit/test_sync.py @@ -1,25 +1,26 @@ import asyncio -from functools import wraps import importlib import os -from datetime import datetime -import pytest import time +from datetime import datetime +from functools import wraps +from platform import python_version from typing import Optional -from semantic_router.encoders import DenseEncoder, CohereEncoder, OpenAIEncoder + +import pytest + +from semantic_router.encoders import CohereEncoder, DenseEncoder, OpenAIEncoder from semantic_router.index import ( - PineconeIndex, HybridLocalIndex, LocalIndex, - QdrantIndex, + PineconeIndex, PostgresIndex, + QdrantIndex, ) -from semantic_router.schema import Utterance -from semantic_router.routers import SemanticRouter, HybridRouter from semantic_router.route import Route +from semantic_router.routers import HybridRouter, SemanticRouter +from semantic_router.schema import Utterance from semantic_router.utils.logger import logger -from platform import python_version - PINECONE_SLEEP = 6 RETRY_COUNT = 5