Fix types

Author: donato-ax

src/axiomatic/axtract/axtract_report.py

@@ -1,5 +1,5 @@
 from .relation_graph import generate_relation_graph
-from .. import EquationProcessingResponse
+from ..types.expression_processing_response import ExpressionProcessingResponse
 import os
 import re
 
@@ -190,7 +190,7 @@
 """
 
 
-def create_report(report_data: EquationProcessingResponse, report_path: str = "./report.html"):
+def create_report(report_data: ExpressionProcessingResponse, report_path: str = "./report.html"):
     """
     Creates an HTML report for the extracted equations.
     """

src/axiomatic/axtract/relation_graph.py

@@ -1,5 +1,5 @@
 from typing import List
-from ..types.response_equation import ResponseEquation
+from ..types.expression import Expression
 from pyvis.network import Network  # type: ignore
 
 
@@ -20,7 +20,7 @@ def normalize_latex_symbol(symbol: str) -> str:
     return symbol
 
 
-def generate_relation_graph(equations: List[ResponseEquation]) -> str:
+def generate_relation_graph(equations: List[Expression]) -> str:
     """
     Generates HTML code for a bipartite graph visualization.
     Green nodes represent equations, red nodes represent variables.

src/axiomatic/client.py

@@ -7,11 +7,12 @@
 from typing import Dict, Optional, Sequence
 
 from .base_client import BaseClient, AsyncBaseClient
-from . import ParseResponse, EquationProcessingResponse
 from .axtract.axtract_report import create_report
 from .axtract.validation_results import display_full_results
-from .types.variable_requirement import VariableRequirement
-from .types.equation_validation_result import EquationValidationResult
+from .types.parse_response import ParseResponse
+from .types.expression_processing_response import ExpressionProcessingResponse
+from .types.validation_requirement import ValidationRequirement
+from .types.expression_validation_result import ExpressionValidationResult
 
 
 class Axiomatic(BaseClient):
@@ -147,7 +148,7 @@ def __init__(self, ax_client: Axiomatic):
         """
         self._ax_client = ax_client
 
-    def create_report(self, response: EquationProcessingResponse, path: str):
+    def create_report(self, response: ExpressionProcessingResponse, path: str):
         """Create a report from equation extraction results.
 
         Args:
@@ -161,7 +162,7 @@ def analyze_equations(
         file_path: Optional[str] = None,
         url_path: Optional[str] = None,
         parsed_paper: Optional[ParseResponse] = None,
-    ) -> Optional[EquationProcessingResponse]:
+    ) -> Optional[ExpressionProcessingResponse]:
         """Analyze and extract equations from a scientific document.
 
         This method supports three input methods:
@@ -214,7 +215,7 @@ def analyze_equations(
             )
 
         elif parsed_paper:
-            response = EquationProcessingResponse.model_validate(
+            response = ExpressionProcessingResponse.model_validate(
                 self._ax_client.document.equation.process(**parsed_paper.model_dump()).model_dump()
             )
 
@@ -226,10 +227,10 @@ def analyze_equations(
 
     def validate_equations(
         self,
-        requirements: Sequence[VariableRequirement],
-        loaded_equations: EquationProcessingResponse,
+        requirements: Sequence[ValidationRequirement],
+        loaded_equations: ExpressionProcessingResponse,
         include_internal_model: bool = False,
-    ) -> EquationValidationResult:
+    ) -> ExpressionValidationResult:
         """Validate equations against a set of variable requirements.
 
         Args:
@@ -249,10 +250,10 @@ def validate_equations(
 
         return api_response
 
-    def display_full_results(self, api_response: EquationValidationResult, user_choice):
+    def display_full_results(self, api_response: ExpressionValidationResult, user_choice):
         display_full_results(api_response, user_choice)
 
-    def set_numerical_requirements(self, extracted_equations: EquationProcessingResponse):
+    def set_numerical_requirements(self, extracted_equations: ExpressionProcessingResponse):
         """Launch an interactive interface for setting numerical requirements for equations.
 
         This method opens an interactive table interface where users can specify