File size: 3,847 Bytes
18e58fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | import os
import sys
from unstructured.partition.pdf import partition_pdf
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from embeddings.embedder import upsert_document
import httpx
def is_api_running() -> bool:
try:
r = httpx.get("http://127.0.0.1:8000/health", timeout=1.0)
return r.status_code == 200
except Exception:
return False
def _local_ingest_pdf(filepath: str, metadata: dict, thorough: bool = True):
print(f"Starting advanced partitioning for {filepath}...")
# Partitioning logic
elements = partition_pdf(
filename=filepath,
strategy="hi_res" if thorough else "fast",
)
chunks = []
current_chunk = ""
for el in elements:
if hasattr(el, 'text'):
if len(current_chunk) + len(el.text) > 2000: # Max chunk size
chunks.append(current_chunk)
current_chunk = el.text
else:
current_chunk += "\n" + el.text
if current_chunk:
chunks.append(current_chunk)
print(f"Ingesting {filepath}: {len(chunks)} high-fidelity chunks")
for i, chunk in enumerate(chunks):
chunk_meta = {**metadata, "chunk_index": i, "total_chunks": len(chunks), "method": "unstructured"}
upsert_document(chunk, chunk_meta)
print(f"Done: {filepath}")
def ingest_pdf(filepath: str, metadata: dict, thorough: bool = True):
"""
Advanced ingestion using unstructured.io for high-fidelity partitioning.
Routes to API if server is active to prevent lock conflict and double models in memory.
"""
if is_api_running():
print(f"API server is active. Routing PDF ingestion for {filepath} through endpoint...")
try:
r = httpx.post("http://127.0.0.1:8000/api/ingest/pdf", json={
"filepath": os.path.abspath(filepath),
"metadata": metadata,
"thorough": thorough
}, timeout=180.0)
r.raise_for_status()
print(f"API successfully ingested {filepath}")
return
except Exception as e:
print(f"API Ingestion failed, falling back to local database write: {e}")
_local_ingest_pdf(filepath, metadata, thorough)
def _local_ingest_text(text: str, metadata: dict):
# Basic chunking for plain text
words = text.split()
chunks = []
chunk_size = 500
overlap = 50
i = 0
while i < len(words):
chunk = " ".join(words[i:i+chunk_size])
chunks.append(chunk)
i += chunk_size - overlap
for i, chunk in enumerate(chunks):
if len(chunk.strip()) > 100:
chunk_meta = {**metadata, "chunk_index": i, "total_chunks": len(chunks)}
upsert_document(chunk, chunk_meta)
def ingest_text(text: str, metadata: dict):
if is_api_running():
try:
r = httpx.post("http://127.0.0.1:8000/api/ingest", json={
"text": text,
"metadata": metadata
}, timeout=300.0)
r.raise_for_status()
return
except Exception as e:
print(f"API Ingestion failed, falling back to local database write: {e}")
_local_ingest_text(text, metadata)
if __name__ == "__main__":
# Test ingestion with a sample text
sample = """
The UAE Federal Tax Authority (FTA) administers Value Added Tax (VAT)
at a standard rate of 5% on most goods and services. Corporate Tax was
introduced in June 2023 at 9% on taxable income exceeding AED 375,000.
"""
ingest_text(sample, {
"source": "sample_v2",
"source_type": "statute",
"jurisdiction": "UAE",
"doc_title": "Tax Overview UAE 2024",
"date": "2024-01-01",
"url": ""
})
print("Sample V2 ingestion complete.") |