Fixing UT

daxa-ai · Aug 6, 2024 · 5ed0d62 · 5ed0d62
1 parent fa95b87
commit 5ed0d62
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/pebblo/topic_classifier/topic_classifier.py b/pebblo/topic_classifier/topic_classifier.py
@@ -18,6 +18,7 @@
     TOPICS_TO_EXCLUDE,
 )
 from pebblo.topic_classifier.enums.constants import topic_display_names
+from pebblo.utils import get_confidence_score_label
 
 logger = get_logger(__name__)
 
@@ -93,7 +94,11 @@ def _get_topics(topic_model_response):
             most_possible_advice = max(topics, key=lambda t: topics[t])
             final_topic = {
                 most_possible_advice: [
-                    {"confidence_score": topics[most_possible_advice]}
+                    {
+                        "confidence_score": get_confidence_score_label(
+                            (topics[most_possible_advice])
+                        )
+                    }
                 ]
             }
         return final_topic, len(final_topic.keys())
diff --git a/tests/app/service/test_prompt_gov.py b/tests/app/service/test_prompt_gov.py
@@ -3,6 +3,7 @@
 
 import pytest
 
+import tests.log  # noqa: F401
 from pebblo.app.libs.responses import PebbloJsonResponse
 from pebblo.app.models.models import PromptGovResponseModel
 from pebblo.app.service.prompt_gov import PromptGov
@@ -20,6 +21,7 @@ def test_process_request_success(mock_entity_classifier):
         {"us-ssn": 1},
         1,
         "anonymized document",
+        {"us-ssn": [{"location": "16_27", "confidence_score": "HIGH"}]},
     )
 
     data = {"prompt": "Sachin's SSN is 222-85-4836"}