fix axtract typing

Author: lsarra-ax

src/axiomatic/axtract/axtract_report.py

@@ -1,5 +1,5 @@
 from .relation_graph import generate_relation_graph
-from .models import EquationExtractionResponse
+from .. import EquationExtractionResponse
 import os
 import re
 

src/axiomatic/axtract/interactive_table.py

@@ -2,8 +2,7 @@
 from IPython.display import display  # type: ignore
 import json  # type: ignore
 import os  # type: ignore
-from .axtract_report import EquationExtractionResponse
-from .models import VariableRequirement
+from .. import EquationExtractionResponse, VariableRequirement
 
 
 def _find_symbol(name, variable_dict):

src/axiomatic/axtract/models.py

@@ -1,5 +1,5 @@
 from typing import List
-from .models import EquationExtraction
+from .. import EquationExtraction
 from pyvis.network import Network  # type: ignore
 
 

src/axiomatic/axtract/relation_graph.py

@@ -2,12 +2,11 @@
 import requests
 import os
 import time
+import json
 from typing import Dict, List, Union, Optional
 
 from .base_client import BaseClient, AsyncBaseClient
-from . import ParseResponse
-from .axtract.models import EquationExtractionResponse
-
+from . import ParseResponse, EquationProcessingResponse, EquationExtractionResponse
 
 class Axiomatic(BaseClient):
     def __init__(self, *args, **kwargs):
@@ -17,6 +16,7 @@ def __init__(self, *args, **kwargs):
 
         self.document_helper = DocumentHelper(self)
         self.tools_helper = ToolsHelper(self)
+        self.axtract_helper = AxtractHelper(self)
 
 
 class AxtractHelper:
@@ -38,6 +38,7 @@ def analyze_equations(
         url_path: Optional[str] = None,
         parsed_paper: Optional[ParseResponse] = None,
     ) -> Optional[EquationExtractionResponse]:
+        response : EquationExtractionResponse | EquationProcessingResponse
         if file_path:
             with open(file_path, "rb") as file:
                 response = self._ax_client.document.equation.from_pdf(document=file)
@@ -47,7 +48,7 @@ def analyze_equations(
 
             response = self._ax_client.document.equation.from_pdf(document=url_path)
         elif parsed_paper:
-            response = self._ax_client.document.equation.process(**parsed_paper)
+            response = self._ax_client.document.equation.process(**parsed_paper.model_dump())
         else:
             print("Please provide either a file path or a URL to analyze.")
             return None
@@ -135,13 +136,19 @@ def navigate_image(change):
     def save_parsed_pdf(self, response: ParseResponse, path: str):
         """Save a parsed PDF response to a file."""
         os.makedirs(path, exist_ok=True)
+        with open(os.path.join(path, "text.md"), "w") as f:
+            f.write(response.markdown)
+
         if response.images:
             for img_name, img in response.images.items():
                 with open(os.path.join(path, f"{img_name}.png"), "wb") as f:
                     f.write(base64.b64decode(img))
 
-        with open(os.path.join(path, "text.md"), "w") as f:
-            f.write(response.markdown)
+        with open(os.path.join(path, "interline_equations.json"), "w") as f:
+            json.dump(response.interline_equations, f)
+
+        with open(os.path.join(path, "inline_equations.json"), "w") as f:
+            json.dump(response.inline_equations, f)
 
     def load_parsed_pdf(self, path: str) -> ParseResponse:
         """Load a parsed PDF response from a file."""
@@ -154,7 +161,18 @@ def load_parsed_pdf(self, path: str) -> ParseResponse:
                 with open(os.path.join(path, img_name), "rb") as img_file:
                     images[img_name] = base64.b64encode(img_file.read()).decode("utf-8")
 
-        return ParseResponse(markdown=markdown, images=images)
+        with open(os.path.join(path, "interline_equations.json"), "r") as f:
+            interline_equations = json.load(f)
+
+        with open(os.path.join(path, "inline_equations.json"), "r") as f:
+            inline_equations = json.load(f)
+
+        return ParseResponse(
+            markdown=markdown,
+            images=images,
+            interline_equations=interline_equations,
+            inline_equations=inline_equations,
+        )
 
 
 class ToolsHelper: