Make main.py into a proper entrypoint (#10)

src/vsp/app/main.py

@@ -1,15 +1,28 @@
-# src/vsp/app/main.py
+"""
+main.py
+
+This module serves as the main executable file entrypoint for the VSP Data Enrichment project.
+It provides functionality to process LinkedIn profiles and classify various aspects of a person's
+educational and professional background.
+
+The main class, VspDataEnrichment, encapsulates all the necessary classifiers and methods
+to perform a comprehensive analysis of a LinkedIn profile.
+
+Usage:
+    from vsp.app.main import VspDataEnrichment
+
+    vsp_enrichment = VspDataEnrichment()
+    results = await vsp_enrichment.process_linkedin_profile(linkedin_profile)
+
+"""
 
 import asyncio
-import json
 from typing import Sequence
 
 from pydantic import BaseModel, Field
 
 from vsp.app.classifiers.education_classifier import EducationClassification, EducationClassifier
 from vsp.app.classifiers.work_experience.general_work_experience_classifier import (
-    PrimaryJobType,
-    SecondaryJobType,
     WorkExperienceClassification,
     WorkExperienceClassifier,
 )
@@ -29,14 +42,29 @@ from vsp.app.model.linkedin.linkedin_models import Education, LinkedinProfile, P
 
 
 class ClassifiedEducation(BaseModel):
-    """Represents a classified education item."""
+    """
+    Represents a classified education item from a LinkedIn profile.
+
+    Attributes:
+        education (Education): The original education item from the LinkedIn profile.
+        classification (EducationClassification): The classification results for the education item.
+    """
 
     education: Education
     classification: EducationClassification
 
 
 class ClassifiedWorkExperience(BaseModel):
-    """Represents a classified work experience item."""
+    """
+    Represents a classified work experience item from a LinkedIn profile.
+
+    Attributes:
+        position (Position): The original position item from the LinkedIn profile.
+        work_experience_classification (WorkExperienceClassification): The general classification results for the work experience.
+        investment_banking_classification (InvestmentBankingGroupClassification | None): The investment banking classification results, if applicable.
+        investing_focus_asset_class_classification (InvestingFocusAssetClassClassification | None): The investing focus asset class classification results, if applicable.
+        investing_focus_sector_classification (InvestingFocusSectorClassification | None): The investing focus sector classification results, if applicable.
+    """
 
     position: Position
     work_experience_classification: WorkExperienceClassification
@@ -46,91 +74,132 @@ class ClassifiedWorkExperience(BaseModel):
 
 
 class LinkedinProfileClassificationResults(BaseModel):
-    """Represents the classification results for a LinkedIn profile."""
+    """
+    Represents the complete classification results for a LinkedIn profile.
+
+    Attributes:
+        classified_educations (Sequence[ClassifiedEducation]): A sequence of classified education items.
+        classified_work_experiences (Sequence[ClassifiedWorkExperience]): A sequence of classified work experience items.
+    """
 
     classified_educations: Sequence[ClassifiedEducation] = Field(default_factory=list)
     classified_work_experiences: Sequence[ClassifiedWorkExperience] = Field(default_factory=list)
 
 
-async def process_linkedin_profile(profile: LinkedinProfile) -> LinkedinProfileClassificationResults:
+class VspDataEnrichment:
     """
-    Process a LinkedIn profile and classify its education and work experiences.
+    Main class for the VSP Data Enrichment project.
 
-    This function maintains the original order of educations and work experiences
-    from the input profile while performing asynchronous classification tasks.
+    This class encapsulates all the necessary classifiers and methods to process
+    and enrich LinkedIn profile data with various classifications.
 
-    Args:
-        profile (LinkedinProfile): The LinkedIn profile to process.
-
-    Returns:
-        LinkedinProfileClassificationResults: The classification results for the profile.
+    Attributes:
+        education_classifier (EducationClassifier): Classifier for education items.
+        work_experience_classifier (WorkExperienceClassifier): Classifier for general work experiences.
+        investment_banking_classifier (InvestmentBankingGroupClassifier): Classifier for investment banking groups.
+        investing_focus_asset_class_classifier (InvestingFocusAssetClassClassifier): Classifier for investing focus asset classes.
+        investing_focus_sector_classifier (InvestingFocusSectorClassifier): Classifier for investing focus sectors.
     """
-    education_classifier = EducationClassifier()
-    work_experience_classifier = WorkExperienceClassifier()
-    investment_banking_classifier = InvestmentBankingGroupClassifier()
-    investing_focus_asset_class_classifier = InvestingFocusAssetClassClassifier()
-    investing_focus_sector_classifier = InvestingFocusSectorClassifier()
-
-    # Create tasks for education classification
-    education_tasks = {
-        education: education_classifier.classify_education(profile, education) for education in profile.educations
-    }
-
-    # Create tasks for work experience classification
-    work_experience_tasks = {
-        position: work_experience_classifier.classify_work_experience(profile, position)
-        for position in profile.positions
-    }
-
-    # Wait for all education and work experience classifications to complete
-    education_results = await asyncio.gather(*education_tasks.values())
-    work_experience_results = await asyncio.gather(*work_experience_tasks.values())
-
-    # Create ClassifiedEducation objects in the original order
-    classified_educations = [
-        ClassifiedEducation(education=education, classification=classification)
-        for education, classification in zip(profile.educations, education_results)
-    ]
-
-    # Process work experiences and create ClassifiedWorkExperience objects
-    classified_work_experiences = []
-    for position, work_classification in zip(profile.positions, work_experience_results):
-        classified_work_experience = ClassifiedWorkExperience(
-            position=position, work_experience_classification=work_classification
-        )
-
-        if work_classification.primary_job_type not in {PrimaryJobType.INTERNSHIP, PrimaryJobType.EXTRACURRICULAR}:
-            if work_classification.secondary_job_type == SecondaryJobType.INVESTMENT_BANKING:
-                ib_classification = await investment_banking_classifier.classify_investment_banking_group(
-                    profile, position
-                )
-                classified_work_experience.investment_banking_classification = ib_classification
-
-            if work_classification.secondary_job_type == SecondaryJobType.INVESTING:
-                asset_class_task = investing_focus_asset_class_classifier.classify_investing_focus_asset_class(
-                    profile, position
-                )
-                sector_task = investing_focus_sector_classifier.classify_investing_focus_sector(profile, position)
-
-                asset_class_result, sector_result = await asyncio.gather(asset_class_task, sector_task)
-
-                classified_work_experience.investing_focus_asset_class_classification = asset_class_result
-                classified_work_experience.investing_focus_sector_classification = sector_result
 
-        classified_work_experiences.append(classified_work_experience)
-
-    return LinkedinProfileClassificationResults(
-        classified_educations=classified_educations, classified_work_experiences=classified_work_experiences
-    )
+    def __init__(self):
+        """Initialize the VspDataEnrichment class with all required classifiers."""
+        self._education_classifier = EducationClassifier()
+        self._work_experience_classifier = WorkExperienceClassifier()
+        self._investment_banking_classifier = InvestmentBankingGroupClassifier()
+        self._investing_focus_asset_class_classifier = InvestingFocusAssetClassClassifier()
+        self._investing_focus_sector_classifier = InvestingFocusSectorClassifier()
+
+    async def process_linkedin_profile(self, profile: LinkedinProfile) -> LinkedinProfileClassificationResults:
+        """
+        Process a LinkedIn profile and classify its education and work experiences.
+
+        This method maintains the original order of educations and work experiences
+        from the input profile while performing asynchronous classification tasks.
+
+        Args:
+            profile (LinkedinProfile): The LinkedIn profile to process.
+
+        Returns:
+            LinkedinProfileClassificationResults: The comprehensive classification results for the profile.
+        """
+        # Create tasks for education classification
+        education_tasks = {
+            education: self._education_classifier.classify_education(profile, education)
+            for education in profile.educations
+        }
+
+        # Create tasks for work experience classification
+        work_experience_tasks = {
+            position: self._work_experience_classifier.classify_work_experience(profile, position)
+            for position in profile.positions
+        }
+
+        # Wait for all education and work experience classifications to complete
+        education_results = await asyncio.gather(*education_tasks.values())
+        work_experience_results = await asyncio.gather(*work_experience_tasks.values())
+
+        # Create ClassifiedEducation objects in the original order
+        classified_educations = [
+            ClassifiedEducation(education=education, classification=classification)
+            for education, classification in zip(profile.educations, education_results)
+        ]
+
+        # Process work experiences and create ClassifiedWorkExperience objects
+        classified_work_experiences = []
+        for position, work_classification in zip(profile.positions, work_experience_results):
+            classified_work_experience = ClassifiedWorkExperience(
+                position=position, work_experience_classification=work_classification
+            )
+
+            if work_classification.primary_job_type not in {
+                work_classification.primary_job_type.INTERNSHIP,
+                work_classification.primary_job_type.EXTRACURRICULAR,
+            }:
+                if work_classification.secondary_job_type == work_classification.secondary_job_type.INVESTMENT_BANKING:
+                    ib_classification = await self._investment_banking_classifier.classify_investment_banking_group(
+                        profile, position
+                    )
+                    classified_work_experience.investment_banking_classification = ib_classification
+
+                if work_classification.secondary_job_type == work_classification.secondary_job_type.INVESTING:
+                    asset_class_task = self._investing_focus_asset_class_classifier.classify_investing_focus_asset_class(
+                        profile, position
+                    )
+                    sector_task = self._investing_focus_sector_classifier.classify_investing_focus_sector(
+                        profile, position
+                    )
+
+                    asset_class_result, sector_result = await asyncio.gather(asset_class_task, sector_task)
+
+                    classified_work_experience.investing_focus_asset_class_classification = asset_class_result
+                    classified_work_experience.investing_focus_sector_classification = sector_result
+
+            classified_work_experiences.append(classified_work_experience)
+
+        return LinkedinProfileClassificationResults(
+            classified_educations=classified_educations, classified_work_experiences=classified_work_experiences
+        )
 
 
 async def main() -> None:
     """
-    Main function to demonstrate the usage of process_linkedin_profile.
+    Main function to demonstrate the usage of VspDataEnrichment.
+
+    This function loads a sample LinkedIn profile from a JSON file,
+    processes it using the VspDataEnrichment class, and prints the results.
     """
+    import json
+
+    # Load a sample LinkedIn profile
     with open("tests/test_data/sample_profiles/eric_armagost.json") as f:
-        profile = LinkedinProfile.model_validate(json.load(f))
-    results = await process_linkedin_profile(profile)
+        profile_data = json.load(f)
+        profile = LinkedinProfile.model_validate(profile_data)
+
+    # Create an instance of VspDataEnrichment and process the profile
+    vsp_enrichment = VspDataEnrichment()
+    results = await vsp_enrichment.process_linkedin_profile(profile)
+
+    # Print the results
     print(results.model_dump_json(indent=2))