feat: Add structured LLM extraction endpoint#97
Conversation
- Adds litellm and instructor dependencies for model integration and structured outputs - Introduces `/extract_structured` API endpoint for handling files and cloud storage payloads - Creates Pydantic models for structured outputs (`StructuredDocumentData`, `StructuredExtractionResult`) - Implements hybrid text extraction pipeline with fallback to local NER using spaCy - Provides tests mocking LLM extraction and fallback behavior - Cleans up legacy Flask and corrupted test files from repository Co-authored-by: gomesrocha <3893269+gomesrocha@users.noreply.github.com>
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
📝 WalkthroughWalkthroughThis PR migrates the application from Flask to FastAPI with modular router architecture, removes Flask-specific configuration and routes, and introduces structured document extraction capabilities with LLM support via instructor and litellm libraries. Legacy OCR scripts are removed, and new Pydantic schemas define extraction outputs. Changes
Sequence DiagramsequenceDiagram
participant Client
participant API as API Endpoint
participant Storage as Object Storage
participant OCR as OCR/PDF Engine
participant NER as NER Module
participant LLM as LLM Provider
participant Response as Response Builder
Client->>API: POST /extract_structured (files + config)
alt source=upload
API->>API: Save uploaded files to temp
else source=object_storage
API->>Storage: Download files by client_id/keys
Storage-->>API: File content
end
loop For each file
alt Detected as PDF
API->>OCR: process_pdf(file_path)
OCR-->>API: extracted_text
else Image file
API->>OCR: read_image(file_path)
OCR-->>API: extracted_text
end
alt provider=local_tesseract
API->>NER: extract_entities(text)
NER-->>API: entities
API->>Response: Build result with local extraction
else LLM provider
API->>LLM: chat.completions.create(text, model, response_model)
LLM-->>API: StructuredDocumentData
API->>Response: Build result with LLM extraction
end
end
API-->>Client: List[StructuredExtractionResult] (200)
API->>API: Clean up temp files
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| if __name__ == '__main__': | ||
| uvicorn.run('main:app', host='127.0.0.1', port=8000, reload=True) No newline at end of file | ||
| if __name__ == "__main__": | ||
| import uvicorn |
| ) | ||
|
|
||
| # Use litellm with instructor to get structured JSON | ||
| log.info(f"Using LLM for structured extraction: provider={provider}, model={model_name}") |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
app/model/StructuredSchema.py (1)
27-32: Use a numeric duration field here.
time_takenasstrpushes parsing work onto every client and loses unit semantics. A numeric field liketime_taken_seconds: floatwill be easier to consume and keeps the contract machine-friendly.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@app/model/StructuredSchema.py` around lines 27 - 32, The StructuredExtractionResult model currently exposes time_taken: str; change it to a numeric seconds field (e.g., time_taken_seconds: float) on the StructuredExtractionResult class so consumers get a machine-friendly duration in seconds; update the field name and type in the model (replace time_taken: str with time_taken_seconds: float), add any necessary Optional/nullable typing if durations can be missing, and search/replace all usages of StructuredExtractionResult.time_taken across the codebase (serializers, constructors, tests, consumers) to convert string durations to the new float seconds value or to accept the new field.tests/test_structured_extract.py (1)
56-85: Add a regression test for the remote-failure branch.This only exercises the happy-path LLM response. It doesn’t cover
instructor/ provider failures, which is where the implementation currently falls back locally and can emit misleading provider metadata. A test withmock_client.chat.completions.create.side_effect = ...would lock that behavior down.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_structured_extract.py` around lines 56 - 85, Add a regression test that simulates a remote LLM failure by configuring the patched instructor client used in test_extract_structured_llm to raise an exception from mock_client.chat.completions.create (e.g., set side_effect to an Exception) so the code exercises the remote-failure branch; then call the same POST to "/extract_structured" and assert the response still succeeds but reports the correct fallback provider/model metadata (verify res_data[0]["provider_used"] / ["model_used"] reflect the local fallback) and that mock_ocr is still called; locate the patch on app.domain.structured_extraction.instructor and the mock_client.chat.completions.create symbol to implement this.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@app/api/structured_extract.py`:
- Around line 132-156: Non-PDF blobs sent to ocr.read_image bypass image
validation; call the existing image validator (fileUpload.validate_image_file or
the corresponding validate function used for uploads) on temp_file in the else
branch before calling ocr.read_image, and if validation fails raise an
HTTPException(400) with a clear message; ensure you catch/propagate validation
exceptions similar to the PDF branch so invalid files are rejected prior to
calling ocr.read_image (refer to temp_file, is_pdf, pdf_ocr.process_pdf,
ocr.read_image and fileUpload.validate_image_file to locate the code).
In `@app/domain/structured_extraction.py`:
- Around line 74-77: The current except block in structured_extraction.py that
catches Exception as e and returns extract_structured_data(text,
"local_tesseract", "", lang_hint) silently hides remote LLM failures; update the
error handling in the except of the LLM extraction path (the block surrounding
extract_structured_data call) to either re-raise the original exception or
return a structured fallback that includes explicit metadata (e.g.,
provider_used="local_tesseract", model_used="" plus an error field or
fallback_flag=true and original_error message) so callers can detect the
downgrade; ensure this change is applied around the function/method performing
LLM structured extraction and that any return type changes are propagated to
callers or documented.
- Around line 45-72: The blocking call to client.chat.completions.create inside
extract_structured_data() (which is invoked from async extract_structured(...))
must be moved off the event loop; wrap the synchronous call to
client.chat.completions.create(..., response_model=StructuredDocumentData, ...)
in a thread executor (e.g. use asyncio.to_thread(...) or
starlette.concurrency.run_in_threadpool(...)) or alternatively convert
extract_structured_data() to async with a non-blocking client; ensure you return
the StructuredDocumentData result from the threaded call and preserve the same
parameters (model=full_model_name, messages, max_tokens) when invoking the
client.
---
Nitpick comments:
In `@app/model/StructuredSchema.py`:
- Around line 27-32: The StructuredExtractionResult model currently exposes
time_taken: str; change it to a numeric seconds field (e.g., time_taken_seconds:
float) on the StructuredExtractionResult class so consumers get a
machine-friendly duration in seconds; update the field name and type in the
model (replace time_taken: str with time_taken_seconds: float), add any
necessary Optional/nullable typing if durations can be missing, and
search/replace all usages of StructuredExtractionResult.time_taken across the
codebase (serializers, constructors, tests, consumers) to convert string
durations to the new float seconds value or to accept the new field.
In `@tests/test_structured_extract.py`:
- Around line 56-85: Add a regression test that simulates a remote LLM failure
by configuring the patched instructor client used in test_extract_structured_llm
to raise an exception from mock_client.chat.completions.create (e.g., set
side_effect to an Exception) so the code exercises the remote-failure branch;
then call the same POST to "/extract_structured" and assert the response still
succeeds but reports the correct fallback provider/model metadata (verify
res_data[0]["provider_used"] / ["model_used"] reflect the local fallback) and
that mock_ocr is still called; locate the patch on
app.domain.structured_extraction.instructor and the
mock_client.chat.completions.create symbol to implement this.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 3cd9f704-a0e1-41b7-b889-f272e35a239b
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (13)
app/__init__.pyapp/api/__init__.pyapp/api/structured_extract.pyapp/config.pyapp/domain/structured_extraction.pyapp/main.pyapp/model/StructuredSchema.pyapp/routes.pypyproject.tomlrun.pysrc/main.pytests/conftest.pytests/test_structured_extract.py
💤 Files with no reviewable changes (6)
- run.py
- app/config.py
- app/init.py
- tests/conftest.py
- src/main.py
- app/routes.py
| if is_pdf: | ||
| try: | ||
| with open(temp_file, "rb") as f: | ||
| header = f.read(5) | ||
| if header != b"%PDF-": | ||
| raise HTTPException(status_code=400, detail="Downloaded file is not a valid PDF.") | ||
| except HTTPException: | ||
| raise | ||
| except Exception as e: | ||
| raise HTTPException(status_code=400, detail=f"Error validating file: {e}") | ||
|
|
||
| raw_text = await pdf_ocr.process_pdf( | ||
| file_path=temp_file, | ||
| lang=lang, | ||
| mode=mode, | ||
| auto=auto_detect, | ||
| force_processing=force_processing | ||
| ) | ||
| else: | ||
| raw_text = await ocr.read_image( | ||
| img_path=temp_file, | ||
| lang=lang, | ||
| mode=mode, | ||
| auto=auto_detect | ||
| ) |
There was a problem hiding this comment.
Validate downloaded non-PDF files before sending them to OCR.
The upload path validates images via fileUpload.validate_image_file(...), but the object-storage path sends any non-.pdf blob straight to ocr.read_image(...). That means a bad key or unexpected file type bypasses the same guardrails as uploads.
🧰 Tools
🪛 Ruff (0.15.9)
[warning] 140-140: Do not catch blind exception: Exception
(BLE001)
[warning] 141-141: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@app/api/structured_extract.py` around lines 132 - 156, Non-PDF blobs sent to
ocr.read_image bypass image validation; call the existing image validator
(fileUpload.validate_image_file or the corresponding validate function used for
uploads) on temp_file in the else branch before calling ocr.read_image, and if
validation fails raise an HTTPException(400) with a clear message; ensure you
catch/propagate validation exceptions similar to the PDF branch so invalid files
are rejected prior to calling ocr.read_image (refer to temp_file, is_pdf,
pdf_ocr.process_pdf, ocr.read_image and fileUpload.validate_image_file to locate
the code).
| client = instructor.from_litellm(completion) | ||
|
|
||
| # Prefix the model with the provider if litellm requires it | ||
| # litellm format is usually "provider/model_name" | ||
| full_model_name = f"{provider}/{model_name}" if "/" not in model_name and provider != "openai" else model_name | ||
|
|
||
| prompt = ( | ||
| "You are an expert data extraction assistant. I will provide you with the raw OCR text of a document. " | ||
| "Your task is to analyze the text and extract structured information, including the type of the document, " | ||
| "any key-value pairs (like invoice numbers, dates, totals, names), any tables you can identify, " | ||
| "and a brief summary. Also extract important named entities.\n\n" | ||
| "Here is the raw document text:\n" | ||
| "------------------\n" | ||
| f"{text}\n" | ||
| "------------------\n" | ||
| ) | ||
|
|
||
| try: | ||
| # instructor will handle forcing the StructuredDocumentData schema | ||
| structured_data = client.chat.completions.create( | ||
| model=full_model_name, | ||
| messages=[ | ||
| {"role": "system", "content": "You are a precise data extraction AI."}, | ||
| {"role": "user", "content": prompt} | ||
| ], | ||
| response_model=StructuredDocumentData, | ||
| max_tokens=4096, | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify that the FastAPI handler is async, calls the helper without `await`,
# and that the helper performs a synchronous Instructor/LiteLLM completion call.
set -euo pipefail
echo "== async endpoint calling structured_extraction.extract_structured_data =="
rg -n -C2 'async def extract_structured|structured_extraction\.extract_structured_data\(' app/api/structured_extract.py
echo
echo "== sync helper using litellm completion + blocking create(...) call =="
rg -n -C2 'def extract_structured_data|from_litellm|client\.chat\.completions\.create' app/domain/structured_extraction.pyRepository: gomesrocha/api_ocr
Length of output: 1561
🏁 Script executed:
# Check if there's any threadpool wrapping around the extract_structured_data calls
rg -n -B5 -A5 'structured_extraction\.extract_structured_data' app/api/structured_extract.py | head -60Repository: gomesrocha/api_ocr
Length of output: 937
🏁 Script executed:
# Check imports in the structured_extract.py to see if asyncio utilities are imported
head -20 app/api/structured_extract.pyRepository: gomesrocha/api_ocr
Length of output: 1260
Don't make a blocking LLM call from the async request path.
extract_structured_data() is a synchronous function that's called directly from async def extract_structured(...) without wrapping in a threadpool executor. The function performs a blocking I/O call with client.chat.completions.create(...) (line 64), which blocks the event loop. Under load, this serializes requests and stalls the entire async handler. Wrap the call with asyncio.to_thread() or starlette.concurrency.run_in_threadpool(), or convert the function to async with an async HTTP client.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@app/domain/structured_extraction.py` around lines 45 - 72, The blocking call
to client.chat.completions.create inside extract_structured_data() (which is
invoked from async extract_structured(...)) must be moved off the event loop;
wrap the synchronous call to client.chat.completions.create(...,
response_model=StructuredDocumentData, ...) in a thread executor (e.g. use
asyncio.to_thread(...) or starlette.concurrency.run_in_threadpool(...)) or
alternatively convert extract_structured_data() to async with a non-blocking
client; ensure you return the StructuredDocumentData result from the threaded
call and preserve the same parameters (model=full_model_name, messages,
max_tokens) when invoking the client.
| except Exception as e: | ||
| log.error(f"Error during LLM structured extraction: {e}") | ||
| # Return fallback on error | ||
| return extract_structured_data(text, "local_tesseract", "", lang_hint) |
There was a problem hiding this comment.
Surface remote-provider failures instead of silently downgrading.
This catch-all fallback returns local data after any LLM error, but the API layer still reports the requested provider_used / model_used. That makes a degraded local result look like a successful remote extraction. Please either propagate an error or return enough metadata for the caller to see that fallback happened.
🧰 Tools
🪛 Ruff (0.15.9)
[warning] 74-74: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@app/domain/structured_extraction.py` around lines 74 - 77, The current except
block in structured_extraction.py that catches Exception as e and returns
extract_structured_data(text, "local_tesseract", "", lang_hint) silently hides
remote LLM failures; update the error handling in the except of the LLM
extraction path (the block surrounding extract_structured_data call) to either
re-raise the original exception or return a structured fallback that includes
explicit metadata (e.g., provider_used="local_tesseract", model_used="" plus an
error field or fallback_flag=true and original_error message) so callers can
detect the downgrade; ensure this change is applied around the function/method
performing LLM structured extraction and that any return type changes are
propagated to callers or documented.
2 participants
This pull request introduces a new
/extract_structuredAPI endpoint which is capable of returning extracted data in a typed and categorized JSON format.The feature uses
litellmandinstructorto leverage remote Models and enforce Pydantic structured schemas. In cases where the configuration does not provide a cloud provider, it gracefully falls back to a fast, heuristic local execution utilizing the previously existing Tesseract OCR and SpaCy NER components, preserving the ability to run multi-cloud without extensive footprint requirements for the Docker execution.Legacy Flask remnants (
app/__init__.py,app/routes.py,src/main.py) which causedpytestloading errors were cleaned up and thetestssuite executed completely cleanly.PR created automatically by Jules for task 13248805800921705453 started by @gomesrocha
Summary by CodeRabbit
Release Notes
New Features
Chores