From 05ae9df1f69e69301bb3311c45dac1bc6b946441 Mon Sep 17 00:00:00 2001
From: Luca Mannini <48441989+italianconcerto@users.noreply.github.com>
Date: Fri, 31 Jan 2025 14:27:50 +0100
Subject: [PATCH] fix: error in base.py and hybrid.py related to the unzipping
of routes fixed
refactor: linter + formatter
---
docs/source/conf.py | 2 +-
semantic_router/__init__.py | 2 +-
semantic_router/encoders/__init__.py | 3 +--
semantic_router/encoders/aurelio.py | 4 ++--
semantic_router/encoders/base.py | 2 +-
semantic_router/encoders/bedrock.py | 4 +++-
semantic_router/encoders/bm25.py | 4 ++--
semantic_router/encoders/clip.py | 4 ++--
semantic_router/encoders/huggingface.py | 6 ++---
semantic_router/encoders/openai.py | 7 +++---
semantic_router/encoders/zure.py | 4 ++--
semantic_router/index/base.py | 9 ++++---
semantic_router/index/hybrid_local.py | 5 ++--
semantic_router/index/local.py | 5 ++--
semantic_router/index/pinecone.py | 7 +++---
semantic_router/index/postgres.py | 2 +-
semantic_router/llms/base.py | 2 +-
semantic_router/llms/llamacpp.py | 2 +-
semantic_router/llms/openai.py | 22 ++++++++---------
semantic_router/llms/openrouter.py | 3 +--
semantic_router/llms/zure.py | 2 +-
semantic_router/routers/__init__.py | 2 +-
semantic_router/routers/base.py | 16 ++++++++-----
semantic_router/routers/hybrid.py | 24 +++++++++++--------
semantic_router/routers/semantic.py | 2 +-
semantic_router/schema.py | 8 ++++---
semantic_router/utils/function_call.py | 3 +--
.../encoders/test_openai_integration.py | 2 ++
tests/unit/encoders/test_azure.py | 3 ++-
tests/unit/encoders/test_bedrock.py | 6 +++--
tests/unit/encoders/test_bm25.py | 4 ++--
tests/unit/encoders/test_clip.py | 4 +++-
tests/unit/encoders/test_fastembed.py | 4 ++--
tests/unit/encoders/test_huggingface.py | 2 +-
tests/unit/encoders/test_mistral.py | 1 -
tests/unit/encoders/test_openai.py | 3 ++-
tests/unit/encoders/test_vit.py | 2 +-
tests/unit/llms/test_llm_base.py | 4 +++-
tests/unit/test_function_schema.py | 1 +
tests/unit/test_router.py | 17 ++++++-------
tests/unit/test_sync.py | 21 ++++++++--------
41 files changed, 123 insertions(+), 107 deletions(-)
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 763e91ce..78249d78 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
diff --git a/semantic_router/__init__.py b/semantic_router/__init__.py
index 2b446694..53c24327 100644
--- a/semantic_router/__init__.py
+++ b/semantic_router/__init__.py
@@ -1,5 +1,5 @@
-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"]
diff --git a/semantic_router/encoders/__init__.py b/semantic_router/encoders/__init__.py
index 07e468d8..5fdeb789 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 c50c6c4e..d1d40aa7 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 993093c1..3e8cba21 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 7adf60c7..f200bea8 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 f42bf9c2..6c180179 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 065ff115..1ecfe41c 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 f553a082..b659ebac 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 1ac5e56e..4865cc29 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 fd8594a9..f16213dc 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 0025a0bb..cad75daf 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 704df257..c473a56c 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 9cccfc4a..d66f6159 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 15889013..75e01902 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 6f4a9f2a..362741ef 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 02df2764..1ebc7aa5 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 8257593c..dda11bea 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 7dc565e8..3e991148 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 03762364..34dc147d 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 ba833d04..fae5149d 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 79987406..754ea17d 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 b5136f97..43f29d5b 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,
@@ -1345,9 +1345,13 @@ class BaseRouter(BaseModel):
remote_routes = 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 route in remote_routes:
+ routes.append(route.route)
+ utterances.append(route.utterance)
+ metadata.append(route.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 94ef33a3..8bdc5a80 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):
@@ -288,9 +288,13 @@ class HybridRouter(BaseRouter):
remote_routes = 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 route in remote_routes:
+ routes.append(route.route)
+ utterances.append(route.utterance)
+ metadata.append(route.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 5efa2e48..a468e0de 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 273043f5..2ba85b49 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 f5fd5a87..e9c0afd1 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 8c557d80..63c6968a 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 5a6841d7..01a495a4 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 619325ed..a1f40cc2 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 a22b0cc2..6c213348 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 29776151..ac87e6a0 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 f46303c1..a8ad3438 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 7f496e39..1677a450 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 a236ad16..25dba6b7 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 96171cce..a564f2e6 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 622b2268..c2488a1e 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 7fe8c0f0..c849901d 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 94a3f8e8..fbd2d7c6 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 a210cc4c..771d44ef 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
diff --git a/tests/unit/test_sync.py b/tests/unit/test_sync.py
index d40a57fb..02f4aa89 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
--
GitLab