Merge pull request #36 from run-llama/logan/fix_everything

.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

@@ -1,6 +1,9 @@
 import asyncio
 import os
 import time
+
+import pydantic.v1 as pydantic_v1
+
 from io import BufferedIOBase, BufferedReader, BytesIO
 from json.decoder import JSONDecodeError
 from pathlib import Path
@@ -23,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
@@ -212,15 +215,16 @@ async def ainfer_schema(
         )
 
         if 200 <= _response.status_code < 300:
-            return pydantic.parse_obj_as(ExtractionSchema, _response.json())  # type: ignore
+            return pydantic_v1.parse_obj_as(ExtractionSchema, _response.json())
         if _response.status_code == 422:
             raise UnprocessableEntityError(
-                pydantic.parse_obj_as(HttpValidationError, _response.json())
-            )  # type: ignore
+                pydantic_v1.parse_obj_as(HttpValidationError, _response.json())
+            )
         try:
             _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(