Spaces:
Sleeping
Sleeping
File size: 5,749 Bytes
9f4cbec | 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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """
Document / file ingestion layer.
Active path (MVP): Claude-native parsing. Uploaded PDFs and images are sent
directly to Claude (which reads PDFs and images natively) with an extraction
prompt that pulls vibration-relevant parameters into structured fields.
Scaffolded paths (for later, per the architecture plan):
- LlamaParse: AI-vision parsing of complex technical PDFs -> Markdown for RAG
- Unstructured.io: PDF/DOCX/PPT/image -> structured JSON elements
Both are stubbed with clear TODOs and a consistent return contract so they can
be dropped in without touching the rest of the app.
"""
import base64
import os
from typing import Optional
# --------------------------------------------------------------------------- #
# Active: Claude-native extraction
# --------------------------------------------------------------------------- #
EXTRACTION_SYSTEM = """\
You are extracting vibration-test-relevant parameters from an uploaded document \
(a mission spec, launch vehicle user's guide excerpt, ICD, or test article \
datasheet). Return ONLY a compact JSON object with any of these keys you can \
find (omit keys you cannot determine; never invent values):
{
"launch_vehicle": str,
"orbit": str,
"test_article_name": str,
"mass_kg": number,
"envelope_mm": str,
"random_vibration_psd": str, // any stated PSD breakpoints / Grms
"sine_environment": str,
"quasi_static_g": str,
"shock_srs": str,
"min_structural_freq_hz": number,
"mounting_interface": str, // bolt pattern, hole count/size if stated
"operating_modes": str,
"notes": str
}
Output the JSON and nothing else."""
def _encode_file(path: str) -> tuple[str, str]:
"""Return (base64_data, media_type) for a PDF or image."""
ext = os.path.splitext(path)[1].lower()
media = {
".pdf": "application/pdf",
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".gif": "image/gif",
".webp": "image/webp",
}.get(ext)
if media is None:
raise ValueError(f"Unsupported file type for extraction: {ext}")
with open(path, "rb") as f:
data = base64.standard_b64encode(f.read()).decode("utf-8")
return data, media
def extract_with_claude(path: str, api_key: Optional[str] = None) -> dict:
"""
Send the file to Claude and parse the returned JSON.
Returns {"ok": bool, "data": dict|None, "raw": str, "error": str|None}.
"""
api_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
if not api_key:
return {"ok": False, "data": None, "raw": "",
"error": "No ANTHROPIC_API_KEY set; cannot parse documents."}
try:
import json
import anthropic
data, media = _encode_file(path)
if media == "application/pdf":
source_block = {"type": "document",
"source": {"type": "base64",
"media_type": media, "data": data}}
else:
source_block = {"type": "image",
"source": {"type": "base64",
"media_type": media, "data": data}}
client = anthropic.Anthropic(api_key=api_key)
msg = client.messages.create(
model="claude-sonnet-4-6",
max_tokens=1500,
system=EXTRACTION_SYSTEM,
messages=[{"role": "user", "content": [
source_block,
{"type": "text",
"text": "Extract the vibration-relevant parameters as JSON."},
]}],
)
raw = "".join(b.text for b in msg.content if b.type == "text").strip()
# strip code fences if present
cleaned = raw.replace("```json", "").replace("```", "").strip()
parsed = json.loads(cleaned)
return {"ok": True, "data": parsed, "raw": raw, "error": None}
except Exception as e: # noqa: BLE001
return {"ok": False, "data": None, "raw": "", "error": str(e)}
# --------------------------------------------------------------------------- #
# Scaffolded: LlamaParse (TODO - requires LLAMA_CLOUD_API_KEY)
# --------------------------------------------------------------------------- #
def extract_with_llamaparse(path: str) -> dict:
"""
TODO (post-MVP): use LlamaParse for AI-vision parsing of complex technical
PDFs into Markdown, then index with LlamaIndex for RAG.
from llama_parse import LlamaParse
parser = LlamaParse(api_key=os.environ["LLAMA_CLOUD_API_KEY"],
result_type="markdown",
parsing_instruction="Extract test requirements "
"and acceptance criteria.")
docs = parser.load_data(path)
return {"ok": True, "markdown": docs[0].text, ...}
"""
return {"ok": False, "data": None, "raw": "",
"error": "LlamaParse not configured (scaffold). "
"Set LLAMA_CLOUD_API_KEY and implement."}
# --------------------------------------------------------------------------- #
# Scaffolded: Unstructured.io (TODO)
# --------------------------------------------------------------------------- #
def extract_with_unstructured(path: str) -> dict:
"""
TODO (post-MVP): use Unstructured.io to break PDF/DOCX/PPT/image into
structured elements (paragraphs, tables, lists, images) -> JSON for
LangChain ingestion.
from unstructured.partition.auto import partition
elements = partition(filename=path)
return {"ok": True, "elements": [el.to_dict() for el in elements], ...}
"""
return {"ok": False, "data": None, "raw": "",
"error": "Unstructured.io not configured (scaffold)."}
|