add tests

.pre-commit-config.yaml

@@ -21,14 +21,14 @@ repos:
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
-        exclude: ".*poetry.lock"
+        exclude: "poetry.lock"
   - repo: https://github.com/psf/black-pre-commit-mirror
     rev: 23.10.1
     hooks:
       - id: black-jupyter
         name: black-src
         alias: black
-        exclude: ".*poetry.lock"
+        exclude: "poetry.lock"
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.0.1
     hooks:

llama_extract/base.py

@@ -7,7 +7,7 @@
 from io import BufferedIOBase, BufferedReader, BytesIO
 from json.decoder import JSONDecodeError
 from pathlib import Path
-from pydantic import BaseModel, Extra, ValidationError
+from pydantic import BaseModel, ValidationError
 from typing import List, Optional, Tuple, Type, Union
 import urllib.parse
 
@@ -26,7 +26,7 @@
 from llama_extract.utils import nest_asyncio_err, nest_asyncio_msg
 from llama_index.core.schema import BaseComponent
 from llama_index.core.async_utils import asyncio_run, run_jobs
-from llama_index.core.bridge.pydantic import Field, pydantic, PrivateAttr
+from llama_index.core.bridge.pydantic import Field, PrivateAttr
 from llama_index.core.constants import DEFAULT_BASE_URL
 
 # can put in a path to the file or the file bytes itself
@@ -224,7 +224,7 @@ async def ainfer_schema(
             _response_json = _response.json()
         except JSONDecodeError:
             raise ApiError(status_code=_response.status_code, body=_response.text)
-        
+
         raise ApiError(status_code=_response.status_code, body=_response_json)
 
     def infer_schema(