wehub-resource-sync
160 lines
6.1 KiB
Python
160 lines
6.1 KiB
Python
import os
|
||
import tempfile
|
||
import time
|
||
import requests
|
||
from pathlib import Path
|
||
from typing import Any, Dict, Optional
|
||
|
||
import streamlit as st
|
||
from groundx import GroundX, Document
|
||
from dotenv import load_dotenv
|
||
|
||
# Load environment configuration
|
||
load_dotenv()
|
||
API_KEY: Optional[str] = os.getenv("GROUNDX_API_KEY")
|
||
|
||
# Ground X API Utility Functions
|
||
@st.cache_resource(show_spinner=False)
|
||
def create_client() -> GroundX:
|
||
"""Initialize and return a Ground X API client instance"""
|
||
if not API_KEY:
|
||
raise ValueError("No `GROUNDX_API_KEY` found in secrets, .env, or environment.")
|
||
return GroundX(api_key=API_KEY)
|
||
|
||
@st.cache_resource(show_spinner=False)
|
||
def ensure_bucket(_gx: GroundX, name: str = "gx_demo") -> str:
|
||
"""Ensure storage bucket exists and return its identifier"""
|
||
buckets_response = _gx.buckets.list()
|
||
if buckets_response.buckets:
|
||
for bucket in buckets_response.buckets:
|
||
if bucket.name == name:
|
||
return bucket.bucket_id
|
||
|
||
create_response = _gx.buckets.create(name=name)
|
||
return create_response.bucket.bucket_id
|
||
|
||
def ingest_document(gx: GroundX, bucket_id: str, path: Path, mime: str) -> str:
|
||
"""Upload and process document in Ground X, return processing identifier"""
|
||
bucket_id_int = int(bucket_id) if isinstance(bucket_id, str) else bucket_id
|
||
|
||
ingest = gx.ingest(
|
||
documents=[
|
||
Document(
|
||
bucket_id=bucket_id_int,
|
||
file_name=path.name,
|
||
file_path=str(path),
|
||
file_type=mime.split("/")[-1],
|
||
)
|
||
]
|
||
)
|
||
return ingest.ingest.process_id
|
||
|
||
def poll_until_complete(gx: GroundX, process_id: str, timeout: int = 600) -> None:
|
||
"""Monitor document processing status until completion"""
|
||
start_time = time.time()
|
||
status_text = st.empty()
|
||
progress_bar = st.progress(0)
|
||
|
||
while True:
|
||
status = gx.documents.get_processing_status_by_id(process_id=process_id).ingest
|
||
|
||
progress_value = 0
|
||
if hasattr(status, 'percent') and status.percent is not None:
|
||
try:
|
||
progress_value = int(status.percent)
|
||
except (ValueError, TypeError):
|
||
progress_value = 0
|
||
elif hasattr(status, 'progress') and status.progress is not None:
|
||
try:
|
||
if hasattr(status.progress, 'percent'):
|
||
progress_value = int(status.progress.percent)
|
||
elif hasattr(status.progress, 'value'):
|
||
progress_value = int(status.progress.value)
|
||
elif hasattr(status.progress, 'percentage'):
|
||
progress_value = int(status.progress.percentage)
|
||
else:
|
||
progress_value = int(status.progress)
|
||
except (ValueError, TypeError, AttributeError):
|
||
progress_value = 0
|
||
elif hasattr(status, 'percentage') and status.percentage is not None:
|
||
try:
|
||
progress_value = int(status.percentage)
|
||
except (ValueError, TypeError):
|
||
progress_value = 0
|
||
|
||
progress_bar.progress(progress_value)
|
||
|
||
status_display = f"**{status.status.capitalize()}**"
|
||
if progress_value > 0:
|
||
status_display += f" – {progress_value}%"
|
||
status_text.write(status_display)
|
||
|
||
if status.status in {"complete", "error", "cancelled"}:
|
||
break
|
||
if time.time() - start_time > timeout:
|
||
raise TimeoutError("Ground X ingest timed out.")
|
||
time.sleep(3)
|
||
|
||
if status.status != "complete":
|
||
raise RuntimeError(f"Ingest finished with status: {status.status!r}")
|
||
|
||
def fetch_xray_json(gx: GroundX, bucket_id: str) -> Dict[str, Any]:
|
||
"""Retrieve X-Ray analysis data for documents in storage bucket"""
|
||
documents = gx.documents.lookup(id=bucket_id).documents
|
||
if not documents:
|
||
raise RuntimeError("No documents found in bucket after ingest.")
|
||
|
||
document = documents[0]
|
||
if hasattr(document, 'xray_url') and document.xray_url:
|
||
response = requests.get(document.xray_url)
|
||
response.raise_for_status()
|
||
return response.json()
|
||
else:
|
||
raise RuntimeError("No X-Ray URL available for this document")
|
||
|
||
def check_file_exists(gx: GroundX, bucket_id: str, file_name: str) -> Optional[str]:
|
||
"""Verify document existence in bucket and return document identifier"""
|
||
documents = gx.documents.lookup(id=bucket_id).documents
|
||
for doc in documents:
|
||
if doc.file_name == file_name:
|
||
return doc.document_id
|
||
return None
|
||
|
||
def get_xray_for_existing_document(gx: GroundX, document_id: str, bucket_id: str) -> Dict[str, Any]:
|
||
"""Retrieve X-Ray analysis data for existing document"""
|
||
# Get the document to access its xray_url
|
||
documents = gx.documents.lookup(id=bucket_id).documents
|
||
for doc in documents:
|
||
if doc.document_id == document_id:
|
||
if hasattr(doc, 'xray_url') and doc.xray_url:
|
||
response = requests.get(doc.xray_url)
|
||
response.raise_for_status()
|
||
return response.json()
|
||
else:
|
||
raise RuntimeError("No X-Ray URL available for this document")
|
||
|
||
raise RuntimeError(f"Document with ID {document_id} not found")
|
||
|
||
def process_document(gx: GroundX, bucket_id: str, file_to_process: Any, file_path: str) -> tuple[Dict[str, Any], bool]:
|
||
"""Process document through Ground X pipeline and return analysis data
|
||
|
||
Args:
|
||
gx: Ground X client instance
|
||
bucket_id: Storage bucket identifier
|
||
file_to_process: File object with name and type attributes
|
||
file_path: Path to the file on disk
|
||
|
||
Returns:
|
||
Tuple of (xray_data, used_existing_file)
|
||
"""
|
||
existing_doc_id = check_file_exists(gx, bucket_id, file_to_process.name)
|
||
|
||
if existing_doc_id:
|
||
# Retrieve analysis for existing document
|
||
return get_xray_for_existing_document(gx, existing_doc_id, bucket_id), True
|
||
else:
|
||
# Process new document through ingestion pipeline
|
||
process_id = ingest_document(gx, bucket_id, Path(file_path), file_to_process.type)
|
||
poll_until_complete(gx, process_id)
|
||
return fetch_xray_json(gx, bucket_id), False
|