diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index e9105edd056771eb02c4eeb052bbcc9836c02065..a69425a5e2da521dad6048bc25bec45be6516fe9 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -14,9 +14,9 @@ jobs:
strategy:
matrix:
python-version:
- - "3.11"
+ - "3.13"
env:
- POETRY_VERSION: "1.8.3"
+ POETRY_VERSION: "1.8.4"
steps:
- uses: actions/checkout@v4
- name: Cache Poetry
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index 424c9036240a76e4c73edc07e5480fabc72059bb..ca5a94e9058cedcab1e00554fea84c296a53d9f8 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -6,7 +6,7 @@ on:
pull_request:
env:
- POETRY_VERSION: "1.5.1"
+ POETRY_VERSION: "1.8.4"
jobs:
build:
@@ -14,7 +14,7 @@ jobs:
strategy:
matrix:
python-version:
- - "3.11"
+ - "3.13"
steps:
- uses: actions/checkout@v3
- name: Cache Poetry
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 57b3eb0bd6a7e86231d6bcc81d402a7e02f5a021..25fcef711ee7b1234783c9ce9fb4e2b50b441ffe 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -13,7 +13,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v2
with:
- python-version: '3.11'
+ python-version: '3.13'
- name: Install Poetry
run: |
curl -sSL https://install.python-poetry.org | python - -y --version 1.5.1
@@ -30,10 +30,10 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v2
with:
- python-version: '3.11'
+ python-version: '3.13'
- name: Install Poetry
run: |
- curl -sSL https://install.python-poetry.org | python - -y --version 1.5.1
+ curl -sSL https://install.python-poetry.org | python - -y --version 1.8.4
- name: Publish to PyPI
run: |
poetry config repositories.remote https://upload.pypi.org/legacy/
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 99d8ed78afda14f15afa9e7e6598dfa6a55e6089..944cd476f17e6d6221ec623d3a0d21b8c3326ae7 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -4,7 +4,7 @@ on:
pull_request:
env:
- POETRY_VERSION: "1.5.1"
+ POETRY_VERSION: "1.8.4"
jobs:
build:
@@ -15,6 +15,7 @@ jobs:
- "3.10"
- "3.11"
- "3.12"
+ - "3.13"
steps:
- uses: actions/checkout@v4
- name: Cache Poetry
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index b24594f35ab4a09c525946915f027ed267ac78c5..1dfe6d342f478deceb8a02ab83e1cad0569b040d 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,5 +1,5 @@
default_language_version:
- python: python3.9
+ python: python3.13
repos:
- repo: meta
hooks:
diff --git a/.python-version b/.python-version
index 2c0733315e415bfb5e5b353f9996ecd964d395b2..24ee5b1be9961e38a503c8e764b7385dbb6ba124 100644
--- a/.python-version
+++ b/.python-version
@@ -1 +1 @@
-3.11
+3.13
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
index 6ad1571dae762366510051c85a744313e8e63184..34e88daf03d6e13aa88604372127a32e213f5b04 100644
--- a/.readthedocs.yaml
+++ b/.readthedocs.yaml
@@ -9,7 +9,7 @@ version: 2
build:
os: ubuntu-22.04
tools:
- python: "3.12"
+ python: "3.13"
# You can also specify other tool versions:
# nodejs: "19"
# rust: "1.64"
diff --git a/LICENSE b/LICENSE
index 3fa741bbaa7903b52deccb2358ad3b0ac6b24759..06edeb37df222cc1accb948a0353addcc8cfdec5 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,6 +1,6 @@
MIT License
-Copyright (c) 2023 Aurelio AI
+Copyright (c) 2024 Aurelio AI
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
diff --git a/Makefile b/Makefile
index 979a97e522253154794c78688035ab5128dd6dba..30e1f6ad0d44a349017eb3b39aa569226bb03f3c 100644
--- a/Makefile
+++ b/Makefile
@@ -1,5 +1,5 @@
format:
- poetry run black --target-version py39 -l 88 .
+ poetry run black --target-version py313 -l 88 .
poetry run ruff --select I --fix .
PYTHON_FILES=.
@@ -7,7 +7,7 @@ lint: PYTHON_FILES=.
lint_diff: PYTHON_FILES=$(shell git diff --name-only --diff-filter=d main | grep -E '\.py$$')
lint lint_diff:
- poetry run black --target-version py311 -l 88 $(PYTHON_FILES) --check
+ poetry run black --target-version py313 -l 88 $(PYTHON_FILES) --check
poetry run ruff check .
poetry run mypy $(PYTHON_FILES)
diff --git a/poetry.lock b/poetry.lock
index 19288adbc6b52ed524059c8bcceaf8cbb9156d9a..1202d622c07a86c3e21d2d6b30bed1500ce7517c 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -13,13 +13,13 @@ files = [
[[package]]
name = "aiohappyeyeballs"
-version = "2.4.3"
+version = "2.4.4"
description = "Happy Eyeballs for asyncio"
optional = false
python-versions = ">=3.8"
files = [
- {file = "aiohappyeyeballs-2.4.3-py3-none-any.whl", hash = "sha256:8a7a83727b2756f394ab2895ea0765a0a8c475e3c71e98d43d76f22b4b435572"},
- {file = "aiohappyeyeballs-2.4.3.tar.gz", hash = "sha256:75cf88a15106a5002a8eb1dab212525c00d1f4c0fa96e551c9fbe6f09a621586"},
+ {file = "aiohappyeyeballs-2.4.4-py3-none-any.whl", hash = "sha256:a980909d50efcd44795c4afeca523296716d50cd756ddca6af8c65b996e27de8"},
+ {file = "aiohappyeyeballs-2.4.4.tar.gz", hash = "sha256:5fdd7d87889c63183afc18ce9271f9b0a7d32c2303e394468dd45d514a757745"},
]
[[package]]
@@ -191,21 +191,18 @@ files = [
[[package]]
name = "asttokens"
-version = "2.4.1"
+version = "3.0.0"
description = "Annotate AST trees with source code positions"
optional = false
-python-versions = "*"
+python-versions = ">=3.8"
files = [
- {file = "asttokens-2.4.1-py2.py3-none-any.whl", hash = "sha256:051ed49c3dcae8913ea7cd08e46a606dba30b79993209636c4875bc1d637bc24"},
- {file = "asttokens-2.4.1.tar.gz", hash = "sha256:b03869718ba9a6eb027e134bfdf69f38a236d681c83c160d510768af11254ba0"},
+ {file = "asttokens-3.0.0-py3-none-any.whl", hash = "sha256:e3078351a059199dd5138cb1c706e6430c05eff2ff136af5eb4790f9d28932e2"},
+ {file = "asttokens-3.0.0.tar.gz", hash = "sha256:0dcd8baa8d62b0c1d118b399b2ddba3c4aff271d0d7a9e0d4c1681c79035bbc7"},
]
-[package.dependencies]
-six = ">=1.12.0"
-
[package.extras]
-astroid = ["astroid (>=1,<2)", "astroid (>=2,<4)"]
-test = ["astroid (>=1,<2)", "astroid (>=2,<4)", "pytest"]
+astroid = ["astroid (>=2,<4)"]
+test = ["astroid (>=2,<4)", "pytest", "pytest-cov", "pytest-xdist"]
[[package]]
name = "async-timeout"
@@ -612,13 +609,13 @@ colorama = {version = "*", markers = "platform_system == \"Windows\""}
[[package]]
name = "cohere"
-version = "5.11.4"
+version = "5.12.0"
description = ""
optional = true
python-versions = "<4.0,>=3.8"
files = [
- {file = "cohere-5.11.4-py3-none-any.whl", hash = "sha256:59fb427e5426e0ee1c25b9deec83f0418a1c082240c57007f41384b34cd41552"},
- {file = "cohere-5.11.4.tar.gz", hash = "sha256:5586335a20de3bf6816f34151f9d9f2928880cdf776c57aae793b5cca58d1826"},
+ {file = "cohere-5.12.0-py3-none-any.whl", hash = "sha256:47f61c6db274f61fb06781da3808d717b4ac4d46b1ee487c2f727450038c14cb"},
+ {file = "cohere-5.12.0.tar.gz", hash = "sha256:52a30edd4f7253b551045eb624df6c14e840c350306c8a69ae322e1f59743969"},
]
[package.dependencies]
@@ -1170,7 +1167,10 @@ grpcio-status = [
{version = ">=1.33.2,<2.0.dev0", optional = true, markers = "python_version < \"3.11\" and extra == \"grpc\""},
{version = ">=1.49.1,<2.0.dev0", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""},
]
-proto-plus = ">=1.22.3,<2.0.0dev"
+proto-plus = [
+ {version = ">=1.22.3,<2.0.0dev", markers = "python_version < \"3.13\""},
+ {version = ">=1.25.0,<2.0.0dev", markers = "python_version >= \"3.13\""},
+]
protobuf = ">=3.19.5,<3.20.0 || >3.20.0,<3.20.1 || >3.20.1,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0.dev0"
requests = ">=2.18.0,<3.0.0.dev0"
@@ -1314,7 +1314,10 @@ files = [
google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0dev", extras = ["grpc"]}
google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0dev"
grpc-google-iam-v1 = ">=0.12.4,<1.0.0dev"
-proto-plus = ">=1.22.3,<2.0.0dev"
+proto-plus = [
+ {version = ">=1.22.3,<2.0.0dev", markers = "python_version < \"3.13\""},
+ {version = ">=1.25.0,<2.0.0dev", markers = "python_version >= \"3.13\""},
+]
protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0dev"
[[package]]
@@ -2790,13 +2793,13 @@ sympy = "*"
[[package]]
name = "openai"
-version = "1.55.2"
+version = "1.55.3"
description = "The official Python library for the openai API"
optional = false
python-versions = ">=3.8"
files = [
- {file = "openai-1.55.2-py3-none-any.whl", hash = "sha256:3027c7fa4a33ed759f4a3d076093fcfa1c55658660c889bec33f651e2dc77922"},
- {file = "openai-1.55.2.tar.gz", hash = "sha256:5cc0b1162b65dcdf670b4b41448f18dd470d2724ca04821ab1e86b6b4e88650b"},
+ {file = "openai-1.55.3-py3-none-any.whl", hash = "sha256:2a235d0e1e312cd982f561b18c27692e253852f4e5fb6ccf08cb13540a9bdaa1"},
+ {file = "openai-1.55.3.tar.gz", hash = "sha256:547e85b94535469f137a779d8770c8c5adebd507c2cc6340ca401a7c4d5d16f0"},
]
[package.dependencies]
@@ -5558,5 +5561,5 @@ vision = ["pillow", "torch", "torchvision", "transformers"]
[metadata]
lock-version = "2.0"
-python-versions = ">=3.9,<3.13"
-content-hash = "d649e614d9c7122c642ed351abf8886cc008960f0fd7664483cfb5dde7865808"
+python-versions = ">=3.9,<3.14"
+content-hash = "8637391565e137e9f5d9d52bd927845242b7dea9217901c1b69f7be6b0b31917"
diff --git a/pyproject.toml b/pyproject.toml
index e1eff0c7657e89c66fc3dfe35fab7b9a88406f68..a914965b233696d97f98eb7bc836b3234455a466 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,8 +8,8 @@ packages = [{include = "semantic_router"}]
license = "MIT"
[tool.poetry.dependencies]
-python = ">=3.9,<3.13"
-pydantic = "^2.5.3"
+python = ">=3.9,<3.14"
+pydantic = "^2.10.2"
openai = ">=1.10.0,<2.0.0"
cohere = {version = ">=5.9.4,<6.00", optional = true}
mistralai= {version = ">=0.0.12,<0.1.0", optional = true}
@@ -18,15 +18,15 @@ colorlog = "^6.8.0"
pyyaml = "^6.0.1"
aurelio-sdk = {version = "^0.0.16"}
pinecone-text = {version = ">=0.7.1,<0.10.0", optional = true}
-torch = {version = ">=2.1.0,<2.6.0", optional = true}
-transformers = {version = ">=4.36.2", optional = true}
-tokenizers = {version = ">=0.19", optional = true}
-llama-cpp-python = {version = ">=0.2.28,<0.2.86", optional = true}
+torch = {version = ">=2.1.0,<2.6.0", optional = true, python = "<3.13" }
+transformers = {version = ">=4.36.2", optional = true, python = "<3.13" }
+tokenizers = {version = ">=0.19", optional = true, python = "<3.13" }
+llama-cpp-python = {version = ">=0.2.28,<0.2.86", optional = true, python = "<3.13" }
colorama = "^0.4.6"
pinecone = {version=">=5.0.0", optional = true}
regex = ">=2023.12.25"
-torchvision = { version = ">=0.17.0,<0.18.0", optional = true}
-pillow = { version = ">=10.2.0,<11.0.0", optional = true}
+torchvision = { version = ">=0.17.0,<0.18.0", optional = true, python = "<3.13" }
+pillow = { version = ">=10.2.0,<11.0.0", optional = true, python = "<3.13" }
tiktoken = ">=0.6.0,<1.0.0"
qdrant-client = {version = "^1.11.1", optional = true}
google-cloud-aiplatform = {version = "^1.45.0", optional = true}
@@ -34,7 +34,7 @@ requests-mock = "^1.12.1"
boto3 = { version = "^1.34.98", optional = true }
botocore = {version = "^1.34.110", optional = true}
aiohttp = "^3.10.11"
-fastembed = {version = "^0.3.0", optional = true}
+fastembed = {version = "^0.3.0", optional = true, python = "<3.13" }
psycopg2-binary = {version = "^2.9.9", optional = true}
sphinx = {version = "^7.0.0", optional = true}
sphinxawesome-theme = {version = "^5.2.0", optional = true}
diff --git a/semantic_router/encoders/aurelio.py b/semantic_router/encoders/aurelio.py
index 8b2501ba4faa0a22a5a18c5500a300e77fa63f95..c50c6c4e51d648221b2421cdf4dcfb3481ef6c35 100644
--- a/semantic_router/encoders/aurelio.py
+++ b/semantic_router/encoders/aurelio.py
@@ -1,6 +1,6 @@
import os
from typing import Any, List, Optional
-from pydantic.v1 import Field
+from pydantic import Field
from aurelio_sdk import AurelioClient, AsyncAurelioClient, EmbeddingResponse
diff --git a/semantic_router/encoders/base.py b/semantic_router/encoders/base.py
index 0a25f21050e2e5da1ea610fb5194324cdeef9a5f..993093c1a2132c58e769b3035ed7588aefe7a689 100644
--- a/semantic_router/encoders/base.py
+++ b/semantic_router/encoders/base.py
@@ -1,6 +1,6 @@
from typing import Any, Coroutine, List, Optional
-from pydantic.v1 import BaseModel, Field, validator
+from pydantic import BaseModel, Field, field_validator
import numpy as np
from semantic_router.schema import SparseEmbedding
@@ -14,7 +14,7 @@ class DenseEncoder(BaseModel):
class Config:
arbitrary_types_allowed = True
- @validator("score_threshold", pre=True, always=True)
+ @field_validator("score_threshold")
def set_score_threshold(cls, v):
return float(v) if v is not None else None
diff --git a/semantic_router/encoders/clip.py b/semantic_router/encoders/clip.py
index 65fbdb8f0312d5339d68ae5900542aea2c896693..065ff115bede8d4597b57a62857627ea8099525f 100644
--- a/semantic_router/encoders/clip.py
+++ b/semantic_router/encoders/clip.py
@@ -1,7 +1,7 @@
from typing import Any, List, Optional
import numpy as np
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
from typing import Dict
from semantic_router.encoders import DenseEncoder
@@ -9,7 +9,6 @@ from semantic_router.encoders import DenseEncoder
class CLIPEncoder(DenseEncoder):
name: str = "openai/clip-vit-base-patch16"
type: str = "huggingface"
- score_threshold: float = 0.2
tokenizer_kwargs: Dict = {}
processor_kwargs: Dict = {}
model_kwargs: Dict = {}
@@ -21,6 +20,8 @@ class CLIPEncoder(DenseEncoder):
_Image: Any = PrivateAttr()
def __init__(self, **data):
+ if data.get("score_threshold") is None:
+ data["score_threshold"] = 0.2
super().__init__(**data)
self._tokenizer, self._processor, self._model = self._initialize_hf_model()
diff --git a/semantic_router/encoders/cohere.py b/semantic_router/encoders/cohere.py
index 04b878141a4b33d742d5bb28a5bfc8bcd1347d19..e919bae110769f87f3c76ff58104ef14a30b34a7 100644
--- a/semantic_router/encoders/cohere.py
+++ b/semantic_router/encoders/cohere.py
@@ -1,7 +1,7 @@
import os
from typing import Any, List, Optional
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
from semantic_router.encoders import DenseEncoder
from semantic_router.utils.defaults import EncoderDefault
diff --git a/semantic_router/encoders/fastembed.py b/semantic_router/encoders/fastembed.py
index 5cda5e643746c660cd6e820cb767a6a333981b2d..2c9977955994f100366d67ef75a419f83fd55bce 100644
--- a/semantic_router/encoders/fastembed.py
+++ b/semantic_router/encoders/fastembed.py
@@ -1,7 +1,7 @@
from typing import Any, List, Optional
import numpy as np
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
from semantic_router.encoders import DenseEncoder
diff --git a/semantic_router/encoders/huggingface.py b/semantic_router/encoders/huggingface.py
index 7c7e56f77f11cf5143437c822ffda338c8f92cf0..f553a082f663bb6dc697c9ecd80eb581fb21be73 100644
--- a/semantic_router/encoders/huggingface.py
+++ b/semantic_router/encoders/huggingface.py
@@ -25,7 +25,7 @@ import time
import os
from typing import Any, List, Optional, Dict
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
from semantic_router.encoders import DenseEncoder
from semantic_router.utils.logger import logger
@@ -34,7 +34,6 @@ from semantic_router.utils.logger import logger
class HuggingFaceEncoder(DenseEncoder):
name: str = "sentence-transformers/all-MiniLM-L6-v2"
type: str = "huggingface"
- score_threshold: float = 0.5
tokenizer_kwargs: Dict = {}
model_kwargs: Dict = {}
device: Optional[str] = None
@@ -43,6 +42,8 @@ class HuggingFaceEncoder(DenseEncoder):
_torch: Any = PrivateAttr()
def __init__(self, **data):
+ if data.get("score_threshold") is None:
+ data["score_threshold"] = 0.5
super().__init__(**data)
self._tokenizer, self._model = self._initialize_hf_model()
@@ -153,7 +154,6 @@ class HFEndpointEncoder(DenseEncoder):
name: str = "hugging_face_custom_endpoint"
huggingface_url: Optional[str] = None
huggingface_api_key: Optional[str] = None
- score_threshold: float = 0.8
def __init__(
self,
@@ -180,6 +180,8 @@ class HFEndpointEncoder(DenseEncoder):
"""
huggingface_url = huggingface_url or os.getenv("HF_API_URL")
huggingface_api_key = huggingface_api_key or os.getenv("HF_API_KEY")
+ if score_threshold is None:
+ score_threshold = 0.8
super().__init__(name=name, score_threshold=score_threshold) # type: ignore
diff --git a/semantic_router/encoders/mistral.py b/semantic_router/encoders/mistral.py
index 46bc89aec0ed9ed41fffc5734fa5a7658f6f5cd7..6c3a2f5ed64f2a2c581b1aa8241fb81930c07c42 100644
--- a/semantic_router/encoders/mistral.py
+++ b/semantic_router/encoders/mistral.py
@@ -4,7 +4,7 @@ import os
from time import sleep
from typing import Any, List, Optional
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
from semantic_router.encoders import DenseEncoder
from semantic_router.utils.defaults import EncoderDefault
diff --git a/semantic_router/encoders/openai.py b/semantic_router/encoders/openai.py
index fb8a83f097e252a7fcd03a356f648077eb267cda..1ac5e56e4e519ffb8f71ba18f7fdcb86d3168f23 100644
--- a/semantic_router/encoders/openai.py
+++ b/semantic_router/encoders/openai.py
@@ -2,7 +2,7 @@ from asyncio import sleep as asleep
import os
from time import sleep
from typing import Any, List, Optional, Union
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
import openai
from openai import OpenAIError
@@ -36,8 +36,8 @@ model_configs = {
class OpenAIEncoder(DenseEncoder):
- client: Optional[openai.Client]
- async_client: Optional[openai.AsyncClient]
+ _client: Optional[openai.Client] = PrivateAttr(default=None)
+ _async_client: Optional[openai.AsyncClient] = PrivateAttr(default=None)
dimensions: Union[int, NotGiven] = NotGiven()
token_limit: int = 8192 # default value, should be replaced by config
_token_encoder: Any = PrivateAttr()
@@ -77,10 +77,10 @@ class OpenAIEncoder(DenseEncoder):
if max_retries is not None:
self.max_retries = max_retries
try:
- self.client = openai.Client(
+ self._client = openai.Client(
base_url=base_url, api_key=api_key, organization=openai_org_id
)
- self.async_client = openai.AsyncClient(
+ self._async_client = openai.AsyncClient(
base_url=base_url, api_key=api_key, organization=openai_org_id
)
except Exception as e:
@@ -103,7 +103,7 @@ class OpenAIEncoder(DenseEncoder):
False and a document exceeds the token limit, an error will be
raised.
:return: List of embeddings for each document."""
- if self.client is None:
+ if self._client is None:
raise ValueError("OpenAI client is not initialized.")
embeds = None
@@ -114,7 +114,7 @@ class OpenAIEncoder(DenseEncoder):
# Exponential backoff
for j in range(self.max_retries + 1):
try:
- embeds = self.client.embeddings.create(
+ embeds = self._client.embeddings.create(
input=docs,
model=self.name,
dimensions=self.dimensions,
@@ -160,7 +160,7 @@ class OpenAIEncoder(DenseEncoder):
return text
async def acall(self, docs: List[str], truncate: bool = True) -> List[List[float]]:
- if self.async_client is None:
+ if self._async_client is None:
raise ValueError("OpenAI async client is not initialized.")
embeds = None
@@ -171,7 +171,7 @@ class OpenAIEncoder(DenseEncoder):
# Exponential backoff
for j in range(self.max_retries + 1):
try:
- embeds = await self.async_client.embeddings.create(
+ embeds = await self._async_client.embeddings.create(
input=docs,
model=self.name,
dimensions=self.dimensions,
diff --git a/semantic_router/encoders/vit.py b/semantic_router/encoders/vit.py
index 73cb058281d3c47d452bfa3d5a8e4c92246ae173..dec768e43d3eff3a8eb6a428ea4e5773a488c9ad 100644
--- a/semantic_router/encoders/vit.py
+++ b/semantic_router/encoders/vit.py
@@ -1,6 +1,6 @@
-from typing import Any, List, Optional, Dict
+from typing import Any, Dict, List, Optional
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
from semantic_router.encoders import DenseEncoder
@@ -8,7 +8,6 @@ from semantic_router.encoders import DenseEncoder
class VitEncoder(DenseEncoder):
name: str = "google/vit-base-patch16-224"
type: str = "huggingface"
- score_threshold: float = 0.5
processor_kwargs: Dict = {}
model_kwargs: Dict = {}
device: Optional[str] = None
@@ -19,6 +18,8 @@ class VitEncoder(DenseEncoder):
_Image: Any = PrivateAttr()
def __init__(self, **data):
+ if data.get("score_threshold") is None:
+ data["score_threshold"] = 0.5
super().__init__(**data)
self._processor, self._model = self._initialize_hf_model()
diff --git a/semantic_router/index/base.py b/semantic_router/index/base.py
index 97fe3bd4eadc4066f172ed159d0adef7f88f4e5b..884106c0fbe5b6fa79f201b761d2b18368282be5 100644
--- a/semantic_router/index/base.py
+++ b/semantic_router/index/base.py
@@ -2,7 +2,7 @@ from typing import Any, List, Optional, Tuple, Union, Dict
import json
import numpy as np
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
from semantic_router.schema import ConfigParameter, SparseEmbedding, Utterance
from semantic_router.route import Route
diff --git a/semantic_router/index/pinecone.py b/semantic_router/index/pinecone.py
index 303b47157533fbe71a74231ce550a2485d1f8e30..b4ba144e89fefe739f5a1d0dbc1ac84be58460f6 100644
--- a/semantic_router/index/pinecone.py
+++ b/semantic_router/index/pinecone.py
@@ -8,7 +8,7 @@ import json
from typing import Any, Dict, List, Optional, Union, Tuple
import numpy as np
-from pydantic.v1 import BaseModel, Field
+from pydantic import BaseModel, Field
from semantic_router.index.base import BaseIndex
from semantic_router.schema import ConfigParameter, SparseEmbedding
diff --git a/semantic_router/index/postgres.py b/semantic_router/index/postgres.py
index 6d445f2bedabeba63a24712db955d6d4757647a2..71ea32e85fb84531f85178fd74663ad8bce61bcb 100644
--- a/semantic_router/index/postgres.py
+++ b/semantic_router/index/postgres.py
@@ -103,7 +103,6 @@ class PostgresIndex(BaseIndex):
connection_string: Optional[str] = None
index_prefix: str = "semantic_router_"
index_name: str = "index"
- dimensions: int = 1536
metric: Metric = Metric.COSINE
namespace: Optional[str] = ""
conn: Optional["psycopg2.extensions.connection"] = None
@@ -115,9 +114,9 @@ class PostgresIndex(BaseIndex):
connection_string: Optional[str] = None,
index_prefix: str = "semantic_router_",
index_name: str = "index",
- dimensions: int = 1536,
metric: Metric = Metric.COSINE,
namespace: Optional[str] = "",
+ dimensions: int | None = None,
):
"""
Initializes the Postgres index with the specified parameters.
@@ -135,6 +134,8 @@ class PostgresIndex(BaseIndex):
:param namespace: An optional namespace for the index.
:type namespace: Optional[str]
"""
+ if dimensions is None:
+ dimensions = 1536
super().__init__()
# try and import psycopg2
try:
diff --git a/semantic_router/index/qdrant.py b/semantic_router/index/qdrant.py
index 0da5c25eb9dd8891a5d1aa9ef1338b8a07c0caf3..1b23753ba886db48947a14108d1fe7a0ced56c42 100644
--- a/semantic_router/index/qdrant.py
+++ b/semantic_router/index/qdrant.py
@@ -1,7 +1,7 @@
from typing import Any, Dict, List, Optional, Tuple, Union
import numpy as np
-from pydantic.v1 import Field
+from pydantic import Field
from semantic_router.index.base import BaseIndex
from semantic_router.schema import ConfigParameter, Metric, SparseEmbedding, Utterance
@@ -40,7 +40,7 @@ class QdrantIndex(BaseIndex):
default=6334,
description="Port of the gRPC interface.",
)
- prefer_grpc: bool = Field(
+ prefer_grpc: Optional[bool] = Field(
default=None,
description="Whether to use gPRC interface whenever possible in methods",
)
diff --git a/semantic_router/llms/base.py b/semantic_router/llms/base.py
index bbd39b4ef4fdfe2fa6527cba6419d91137fd7946..02df2764502d2d2bdd70135406c5e27c5a900fda 100644
--- a/semantic_router/llms/base.py
+++ b/semantic_router/llms/base.py
@@ -1,7 +1,7 @@
import json
from typing import Any, List, Optional, Dict
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
from semantic_router.schema import Message
from semantic_router.utils.logger import logger
@@ -9,6 +9,8 @@ from semantic_router.utils.logger import logger
class BaseLLM(BaseModel):
name: str
+ temperature: Optional[float] = 0.0
+ max_tokens: Optional[int] = None
class Config:
arbitrary_types_allowed = True
diff --git a/semantic_router/llms/cohere.py b/semantic_router/llms/cohere.py
index 05a9b1bdb0f84f175103e741742777525469a3f1..d37f979d623613c91bbc9a433b6b03d7a9f9cb3f 100644
--- a/semantic_router/llms/cohere.py
+++ b/semantic_router/llms/cohere.py
@@ -1,7 +1,7 @@
import os
from typing import Any, List, Optional
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
from semantic_router.llms import BaseLLM
from semantic_router.schema import Message
diff --git a/semantic_router/llms/llamacpp.py b/semantic_router/llms/llamacpp.py
index 102f7fff7d253533a9b6e195cf5aaa8f7a9d3984..8257593c6f250a8bc1ded5f3af0dbc30a8b44141 100644
--- a/semantic_router/llms/llamacpp.py
+++ b/semantic_router/llms/llamacpp.py
@@ -2,7 +2,7 @@ from contextlib import contextmanager
from pathlib import Path
from typing import Any, Optional, List, Dict
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
from semantic_router.llms.base import BaseLLM
from semantic_router.schema import Message
@@ -11,8 +11,6 @@ from semantic_router.utils.logger import logger
class LlamaCppLLM(BaseLLM):
llm: Any
- temperature: float
- max_tokens: Optional[int] = 200
grammar: Optional[Any] = None
_llama_cpp: Any = PrivateAttr()
diff --git a/semantic_router/llms/mistral.py b/semantic_router/llms/mistral.py
index 8ddd1482975f4cb091816bcd197c416ac3e39da7..370fba65756224ded1af6841382f02f63578851d 100644
--- a/semantic_router/llms/mistral.py
+++ b/semantic_router/llms/mistral.py
@@ -1,7 +1,7 @@
import os
from typing import Any, List, Optional
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
from semantic_router.llms import BaseLLM
from semantic_router.schema import Message
@@ -11,8 +11,6 @@ from semantic_router.utils.logger import logger
class MistralAILLM(BaseLLM):
_client: Any = PrivateAttr()
- temperature: Optional[float]
- max_tokens: Optional[int]
_mistralai: Any = PrivateAttr()
def __init__(
diff --git a/semantic_router/llms/ollama.py b/semantic_router/llms/ollama.py
index df35ac06a97748ece6a2da86fad58bb7e4a0a52f..f6e9779ea65197e08f47135af8fec0c2dba8f295 100644
--- a/semantic_router/llms/ollama.py
+++ b/semantic_router/llms/ollama.py
@@ -8,22 +8,17 @@ from semantic_router.utils.logger import logger
class OllamaLLM(BaseLLM):
- temperature: Optional[float]
- llm_name: Optional[str]
- max_tokens: Optional[int]
- stream: Optional[bool]
+ stream: bool = False
def __init__(
self,
- name: str = "ollama",
+ name: str = "openhermes",
temperature: float = 0.2,
- llm_name: str = "openhermes",
max_tokens: Optional[int] = 200,
stream: bool = False,
):
super().__init__(name=name)
self.temperature = temperature
- self.llm_name = llm_name
self.max_tokens = max_tokens
self.stream = stream
@@ -31,19 +26,19 @@ class OllamaLLM(BaseLLM):
self,
messages: List[Message],
temperature: Optional[float] = None,
- llm_name: Optional[str] = None,
+ name: Optional[str] = None,
max_tokens: Optional[int] = None,
stream: Optional[bool] = None,
) -> str:
# Use instance defaults if not overridden
temperature = temperature if temperature is not None else self.temperature
- llm_name = llm_name if llm_name is not None else self.llm_name
+ name = name if name is not None else self.name
max_tokens = max_tokens if max_tokens is not None else self.max_tokens
stream = stream if stream is not None else self.stream
try:
payload = {
- "model": llm_name,
+ "model": name,
"messages": [m.to_openai() for m in messages],
"options": {"temperature": temperature, "num_predict": max_tokens},
"format": "json",
diff --git a/semantic_router/llms/openai.py b/semantic_router/llms/openai.py
index dfff80968a588cc1dbf387fc049219b64718f9ec..7dc565e8877a46d74e5fd86d4e3d05e04c52175c 100644
--- a/semantic_router/llms/openai.py
+++ b/semantic_router/llms/openai.py
@@ -1,5 +1,6 @@
import os
from typing import List, Optional, Any, Callable, Dict, Union
+from pydantic import PrivateAttr
import openai
from openai._types import NotGiven, NOT_GIVEN
@@ -21,10 +22,8 @@ from openai.types.chat.chat_completion_message_tool_call import (
class OpenAILLM(BaseLLM):
- client: Optional[openai.OpenAI]
- async_client: Optional[openai.AsyncOpenAI]
- temperature: Optional[float]
- max_tokens: Optional[int]
+ _client: Optional[openai.OpenAI] = PrivateAttr(default=None)
+ _async_client: Optional[openai.AsyncOpenAI] = PrivateAttr(default=None)
def __init__(
self,
@@ -40,8 +39,8 @@ class OpenAILLM(BaseLLM):
if api_key is None:
raise ValueError("OpenAI API key cannot be 'None'.")
try:
- self.async_client = openai.AsyncOpenAI(api_key=api_key)
- self.client = openai.OpenAI(api_key=api_key)
+ self._async_client = openai.AsyncOpenAI(api_key=api_key)
+ self._client = openai.OpenAI(api_key=api_key)
except Exception as e:
raise ValueError(
f"OpenAI API client failed to initialize. Error: {e}"
@@ -88,14 +87,14 @@ class OpenAILLM(BaseLLM):
messages: List[Message],
function_schemas: Optional[List[Dict[str, Any]]] = None,
) -> str:
- if self.client is None:
+ if self._client is None:
raise ValueError("OpenAI client is not initialized.")
try:
tools: Union[List[Dict[str, Any]], NotGiven] = (
function_schemas if function_schemas else NOT_GIVEN
)
- completion = self.client.chat.completions.create(
+ completion = self._client.chat.completions.create(
model=self.name,
messages=[m.to_openai() for m in messages],
temperature=self.temperature,
@@ -132,14 +131,14 @@ class OpenAILLM(BaseLLM):
messages: List[Message],
function_schemas: Optional[List[Dict[str, Any]]] = None,
) -> str:
- if self.async_client is None:
+ if self._async_client is None:
raise ValueError("OpenAI async_client is not initialized.")
try:
tools: Union[List[Dict[str, Any]], NotGiven] = (
function_schemas if function_schemas is not None else NOT_GIVEN
)
- completion = await self.async_client.chat.completions.create(
+ completion = await self._async_client.chat.completions.create(
model=self.name,
messages=[m.to_openai() for m in messages],
temperature=self.temperature,
diff --git a/semantic_router/llms/openrouter.py b/semantic_router/llms/openrouter.py
index b00d68a4730c6bf681c0ba4d90a5a79c7febe603..0376236486c06f32775730bb332f13c5f87b0ab8 100644
--- a/semantic_router/llms/openrouter.py
+++ b/semantic_router/llms/openrouter.py
@@ -1,6 +1,8 @@
import os
from typing import List, Optional
+from pydantic import PrivateAttr
+
import openai
from semantic_router.llms import BaseLLM
@@ -9,10 +11,8 @@ from semantic_router.utils.logger import logger
class OpenRouterLLM(BaseLLM):
- client: Optional[openai.OpenAI]
- base_url: Optional[str]
- temperature: Optional[float]
- max_tokens: Optional[int]
+ _client: Optional[openai.OpenAI] = PrivateAttr(default=None)
+ _base_url: str = PrivateAttr(default="https://openrouter.ai/api/v1")
def __init__(
self,
@@ -27,12 +27,12 @@ class OpenRouterLLM(BaseLLM):
"OPENROUTER_CHAT_MODEL_NAME", "mistralai/mistral-7b-instruct"
)
super().__init__(name=name)
- self.base_url = base_url
+ 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)
+ 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}"
@@ -41,10 +41,10 @@ class OpenRouterLLM(BaseLLM):
self.max_tokens = max_tokens
def __call__(self, messages: List[Message]) -> str:
- if self.client is None:
+ if self._client is None:
raise ValueError("OpenRouter client is not initialized.")
try:
- completion = self.client.chat.completions.create(
+ completion = self._client.chat.completions.create(
model=self.name,
messages=[m.to_openai() for m in messages],
temperature=self.temperature,
diff --git a/semantic_router/llms/zure.py b/semantic_router/llms/zure.py
index 26b7901f7b595ed1705a81484598f4389a242d26..ba833d044785b775b0da123db554a79f55dd5be8 100644
--- a/semantic_router/llms/zure.py
+++ b/semantic_router/llms/zure.py
@@ -1,5 +1,6 @@
import os
from typing import List, Optional
+from pydantic import PrivateAttr
import openai
@@ -10,9 +11,7 @@ from semantic_router.utils.logger import logger
class AzureOpenAILLM(BaseLLM):
- client: Optional[openai.AzureOpenAI]
- temperature: Optional[float]
- max_tokens: Optional[int]
+ _client: Optional[openai.AzureOpenAI] = PrivateAttr(default=None)
def __init__(
self,
@@ -33,7 +32,7 @@ class AzureOpenAILLM(BaseLLM):
if azure_endpoint is None:
raise ValueError("Azure endpoint API key cannot be 'None'.")
try:
- self.client = openai.AzureOpenAI(
+ self._client = openai.AzureOpenAI(
api_key=api_key, azure_endpoint=azure_endpoint, api_version=api_version
)
except Exception as e:
@@ -42,10 +41,10 @@ class AzureOpenAILLM(BaseLLM):
self.max_tokens = max_tokens
def __call__(self, messages: List[Message]) -> str:
- if self.client is None:
+ if self._client is None:
raise ValueError("AzureOpenAI client is not initialized.")
try:
- completion = self.client.chat.completions.create(
+ completion = self._client.chat.completions.create(
model=self.name,
messages=[m.to_openai() for m in messages],
temperature=self.temperature,
diff --git a/semantic_router/route.py b/semantic_router/route.py
index 50d516fa5bbe31d7eff25b8ece8372801649abbd..b1fc5e8b699bcaa7a117fabc7f8654a2ae4f4f9c 100644
--- a/semantic_router/route.py
+++ b/semantic_router/route.py
@@ -2,18 +2,13 @@ import json
import re
from typing import Any, Callable, Dict, List, Optional, Union
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
from semantic_router.llms import BaseLLM
from semantic_router.schema import Message, RouteChoice
from semantic_router.utils import function_call
from semantic_router.utils.logger import logger
-try:
- from PIL.Image import Image
-except ImportError:
- pass
-
def is_valid(route_config: str) -> bool:
try:
@@ -45,7 +40,7 @@ def is_valid(route_config: str) -> bool:
class Route(BaseModel):
name: str
- utterances: Union[List[str], List[Union[Any, "Image"]]]
+ utterances: Union[List[str], List[Any]]
description: Optional[str] = None
function_schemas: Optional[List[Dict[str, Any]]] = None
llm: Optional[BaseLLM] = None
diff --git a/semantic_router/routers/base.py b/semantic_router/routers/base.py
index e4376b421b1d9a87b964443795c2ccd4d2e1cab5..3628fda8553a8a6e6979d6684e2fdfaa1611ead5 100644
--- a/semantic_router/routers/base.py
+++ b/semantic_router/routers/base.py
@@ -4,7 +4,7 @@ import os
import random
import hashlib
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
-from pydantic.v1 import BaseModel, Field
+from pydantic import BaseModel, Field
import numpy as np
import yaml # type: ignore
diff --git a/semantic_router/routers/hybrid.py b/semantic_router/routers/hybrid.py
index 36ccd8f94c5e4aa1360462a34ad10382f0c92171..994fcb2dfc8c5e4d85589ca5cf777c2ab7a26d26 100644
--- a/semantic_router/routers/hybrid.py
+++ b/semantic_router/routers/hybrid.py
@@ -1,6 +1,6 @@
from typing import Dict, List, Optional
import asyncio
-from pydantic.v1 import Field
+from pydantic import Field
import numpy as np
diff --git a/semantic_router/schema.py b/semantic_router/schema.py
index c9e63943823aef590f993c0a2910197d2b78a0f9..adb881fa1cedf56a2e103b7208e2b518e1bf0799 100644
--- a/semantic_router/schema.py
+++ b/semantic_router/schema.py
@@ -3,7 +3,7 @@ from difflib import Differ
from enum import Enum
import numpy as np
from typing import List, Optional, Union, Any, Dict, Tuple
-from pydantic.v1 import BaseModel, Field
+from pydantic import BaseModel, Field
from semantic_router.utils.logger import logger
from aurelio_sdk.schema import BM25Embedding
diff --git a/semantic_router/utils/function_call.py b/semantic_router/utils/function_call.py
index d03c0798836294400c44ddca280a75e09e500988..f5fd5a87fd100482cd65f3502c3f8f67349e36c2 100644
--- a/semantic_router/utils/function_call.py
+++ b/semantic_router/utils/function_call.py
@@ -1,7 +1,7 @@
import inspect
from typing import Any, Callable, Dict, List, Optional, Union
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
from semantic_router.llms import BaseLLM
from semantic_router.schema import Message, RouteChoice
diff --git a/tests/integration/encoders/test_openai_integration.py b/tests/integration/encoders/test_openai_integration.py
index c7bd8d9add219c14a9937ba2325c34ebf5b679e7..8c557d809d54413ae28a0676d5abe56026d29022 100644
--- a/tests/integration/encoders/test_openai_integration.py
+++ b/tests/integration/encoders/test_openai_integration.py
@@ -21,7 +21,7 @@ class TestOpenAIEncoder:
os.environ.get("OPENAI_API_KEY") is None, reason="OpenAI API key required"
)
def test_openai_encoder_init_success(self, openai_encoder):
- assert openai_encoder.client is not None
+ assert openai_encoder._client is not None
@pytest.mark.skipif(
os.environ.get("OPENAI_API_KEY") is None, reason="OpenAI API key required"
@@ -50,7 +50,7 @@ class TestOpenAIEncoder:
)
def test_openai_encoder_call_uninitialized_client(self, openai_encoder):
# Set the client to None to simulate an uninitialized client
- openai_encoder.client = None
+ openai_encoder._client = None
with pytest.raises(ValueError) as e:
openai_encoder(["test document"])
assert "OpenAI client is not initialized." in str(e.value)
diff --git a/tests/unit/encoders/test_clip.py b/tests/unit/encoders/test_clip.py
index 1cfc7d4313a1a17302dff195ac71e137e3ae4581..297761519f29870723970da7f6ece10e98ed9ff8 100644
--- a/tests/unit/encoders/test_clip.py
+++ b/tests/unit/encoders/test_clip.py
@@ -1,11 +1,14 @@
import os
import numpy as np
import pytest
-import torch
-from PIL import Image
-from unittest.mock import patch
-from semantic_router.encoders import CLIPEncoder
+_ = pytest.importorskip("torch")
+
+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
test_model_name = "aurelio-ai/sr-test-clip"
embed_dim = 64
diff --git a/tests/unit/encoders/test_fastembed.py b/tests/unit/encoders/test_fastembed.py
index 9b0f32296117db7f9a16ea09284082ce5a24498f..f46303c1bb069653c792fcf00c05ee9ced112923 100644
--- a/tests/unit/encoders/test_fastembed.py
+++ b/tests/unit/encoders/test_fastembed.py
@@ -1,5 +1,9 @@
from semantic_router.encoders import FastEmbedEncoder
+import pytest
+
+_ = pytest.importorskip("fastembed")
+
class TestFastEmbedEncoder:
def test_fastembed_encoder(self):
diff --git a/tests/unit/encoders/test_hfendpointencoder.py b/tests/unit/encoders/test_hfendpointencoder.py
index cb8dd16a661c4194a76935d51390ed4f5b3d3fcc..840df2a9fd0f27ce2cd0dab4ca46b15c994cd5ce 100644
--- a/tests/unit/encoders/test_hfendpointencoder.py
+++ b/tests/unit/encoders/test_hfendpointencoder.py
@@ -1,4 +1,5 @@
import pytest
+
from semantic_router.encoders.huggingface import HFEndpointEncoder
diff --git a/tests/unit/encoders/test_huggingface.py b/tests/unit/encoders/test_huggingface.py
index b615a87d13d64d0a45f6542f137d31ef9fc9a05f..7f496e39f8117dbe4242070e59517acf751ca9a0 100644
--- a/tests/unit/encoders/test_huggingface.py
+++ b/tests/unit/encoders/test_huggingface.py
@@ -4,7 +4,9 @@ import os
import numpy as np
import pytest
-from semantic_router.encoders.huggingface import HuggingFaceEncoder
+_ = pytest.importorskip("transformers")
+
+from semantic_router.encoders.huggingface import HuggingFaceEncoder # noqa: E402
test_model_name = "aurelio-ai/sr-test-huggingface"
diff --git a/tests/unit/encoders/test_mistral.py b/tests/unit/encoders/test_mistral.py
index 25dba6b759bc2a033c64ae0055e0b1b45d626f48..a236ad167139b5d5994ed8ad2d3610e4b49fc349 100644
--- a/tests/unit/encoders/test_mistral.py
+++ b/tests/unit/encoders/test_mistral.py
@@ -1,6 +1,7 @@
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 538e969250051c3da305702a42944903d1cb71cf..96171cce5b462f76dfbd88d7795bfc684ca100c8 100644
--- a/tests/unit/encoders/test_openai.py
+++ b/tests/unit/encoders/test_openai.py
@@ -30,7 +30,7 @@ class TestOpenAIEncoder:
side_effect = ["fake-model-name", "fake-api-key", "fake-org-id"]
mocker.patch("os.getenv", side_effect=side_effect)
encoder = OpenAIEncoder()
- assert encoder.client is not None
+ assert encoder._client is not None
def test_openai_encoder_init_no_api_key(self, mocker):
mocker.patch("os.getenv", return_value=None)
@@ -39,7 +39,7 @@ class TestOpenAIEncoder:
def test_openai_encoder_call_uninitialized_client(self, openai_encoder):
# Set the client to None to simulate an uninitialized client
- openai_encoder.client = None
+ openai_encoder._client = None
with pytest.raises(ValueError) as e:
openai_encoder(["test document"])
assert "OpenAI client is not initialized." in str(e.value)
@@ -74,7 +74,7 @@ class TestOpenAIEncoder:
responses = [OpenAIError("OpenAI error"), mock_response]
mocker.patch.object(
- openai_encoder.client.embeddings, "create", side_effect=responses
+ openai_encoder._client.embeddings, "create", side_effect=responses
)
with patch("semantic_router.encoders.openai.sleep", return_value=None):
embeddings = openai_encoder(["test document"])
@@ -84,7 +84,7 @@ class TestOpenAIEncoder:
mocker.patch("os.getenv", return_value="fake-api-key")
mocker.patch("time.sleep", return_value=None) # To speed up the test
mocker.patch.object(
- openai_encoder.client.embeddings,
+ openai_encoder._client.embeddings,
"create",
side_effect=Exception("Non-OpenAIError"),
)
@@ -114,7 +114,7 @@ class TestOpenAIEncoder:
responses = [OpenAIError("OpenAI error"), mock_response]
mocker.patch.object(
- openai_encoder.client.embeddings, "create", side_effect=responses
+ openai_encoder._client.embeddings, "create", side_effect=responses
)
with patch("semantic_router.encoders.openai.sleep", return_value=None):
embeddings = openai_encoder(["test document"])
diff --git a/tests/unit/encoders/test_vit.py b/tests/unit/encoders/test_vit.py
index b598a818a3af0753a0627fead4801a56c5b03f54..622b22689f327054392a98cb336ec3f40adee6cd 100644
--- a/tests/unit/encoders/test_vit.py
+++ b/tests/unit/encoders/test_vit.py
@@ -3,10 +3,13 @@ from unittest.mock import patch
import os
import numpy as np
import pytest
-import torch
-from PIL import Image
-from semantic_router.encoders import VitEncoder
+_ = pytest.importorskip("torch")
+
+import torch # noqa: E402
+from PIL import Image # noqa: E402
+
+from semantic_router.encoders import VitEncoder # noqa: E402
test_model_name = "aurelio-ai/sr-test-vit"
embed_dim = 32
diff --git a/tests/unit/llms/test_llm_azure_openai.py b/tests/unit/llms/test_llm_azure_openai.py
index 793091957df7fdd89792aebec43a6fe87763e02b..def792b372790a912b68956bd83e5aa5023a5b7e 100644
--- a/tests/unit/llms/test_llm_azure_openai.py
+++ b/tests/unit/llms/test_llm_azure_openai.py
@@ -12,13 +12,13 @@ def azure_openai_llm(mocker):
class TestOpenAILLM:
def test_azure_openai_llm_init_with_api_key(self, azure_openai_llm):
- assert azure_openai_llm.client is not None, "Client should be initialized"
+ assert azure_openai_llm._client is not None, "Client should be initialized"
assert azure_openai_llm.name == "gpt-4o", "Default name not set correctly"
def test_azure_openai_llm_init_success(self, mocker):
mocker.patch("os.getenv", return_value="fake-api-key")
llm = AzureOpenAILLM()
- assert llm.client is not None
+ assert llm._client is not None
def test_azure_openai_llm_init_without_api_key(self, mocker):
mocker.patch("os.getenv", return_value=None)
@@ -44,7 +44,7 @@ class TestOpenAILLM:
def test_azure_openai_llm_call_uninitialized_client(self, azure_openai_llm):
# Set the client to None to simulate an uninitialized client
- azure_openai_llm.client = None
+ azure_openai_llm._client = None
with pytest.raises(ValueError) as e:
llm_input = [Message(role="user", content="test")]
azure_openai_llm(llm_input)
@@ -83,7 +83,7 @@ class TestOpenAILLM:
mocker.patch("os.getenv", return_value="fake-api-key")
mocker.patch.object(
- azure_openai_llm.client.chat.completions,
+ azure_openai_llm._client.chat.completions,
"create",
return_value=mock_completion,
)
diff --git a/tests/unit/llms/test_llm_llamacpp.py b/tests/unit/llms/test_llm_llamacpp.py
index 04a3ad3bc5528d8797ec326b444ba6af2db1c236..fc9813faad19f8c69cf7a21cff9d614d8fc62888 100644
--- a/tests/unit/llms/test_llm_llamacpp.py
+++ b/tests/unit/llms/test_llm_llamacpp.py
@@ -1,10 +1,13 @@
from unittest.mock import patch
import pytest
-from llama_cpp import Llama
-from semantic_router.llms.llamacpp import LlamaCppLLM
-from semantic_router.schema import Message
+_ = pytest.importorskip("llama_cpp")
+
+from llama_cpp import Llama # noqa: E402
+
+from semantic_router.llms.llamacpp import LlamaCppLLM # noqa: E402
+from semantic_router.schema import Message # noqa: E402
@pytest.fixture
diff --git a/tests/unit/llms/test_llm_ollama.py b/tests/unit/llms/test_llm_ollama.py
index 369e5f4444789c3bb70988455bb49c7df4792445..91af8a3c7b3a565ade17c5ecb529b080ecebf978 100644
--- a/tests/unit/llms/test_llm_ollama.py
+++ b/tests/unit/llms/test_llm_ollama.py
@@ -11,9 +11,8 @@ def ollama_llm():
class TestOllamaLLM:
def test_ollama_llm_init_success(self, ollama_llm):
- assert ollama_llm.name == "ollama"
assert ollama_llm.temperature == 0.2
- assert ollama_llm.llm_name == "openhermes"
+ assert ollama_llm.name == "openhermes"
assert ollama_llm.max_tokens == 200
assert ollama_llm.stream is False
diff --git a/tests/unit/llms/test_llm_openai.py b/tests/unit/llms/test_llm_openai.py
index 4287fc882214c8524c5ce611c11af136eea935c5..ee7e82a9d741d31cc7303756a87e5dd007d32f72 100644
--- a/tests/unit/llms/test_llm_openai.py
+++ b/tests/unit/llms/test_llm_openai.py
@@ -42,13 +42,13 @@ example_function_schema = {
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._client is not None, "Client should be initialized"
assert openai_llm.name == "gpt-4o", "Default name not set correctly"
def test_openai_llm_init_success(self, mocker):
mocker.patch("os.getenv", return_value="fake-api-key")
llm = OpenAILLM()
- assert llm.client is not None
+ assert llm._client is not None
def test_openai_llm_init_without_api_key(self, mocker):
mocker.patch("os.getenv", return_value=None)
@@ -57,7 +57,7 @@ class TestOpenAILLM:
def test_openai_llm_call_uninitialized_client(self, openai_llm):
# Set the client to None to simulate an uninitialized client
- openai_llm.client = None
+ openai_llm._client = None
with pytest.raises(ValueError) as e:
llm_input = [Message(role="user", content="test")]
openai_llm(llm_input)
@@ -79,7 +79,7 @@ class TestOpenAILLM:
mocker.patch("os.getenv", return_value="fake-api-key")
mocker.patch.object(
- openai_llm.client.chat.completions, "create", return_value=mock_completion
+ openai_llm._client.chat.completions, "create", return_value=mock_completion
)
llm_input = [Message(role="user", content="test")]
output = openai_llm(llm_input)
@@ -127,7 +127,7 @@ class TestOpenAILLM:
# mocker.MagicMock(function=mocker.MagicMock(arguments="result"))
# ]
# mocker.patch.object(
- # openai_llm.client.chat.completions, "create", return_value=mock_completion
+ # openai_llm._client.chat.completions, "create", return_value=mock_completion
# )
# llm_input = [Message(role="user", content="test")]
# function_schemas = [{"type": "function", "name": "sample_function"}]
@@ -145,7 +145,7 @@ class TestOpenAILLM:
mock_completion.choices[0].message.tool_calls = [mock_tool_call]
mocker.patch.object(
- openai_llm.client.chat.completions, "create", return_value=mock_completion
+ openai_llm._client.chat.completions, "create", return_value=mock_completion
)
llm_input = [Message(role="user", content="test")]
@@ -160,7 +160,7 @@ class TestOpenAILLM:
mock_completion = mocker.MagicMock()
mock_completion.choices[0].message.tool_calls = None
mocker.patch.object(
- openai_llm.client.chat.completions, "create", return_value=mock_completion
+ openai_llm._client.chat.completions, "create", return_value=mock_completion
)
llm_input = [Message(role="user", content="test")]
function_schemas = [{"type": "function", "name": "sample_function"}]
@@ -180,7 +180,7 @@ class TestOpenAILLM:
mocker.MagicMock(function=mocker.MagicMock(arguments=None))
]
mocker.patch.object(
- openai_llm.client.chat.completions, "create", return_value=mock_completion
+ openai_llm._client.chat.completions, "create", return_value=mock_completion
)
llm_input = [Message(role="user", content="test")]
function_schemas = [{"type": "function", "name": "sample_function"}]
@@ -230,7 +230,7 @@ class TestOpenAILLM:
# Patching the completions.create method to return the mocked completion
mocker.patch.object(
- openai_llm.client.chat.completions, "create", return_value=mock_completion
+ openai_llm._client.chat.completions, "create", return_value=mock_completion
)
# Input message list
diff --git a/tests/unit/llms/test_llm_openrouter.py b/tests/unit/llms/test_llm_openrouter.py
index 9b1ee150f2b301984c24eeb144453cd6a5ea0973..71e874f10124748c5e32e869807dcec2fcb0960d 100644
--- a/tests/unit/llms/test_llm_openrouter.py
+++ b/tests/unit/llms/test_llm_openrouter.py
@@ -12,7 +12,7 @@ def openrouter_llm(mocker):
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._client is not None, "Client should be initialized"
assert (
openrouter_llm.name == "mistralai/mistral-7b-instruct"
), "Default name not set correctly"
@@ -20,7 +20,7 @@ class TestOpenRouterLLM:
def test_openrouter_llm_init_success(self, mocker):
mocker.patch("os.getenv", return_value="fake-api-key")
llm = OpenRouterLLM()
- assert llm.client is not None
+ assert llm._client is not None
def test_openrouter_llm_init_without_api_key(self, mocker):
mocker.patch("os.getenv", return_value=None)
@@ -29,7 +29,7 @@ class TestOpenRouterLLM:
def test_openrouter_llm_call_uninitialized_client(self, openrouter_llm):
# Set the client to None to simulate an uninitialized client
- openrouter_llm.client = None
+ openrouter_llm._client = None
with pytest.raises(ValueError) as e:
llm_input = [Message(role="user", content="test")]
openrouter_llm(llm_input)
@@ -51,7 +51,7 @@ class TestOpenRouterLLM:
mocker.patch("os.getenv", return_value="fake-api-key")
mocker.patch.object(
- openrouter_llm.client.chat.completions,
+ openrouter_llm._client.chat.completions,
"create",
return_value=mock_completion,
)
diff --git a/tests/unit/test_schema.py b/tests/unit/test_schema.py
index d0fce781be2bbf948b43b7a48d3047ab7318b92a..bd5d9729a53d40acbbf76ba48b39bd22b92b214e 100644
--- a/tests/unit/test_schema.py
+++ b/tests/unit/test_schema.py
@@ -1,5 +1,5 @@
import pytest
-from pydantic.v1 import ValidationError
+from pydantic import ValidationError
from semantic_router.schema import (
Message,