WIP AF

Author: tymekaxai

client_db.pkl

@@ -2,7 +2,7 @@
 from IPython.display import display  # type: ignore
 import json  # type: ignore
 import os  # type: ignore
-from .. import EquationExtractionResponse, VariableRequirement
+from .. import EquationProcessingResponse, VariableRequirement
 
 
 def _find_symbol(name, variable_dict):
@@ -33,15 +33,15 @@ def _requirements_from_table(results, variable_dict):
     return requirements
 
 
-def interactive_table(loaded_equations, file_path="./custom_presets.json"):
+def interactive_table(loaded_equations: EquationProcessingResponse, file_path: str = "./custom_presets.json"):
     """
     Creates an interactive table for IMAGING_TELESCOPE,
     PAYLOAD, and user-defined custom templates.
     Adds or deletes rows, and can save custom templates persistently in JSON.
 
     Parameters
     ----------
-    loaded_equations : EquationExtractionResponse
+    loaded_equations : EquationProcessingResponse
         The extracted equations containing variable information
     file_path : str, optional
         JSON file path where we load and save user-created custom templates.
@@ -382,14 +382,14 @@ def save_requirements(_):
     return requirements_result
 
 
-def _create_variable_dict(equation_response: EquationExtractionResponse) -> dict:
+def _create_variable_dict(equation_response: EquationProcessingResponse) -> dict:
     """
-    Creates a variable dictionary from an EquationExtractionResponse object
+    Creates a variable dictionary from an EquationProcessingResponse object
     for use with the interactive_table function.
 
     Parameters
     ----------
-    equation_response : EquationExtractionResponse
+    equation_response : EquationProcessingResponse
         The equation extraction response containing equations and their symbols
 
     Returns

src/axiomatic/axtract/interactive_table.py

@@ -8,11 +8,12 @@
 from typing import Dict, List, Optional, Union
 
 from .base_client import BaseClient, AsyncBaseClient
-from . import ParseResponse, EquationExtractionResponse, EquationProcessingResponse
+from . import ParseResponse, EquationProcessingResponse
 from .axtract.axtract_report import create_report
 from .axtract.validation_results import display_full_results
-from .axtract.interactive_table import _create_variable_dict
-from .types.variable_requirement import VariableRequirement as ApiVariableRequirement
+from .axtract.interactive_table import _create_variable_dict, VariableRequirement
+from .types.variable_requirement import VariableRequirement
+from .types.validate_equations_body import ValidateEquationsBody
 
 
 class Axiomatic(BaseClient):
@@ -27,36 +28,83 @@ def __init__(self, *args, **kwargs):
 
 
 class AxtractHelper:
-    from .axtract.interactive_table import VariableRequirement
-
     _ax_client: Axiomatic
 
     def __init__(self, ax_client: Axiomatic):
+        """Initialize the AxtractHelper with an Axiomatic client.
+
+        Args:
+            ax_client (Axiomatic): The Axiomatic client instance to use for API calls
+        """
         self._ax_client = ax_client
 
-    def create_report(self, response: EquationExtractionResponse, path: str):
+    def create_report(self, response: EquationProcessingResponse, path: str):
+        """Create a report from equation extraction results.
+
+        Args:
+            response (EquationExtractionResponse): The extracted equations and their metadata
+            path (str): Directory path where the report should be saved
+        """
         create_report(response, path)
 
     def analyze_equations(
         self,
         file_path: Optional[str] = None,
         url_path: Optional[str] = None,
         parsed_paper: Optional[ParseResponse] = None,
-    ) -> Optional[EquationExtractionResponse]:        
+    ) -> Optional[EquationProcessingResponse]:
+        """Analyze and extract equations from a scientific document.
+
+        This method supports three input methods:
+        1. Local PDF file path
+        2. URL to a PDF (with special handling for arXiv URLs)
+        3. Pre-parsed paper data
+
+        Args:
+            file_path (Optional[str]): Path to a local PDF file
+            url_path (Optional[str]): URL to a PDF file (supports arXiv links)
+            parsed_paper (Optional[ParseResponse]): Pre-parsed paper data
+
+        Returns:
+            Optional[EquationExtractionResponse]: Extracted equations and their metadata.
+            Returns None if no valid input is provided.
+
+        Examples:
+            # From local file
+            client.analyze_equations(file_path="path/to/paper.pdf")
+            
+            # From URL
+            client.analyze_equations(url_path="https://arxiv.org/pdf/2203.00001.pdf")
+            
+            # From parsed paper
+            client.analyze_equations(parsed_paper=parsed_data)
+        """
         if file_path:
             with open(file_path, "rb") as pdf_file:
-                response = self._ax_client.document.equation.from_pdf(document=pdf_file)
+                parsed_document = self._ax_client.document.parse(file=pdf_file)
+                print("We are almost there")
+                response = self._ax_client.document.equation.process(
+                    markdown=parsed_document.markdown, 
+                    interline_equations=parsed_document.interline_equations,
+                    inline_equations=parsed_document.inline_equations
+                    )
 
         elif url_path:
             if "arxiv" in url_path and "abs" in url_path:
                 url_path = url_path.replace("abs", "pdf")
             url_file = requests.get(url_path)
             from io import BytesIO
             pdf_stream = BytesIO(url_file.content)
-            response = self._ax_client.document.equation.from_pdf(document=pdf_stream)
+            parsed_document = self._ax_client.document.parse(file=url_file.content)
+            print("We are almost there")
+            response = self._ax_client.document.equation.process(
+                markdown=parsed_document.markdown, 
+                interline_equations=parsed_document.interline_equations,
+                inline_equations=parsed_document.inline_equations
+                )
 
         elif parsed_paper:
-            response = EquationExtractionResponse.model_validate(
+            response = EquationProcessingResponse.model_validate(
                 self._ax_client.document.equation.process(**parsed_paper.model_dump()).model_dump()
             )
 
@@ -68,22 +116,55 @@ def analyze_equations(
 
     def validate_equations(
         self,
-        requirements: List[VariableRequirement],
-        loaded_equations: EquationExtractionResponse,
+        requirements: list[VariableRequirement],
+        loaded_equations: EquationProcessingResponse,
         show_hypergraph: bool = True,
+        include_internal_model: bool = False,
     ):
-        api_requirements = [
-            ApiVariableRequirement(
-                symbol=req.symbol, name=req.name, value=req.value, units=req.units, tolerance=req.tolerance
-            )
-            for req in requirements
-        ]
+        """Validate equations against a set of variable requirements.
+
+        Args:
+            requirements: List of variable requirements to validate
+            loaded_equations: Previously processed equations to validate
+            show_hypergraph: Whether to display the validation results graph (default: True)
+            include_internal_model: Whether to include internal model equations in validation (default: False)
+
+        Returns:
+            EquationValidationResult containing the validation results
+        """
+        # Convert loaded_equations to dict first to ensure proper serialization
+        equations_dict = loaded_equations.model_dump() if hasattr(loaded_equations, 'model_dump') else loaded_equations.dict()
+        
+        request_body = ValidateEquationsBody(
+            variables=requirements,
+            paper_equations=equations_dict,
+            include_internal_model=include_internal_model
+        )
+        
+        print(request_body.model_dump_json())
+
+        api_response = self._ax_client.document.equation.validate(request_body.model_dump_json())
+        
+        if show_hypergraph:
+            pass
+        
+        return api_response
+    
+
+    
+
+    def set_numerical_requirements(self, extracted_equations: EquationProcessingResponse):
+        """Launch an interactive interface for setting numerical requirements for equations.
 
-        variable_dict = _create_variable_dict(loaded_equations)
-        api_response = self._ax_client.document.equation.validate(request=api_requirements)
-        display_full_results(api_response.model_dump(), variable_dict, show_hypergraph=show_hypergraph)
+        This method opens an interactive table interface where users can specify
+        requirements for variables found in the extracted equations.
 
-    def set_numerical_requirements(self, extracted_equations):
+        Args:
+            extracted_equations: The equations to set requirements for
+
+        Returns:
+            The requirements set through the interactive interface
+        """
         from .axtract.interactive_table import interactive_table
 
         result = interactive_table(extracted_equations)
@@ -255,4 +336,4 @@ def _load_objects_from_base64(self, encoded_dict):
         return loaded_objects
 
 
-class AsyncAxiomatic(AsyncBaseClient): ...
+class AsyncAxiomatic(AsyncBaseClient): ...

src/axiomatic/client.py

@@ -0,0 +1,38 @@
+from pydantic import BaseModel, Field
+from typing import List, Dict, Optional
+
+class DictItem(BaseModel):
+    key: str
+    value: str
+
+class ResponseEquation(BaseModel):
+    """
+    Pydantic model representing a scientific or mathematical equation.
+    """
+
+    name: str = Field(
+        ...,
+        description="Human-readable name of the equation (e.g., 'Newton's Law')",
+    )
+    description: str = Field(..., description="Text describing the equation")
+    original_format: str = Field(..., description="Original form of the equation (e.g., LaTeX)")
+    wolfram_expressions: str = Field(
+        ..., description="Equation in Wolfram Language (or other symbolic form)"
+    )
+    latex_symbols: list[DictItem] = Field(
+        ..., description="Detailed metadata for each latex variable in original format"
+    )
+    wolfram_symbols: list[str] = Field(
+        ..., description="List of symbols used in the wolfram expression"
+    )
+    narrative_assumptions: list[str] = Field(
+        ..., description="Narrative text describing assumptions/approximations"
+    )
+    type: list[str] = Field(
+        ...,
+        description="List of equation classifications (e.g., ['scalar','tensor'])",
+    )
+    field_tags: list[str] = Field(
+        ...,
+        description="Classification tags (e.g., ['quantum mechanics','gravitation'])",
+    )

src/axiomatic/types/equation_processing.py

@@ -2,13 +2,13 @@
 
 from ..core.pydantic_utilities import UniversalBaseModel
 import typing
-from .equation_extraction import EquationExtraction
+from .equation_processing import ResponseEquation
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import pydantic
 
 
 class EquationProcessingResponse(UniversalBaseModel):
-    equations: typing.List[EquationExtraction]
+    equations: typing.List[ResponseEquation]
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2

src/axiomatic/types/equation_processing_response.py

@@ -0,0 +1,8 @@
+from .variable_requirement import VariableRequirement
+from .equation_processing_response import EquationProcessingResponse
+from pydantic import BaseModel
+
+class ValidateEquationsBody(BaseModel):
+    variables: list[VariableRequirement]
+    paper_equations: EquationProcessingResponse
+    include_internal_model: bool = False