Merge pull request #61 from ai-cfia/60-format-website-fields

issue #60: format website fields

pipeline/inspection.py

@@ -1,8 +1,15 @@
 import re
-from typing import List, Optional
+from typing import Annotated, List, Optional
 
 import phonenumbers
-from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    StringConstraints,
+    field_validator,
+    model_validator,
+)
 
 
 class npkError(ValueError):
@@ -89,13 +96,21 @@ def convert_specification_values(cls, v):
 class FertilizerInspection(BaseModel):
     company_name: Optional[str] = None
     company_address: Optional[str] = None
-    company_website: Optional[str] = None
+    company_website: Annotated[str | None, StringConstraints(to_lower=True)] = Field(
+        None,
+        description="Return the distributor's website, ensuring 'www.' prefix is added.",
+    )
     company_phone_number: Optional[str] = Field(
         None, description="The distributor's primary phone number. Return only one."
     )
     manufacturer_name: Optional[str] = None
     manufacturer_address: Optional[str] = None
-    manufacturer_website: Optional[str] = None
+    manufacturer_website: Annotated[str | None, StringConstraints(to_lower=True)] = (
+        Field(
+            None,
+            description="Return the manufacturer's website, ensuring 'www.' prefix is added.",
+        )
+    )
     manufacturer_phone_number: Optional[str] = Field(
         None, description="The manufacturer's primary phone number. Return only one."
     )

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fertiscan_pipeline"
-version = "0.0.5"
+version = "0.0.6"
 description = "A pipeline for the FertiScan project"
 authors = [
     { name = "Albert Bryan Ndjeutcha", email = "[email protected]" }

tests/test_pipeline.py

@@ -92,7 +92,7 @@ def test_analyze(self):
         self.assertFalse(os.path.exists(txt_log_path))
 
 
-class TestPhoneNumbers(unittest.TestCase):
+class TestInspectionAnnotatedFields(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
         # Load environment variables
@@ -139,7 +139,7 @@ def add_images_to_storage(self, image_paths, label_storage):
         for image_path in image_paths:
             label_storage.add_image(image_path)
 
-    def test_label_008_inspection(self):
+    def test_label_008_phone_number_inspection(self):
         label_folder = "test_data/labels/label_008"
         label_storage = LabelStorage()
 
@@ -154,7 +154,7 @@ def test_label_008_inspection(self):
         self.assertEqual(inspection.company_phone_number, "+18003279462")
         self.assertIsNone(inspection.manufacturer_phone_number)
 
-    def test_label_024_inspection(self):
+    def test_label_024_phone_number_inspection(self):
         label_folder = "test_data/labels/label_024"
         label_storage = LabelStorage()
 
@@ -169,6 +169,50 @@ def test_label_024_inspection(self):
         self.assertEqual(inspection.company_phone_number, "+14506556147")
         self.assertIsNone(inspection.manufacturer_phone_number)
 
+    def test_label_001_website_inspection(self):
+        label_folder = "test_data/labels/label_001"
+        label_storage = LabelStorage()
+
+        # Copy images to temporary directory and add to storage
+        image_paths = self.copy_images_to_temp_dir(label_folder)
+        self.add_images_to_storage(image_paths, label_storage)
+
+        # Run the analyze function
+        inspection = analyze(label_storage, self.ocr, self.gpt)
+
+        # Assertions for website fields
+        self.assertEqual(inspection.company_website, "www.soil-aid.com")
+
+    def test_label_006_website_inspection(self):
+        label_folder = "test_data/labels/label_006"
+        label_storage = LabelStorage()
+
+        # Copy images to temporary directory and add to storage
+        image_paths = self.copy_images_to_temp_dir(label_folder)
+        self.add_images_to_storage(image_paths, label_storage)
+
+        # Run the analyze function
+        inspection = analyze(label_storage, self.ocr, self.gpt)
+
+        # Assertions for website fields
+        self.assertEqual(inspection.company_website, "www.activeagriscience.com")
+
+    def test_label_034_website_inspection(self):
+        label_folder = "test_data/labels/label_034"
+        label_storage = LabelStorage()
+
+        # Copy images to temporary directory and add to storage
+        image_paths = self.copy_images_to_temp_dir(label_folder)
+        self.add_images_to_storage(image_paths, label_storage)
+
+        # Run the analyze function
+        inspection = analyze(label_storage, self.ocr, self.gpt)
+
+        # Assertions for website fields
+        self.assertEqual(
+            inspection.company_website, "www.advancednutrients.com/growersupport"
+        )
+
 
 if __name__ == "__main__":
     unittest.main()