Enable new processing modules

bc2/core/analyze/__init__.py

@@ -0,0 +1,5 @@
+from typing import Union
+
+from .azuredi import AzureDIAnalyzeConfig
+
+AnalyzeConfig = Union[AzureDIAnalyzeConfig,]

bc2/core/analyze/azuredi.py

@@ -19,7 +19,7 @@
 class AzureDIAnalyzeConfig(BaseModel):
     """Azure DI Analyze config."""
 
-    engine: Literal["analyze:azuredi"]
+    engine: Literal["analyze:azuredi"] = "analyze:azuredi"
     endpoint: str
     api_key: str
     # Todo: Add api_version, since we'll need to match what's on GovCloud,

bc2/core/common/all.py

@@ -1,9 +1,12 @@
 import typing
 
+from ..analyze import AnalyzeConfig
 from ..extract import ExtractConfig
 from ..input import InputConfig
 from ..inspect import InspectConfig
+from ..ontology import OntologyConfig
 from ..output import OutputConfig
+from ..paint import PaintConfig
 from ..parse import ParseConfig
 from ..redact import RedactConfig
 from ..render import RenderConfig
@@ -14,7 +17,10 @@
 ]
 
 AnyProcessingConfig = typing.Union[
+    AnalyzeConfig,
     ExtractConfig,
+    OntologyConfig,
+    PaintConfig,
     RedactConfig,
     InspectConfig,
     ParseConfig,

bc2/core/ontology/__init__.py

@@ -0,0 +1,5 @@
+from typing import Union
+
+from .openai import OpenAIOntologyConfig
+
+OntologyConfig = Union[OpenAIOntologyConfig,]

bc2/core/ontology/openai.py

@@ -24,7 +24,7 @@
 class OpenAIOntologyConfig(OpenAIConfig):
     """OpenAI Ontology config."""
 
-    engine: Literal["ontology:openai"]
+    engine: Literal["ontology:openai"] = "ontology:openai"
     generator: OpenAIChatConfig[PoliceReport]
 
     @cached_property

bc2/core/paint/__init__.py

@@ -0,0 +1,5 @@
+from typing import Union
+
+from .ontology import OntologyPainterConfig
+
+PaintConfig = Union[OntologyPainterConfig,]