fix: types

semantic_router/llms/base.py

@@ -9,6 +9,8 @@
 
 class BaseLLM(BaseModel):
     name: str
+    temperature: Optional[float] = 0.0
+    max_tokens: Optional[int] = None
 
     class Config:
         arbitrary_types_allowed = True

semantic_router/llms/llamacpp.py

@@ -11,8 +11,6 @@
 
 class LlamaCppLLM(BaseLLM):
     llm: Any
-    temperature: float
-    max_tokens: Optional[int] = 200
     grammar: Optional[Any] = None
     _llama_cpp: Any = PrivateAttr()
 

semantic_router/llms/mistral.py

@@ -11,8 +11,6 @@
 
 class MistralAILLM(BaseLLM):
     _client: Any = PrivateAttr()
-    temperature: Optional[float]
-    max_tokens: Optional[int]
     _mistralai: Any = PrivateAttr()
 
     def __init__(

semantic_router/llms/ollama.py

@@ -8,42 +8,37 @@
 
 
 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
 
     def __call__(
         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",

semantic_router/llms/openai.py

@@ -23,8 +23,6 @@
 class OpenAILLM(BaseLLM):
     client: Optional[openai.OpenAI]
     async_client: Optional[openai.AsyncOpenAI]
-    temperature: Optional[float]
-    max_tokens: Optional[int]
 
     def __init__(
         self,

semantic_router/llms/openrouter.py

@@ -11,8 +11,6 @@
 class OpenRouterLLM(BaseLLM):
     client: Optional[openai.OpenAI]
     base_url: Optional[str]
-    temperature: Optional[float]
-    max_tokens: Optional[int]
 
     def __init__(
         self,

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 @@
 
 
 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 @@ def __init__(
         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 @@ def __init__(
         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,

tests/unit/encoders/test_fastembed.py

@@ -1,5 +1,8 @@
 from semantic_router.encoders import FastEmbedEncoder
 
+import pytest
+
+_ = pytest.importorskip("fastembed")
 
 class TestFastEmbedEncoder:
     def test_fastembed_encoder(self):

tests/unit/encoders/test_hfendpointencoder.py

@@ -1,4 +1,5 @@
 import pytest
+
 from semantic_router.encoders.huggingface import HFEndpointEncoder
 
 

tests/unit/encoders/test_huggingface.py

@@ -4,7 +4,9 @@
 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"
 

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
 

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