Josh/clean up opensearch classes (#381)

## Description

Remove some duplicated models being used in lambda_handler.

## Related Issues

Closes n/a

Author: JNygaard-Skylight

packages/lambda-handler/src/lambda_handler/models/opensearch.py

@@ -1,77 +1,66 @@
-import pydantic
+from pydantic import BaseModel
+from pydantic import Field
 
 
-class S3Location(pydantic.BaseModel):
+class S3Location(BaseModel):
     """Represents the location of a file in S3, indicating which file contained the relevant data."""
 
-    bucket: str = pydantic.Field(description="The S3 bucket where the file is located.")
-    key: str = pydantic.Field(description="The S3 key (path) where the file is located.")
+    bucket: str = Field(description="The S3 bucket where the file is located.")
+    key: str = Field(description="The S3 key (path) where the file is located.")
 
 
-class OpenSearchHitSource(pydantic.BaseModel):
+class OpenSearchHitSource(BaseModel):
     """Represents a single search result _source returned from OpenSearch."""
 
-    id: int = pydantic.Field(
-        description="The unique ID from the embedding data of the search result hit."
-    )
-    loinc_code: str = pydantic.Field(description="The LOINC code of the search result hit.")
-    loinc_name_type: str = pydantic.Field(
-        description="The LOINC name type of the search result hit."
-    )
-    description: str = pydantic.Field(description="The description of the search result hit.")
-    loinc_type: str = pydantic.Field(description="The LOINC type of the search result hit.")
-    s3: S3Location = pydantic.Field(description="The S3 location of the search result hit.")
+    id: int = Field(description="The unique ID from the embedding data of the search result hit.")
+    loinc_code: str = Field(description="The LOINC code of the search result hit.")
+    loinc_name_type: str = Field(description="The LOINC name type of the search result hit.")
+    description: str = Field(description="The description of the search result hit.")
+    loinc_type: str = Field(description="The LOINC type of the search result hit.")
+    s3: S3Location = Field(description="The S3 location of the search result hit.")
 
 
-class OpenSearchHit(pydantic.BaseModel):
+class OpenSearchHit(BaseModel):
     """Represents a single search result hit returned from OpenSearch."""
 
-    index: str = pydantic.Field(
+    index: str = Field(
         description="The index that the search result hit came from.", alias="_index"
     )
-    id: str = pydantic.Field(
-        description="The unique OpenSearch ID of the search result hit.", alias="_id"
-    )
-    score: float = pydantic.Field(
+    id: str = Field(description="The unique OpenSearch ID of the search result hit.", alias="_id")
+    score: float = Field(
         description="The cosine similarity score of the search result hit.", alias="_score"
     )
-    source: OpenSearchHitSource = pydantic.Field(
+    source: OpenSearchHitSource = Field(
         description="The source of the search result hit.", alias="_source"
     )
 
 
-class OpenSearchHits(pydantic.BaseModel):
+class OpenSearchHits(BaseModel):
     """Represents all of the search result hits returned from OpenSearch."""
 
-    total_hits: dict[str, int] = pydantic.Field(
+    total_hits: dict[str, int] = Field(
         alias="total", description="The total number of hits returned from OpenSearch."
     )
-    hits: list[OpenSearchHit] = pydantic.Field(
+    hits: list[OpenSearchHit] = Field(
         description="The list of search result hits returned from OpenSearch."
     )
 
 
-class OpenSearchShards(pydantic.BaseModel):
+class OpenSearchShards(BaseModel):
     """Represents the shard information returned from OpenSearch."""
 
-    total: int = pydantic.Field(description="The total number of shards involved in the search.")
-    successful: int = pydantic.Field(
-        description="The number of shards that successfully returned results."
-    )
-    skipped: int = pydantic.Field(
-        description="The number of shards that were skipped during the search."
-    )
-    failed: int = pydantic.Field(description="The number of shards that failed to return results.")
+    total: int = Field(description="The total number of shards involved in the search.")
+    successful: int = Field(description="The number of shards that successfully returned results.")
+    skipped: int = Field(description="The number of shards that were skipped during the search.")
+    failed: int = Field(description="The number of shards that failed to return results.")
 
 
-class OpenSearchResult(pydantic.BaseModel):
+class OpenSearchResult(BaseModel):
     """Represents the overall search result returned from OpenSearch, including hits and shard information."""
 
-    took: int = pydantic.Field(description="The time taken to execute the search in milliseconds.")
-    timed_out: bool = pydantic.Field(description="Indicates whether the search timed out.")
-    shards: OpenSearchShards = pydantic.Field(
+    took: int = Field(description="The time taken to execute the search in milliseconds.")
+    timed_out: bool = Field(description="Indicates whether the search timed out.")
+    shards: OpenSearchShards = Field(
         description="The shard information for the search.", alias="_shards"
     )
-    hits: OpenSearchHits = pydantic.Field(
-        description="The search result hits returned from OpenSearch."
-    )
+    hits: OpenSearchHits = Field(description="The search result hits returned from OpenSearch.")

packages/text-to-code/src/text_to_code/models/query.py

@@ -1,6 +1,8 @@
 import typing
 
-import pydantic
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import model_validator
 
 from shared_models import DataField
 
@@ -22,110 +24,32 @@ def to_filter_values(cls, data_field: DataField) -> list[str]:
             raise ValueError(f"No type mapping defined for {data_field}") from err
 
 
-class VectorSearchParams(pydantic.BaseModel):
+class VectorSearchParams(BaseModel):
     """Parameters for performing a vector search."""
 
-    vector: list[float] = pydantic.Field(description="The vector to search for.")
-    vector_field: str = pydantic.Field(
+    vector: list[float] = Field(description="The vector to search for.")
+    vector_field: str = Field(
         default="descriptionVector", description="The field to perform the vector search on."
     )
-    filter_field: str = pydantic.Field(
-        default="type", description="The field to filter on, e.g., 'type'."
-    )
-    data_field: DataField = pydantic.Field(
+    filter_field: str = Field(default="type", description="The field to filter on, e.g., 'type'.")
+    data_field: DataField = Field(
         description="The value of the field to filter on, e.g., 'Lab Test Name Ordered' or 'Lab Test Name Resulted'."
     )
-    size: int = pydantic.Field(default=10, description="The number of results to retrieve.")
-    k: int = pydantic.Field(
+    size: int = Field(default=10, description="The number of results to retrieve.")
+    k: int = Field(
         default=10, description="The number of nearest neighbors to examine during the query."
     )
-    filter_value: list[str] = pydantic.Field(
+    filter_value: list[str] = Field(
         default_factory=list,
         init=False,
         description="The list of filter values corresponding to the data_field, computed after initialization.",
     )
 
-    @pydantic.model_validator(mode="after")
+    @model_validator(mode="after")
     def compute_filter_value(self) -> "VectorSearchParams":
         """Uses the DataFieldTypeMapping to get the filter values corresponding to the data_field."""
         if self.filter_field == type(self).model_fields["filter_field"].default:
             self.filter_value = DataFieldTypeMapping.to_filter_values(self.data_field)
         else:
             raise ValueError(f"Unsupported filter field: {self.filter_field}")
         return self
-
-
-class S3Location(pydantic.BaseModel):
-    """Represents the location of a file in S3, indicating which file contained the relevant data."""
-
-    bucket: str = pydantic.Field(description="The S3 bucket where the file is located.")
-    key: str = pydantic.Field(description="The S3 key (path) where the file is located.")
-
-
-class OpenSearchHitSource(pydantic.BaseModel):
-    """Represents a single search result _source returned from OpenSearch."""
-
-    id: int = pydantic.Field(
-        description="The unique ID from the embedding data of the search result hit."
-    )
-    loinc_code: str = pydantic.Field(description="The LOINC code of the search result hit.")
-    loinc_name_type: str = pydantic.Field(
-        description="The LOINC name type of the search result hit."
-    )
-    description: str = pydantic.Field(description="The description of the search result hit.")
-    loinc_type: str = pydantic.Field(description="The LOINC type of the search result hit.")
-    s3: S3Location = pydantic.Field(description="The S3 location of the search result hit.")
-
-
-class OpenSearchHit(pydantic.BaseModel):
-    """Represents a single search result hit returned from OpenSearch."""
-
-    index: str = pydantic.Field(
-        description="The index that the search result hit came from.", alias="_index"
-    )
-    id: str = pydantic.Field(
-        description="The unique OpenSearch ID of the search result hit.", alias="_id"
-    )
-    score: float = pydantic.Field(
-        description="The cosine similarity score of the search result hit.", alias="_score"
-    )
-    source: OpenSearchHitSource = pydantic.Field(
-        description="The source of the search result hit.", alias="_source"
-    )
-
-
-class OpenSearchHits(pydantic.BaseModel):
-    """Represents all of the search result hits returned from OpenSearch."""
-
-    total_hits: dict[str, int] = pydantic.Field(
-        alias="total", description="The total number of hits returned from OpenSearch."
-    )
-    hits: list[OpenSearchHit] = pydantic.Field(
-        description="The list of search result hits returned from OpenSearch."
-    )
-
-
-class OpenSearchShards(pydantic.BaseModel):
-    """Represents the shard information returned from OpenSearch."""
-
-    total: int = pydantic.Field(description="The total number of shards involved in the search.")
-    successful: int = pydantic.Field(
-        description="The number of shards that successfully returned results."
-    )
-    skipped: int = pydantic.Field(
-        description="The number of shards that were skipped during the search."
-    )
-    failed: int = pydantic.Field(description="The number of shards that failed to return results.")
-
-
-class OpenSearchResult(pydantic.BaseModel):
-    """Represents the overall search result returned from OpenSearch, including hits and shard information."""
-
-    took: int = pydantic.Field(description="The time taken to execute the search in milliseconds.")
-    timed_out: bool = pydantic.Field(description="Indicates whether the search timed out.")
-    shards: OpenSearchShards = pydantic.Field(
-        description="The shard information for the search.", alias="_shards"
-    )
-    hits: OpenSearchHits = pydantic.Field(
-        description="The search result hits returned from OpenSearch."
-    )