lint

aurelio-labs · Dec 2, 2023 · 1eb511c · 1eb511c
1 parent 839c259
commit 1eb511c
Showing 1 changed file with 18 additions and 5 deletions.
diff --git a/tests/unit/test_layer.py b/tests/unit/test_layer.py
@@ -1,7 +1,10 @@
 import pytest
 
 from semantic_router.encoders import BaseEncoder, CohereEncoder, OpenAIEncoder
-from semantic_router.layer import DecisionLayer, HybridDecisionLayer  # Replace with the actual module name
+from semantic_router.layer import (
+    DecisionLayer,
+    HybridDecisionLayer,
+)  # Replace with the actual module name
 from semantic_router.schema import Decision
 
 
@@ -111,9 +114,12 @@ def test_failover_score_threshold(self, base_encoder):
         decision_layer = DecisionLayer(encoder=base_encoder)
         assert decision_layer.score_threshold == 0.82
 
+
 class TestHybridDecisionLayer:
     def test_initialization(self, openai_encoder, decisions):
-        decision_layer = HybridDecisionLayer(encoder=openai_encoder, decisions=decisions)
+        decision_layer = HybridDecisionLayer(
+            encoder=openai_encoder, decisions=decisions
+        )
         assert decision_layer.score_threshold == 0.82
         assert len(decision_layer.index) == 5
         assert len(set(decision_layer.categories)) == 2
@@ -140,7 +146,9 @@ def test_add_multiple_decisions(self, openai_encoder, decisions):
         assert len(set(decision_layer.categories)) == 2
 
     def test_query_and_classification(self, openai_encoder, decisions):
-        decision_layer = HybridDecisionLayer(encoder=openai_encoder, decisions=decisions)
+        decision_layer = HybridDecisionLayer(
+            encoder=openai_encoder, decisions=decisions
+        )
         query_result = decision_layer("Hello")
         assert query_result in ["Decision 1", "Decision 2"]
 
@@ -149,7 +157,9 @@ def test_query_with_no_index(self, openai_encoder):
         assert decision_layer("Anything") is None
 
     def test_semantic_classify(self, openai_encoder, decisions):
-        decision_layer = HybridDecisionLayer(encoder=openai_encoder, decisions=decisions)
+        decision_layer = HybridDecisionLayer(
+            encoder=openai_encoder, decisions=decisions
+        )
         classification, score = decision_layer._semantic_classify(
             [
                 {"decision": "Decision 1", "score": 0.9},
@@ -160,7 +170,9 @@ def test_semantic_classify(self, openai_encoder, decisions):
         assert score == [0.9]
 
     def test_semantic_classify_multiple_decisions(self, openai_encoder, decisions):
-        decision_layer = HybridDecisionLayer(encoder=openai_encoder, decisions=decisions)
+        decision_layer = HybridDecisionLayer(
+            encoder=openai_encoder, decisions=decisions
+        )
         classification, score = decision_layer._semantic_classify(
             [
                 {"decision": "Decision 1", "score": 0.9},
@@ -180,4 +192,5 @@ def test_failover_score_threshold(self, base_encoder):
         decision_layer = HybridDecisionLayer(encoder=base_encoder)
         assert decision_layer.score_threshold == 0.82
 
+
 # Add more tests for edge cases and error handling as needed.