Skip to content

Commit

Permalink
TLDR-844 add auth key to grobid (#503)
Browse files Browse the repository at this point in the history
  • Loading branch information
NastyBoget authored and sunveil committed Nov 6, 2024
1 parent d10f023 commit aff3468
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions dedoc/readers/article_reader/article_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,9 @@ def __init__(self, config: Optional[dict] = None) -> None:
else:
self.grobid_url = f"http://{os.environ.get('GROBID_HOST', 'localhost')}:{os.environ.get('GROBID_PORT', '8070')}"
self.url = f"{self.grobid_url}/api/processFulltextDocument"

auth_key = os.environ.get("GROBID_AUTH_KEY", "")
self.request_headers = {"Authorization": auth_key} if auth_key else {}
self.grobid_is_alive = False

def read(self, file_path: str, parameters: Optional[dict] = None) -> UnstructuredDocument:
Expand All @@ -48,7 +51,7 @@ def read(self, file_path: str, parameters: Optional[dict] = None) -> Unstructure
with open(file_path, "rb") as file:
files = {"input": file}
try:
response = requests.post(self.url, files=files, data={"teiCoordinates": "figure"})
response = requests.post(self.url, files=files, data={"teiCoordinates": "figure"}, headers=self.request_headers)
if response.status_code != 200:
warning = f"GROBID returns code {response.status_code}."
self.logger.warning(warning)
Expand Down Expand Up @@ -106,7 +109,7 @@ def __update_grobid_alive(self, grobid_url: str, max_attempts: int = 2) -> None:
attempt = max_attempts
while attempt > 0:
try:
response = requests.get(f"{grobid_url}/api/isalive")
response = requests.get(f"{grobid_url}/api/isalive", headers=self.request_headers)
if response.status_code == 200:
self.logger.info(f"GROBID up on {grobid_url}.")
self.grobid_is_alive = True
Expand Down

0 comments on commit aff3468

Please sign in to comment.