feat: key extractor

app/api/v1/keys_extract.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any, Dict, List, Optional
+
+from fastapi import APIRouter, Depends
+from pydantic import BaseModel, Field, field_validator
+
+from app.api.deps import require_auth
+from app.core.logger import get_logger
+from app.core.thread_pool import thread_pool
+from app.services.keys_extractor_service import KeysExtractor
+
+logger = get_logger(__name__)
+
+router = APIRouter()
+
+
+class KeysExtractRequest(BaseModel):
+    data: Any = Field(..., description="JSON object or array to search")
+    key_names: Optional[List[str]] = Field(
+        default=None,
+        description=(
+            "Key names to look up at any depth. Mutually exclusive with `query`; "
+            "provide one of the two."
+        ),
+    )
+    query: Optional[str] = Field(
+        default=None,
+        description=(
+            "Extended JSONPath expression (parsed with `jsonpath_ng.ext`) used to "
+            "select values. Supports filters, arithmetic and the union operator, e.g. "
+            "`$.store.book[?price > 10].title`. Each match is returned with its "
+            "full path inside the source document. Mutually exclusive with `key_names`."
+        ),
+    )
+    result_limit: Optional[int] = Field(
+        None,
+        ge=1,
+        description=(
+            "Maximum results to return. For `key_names` mode this caps values per "
+            "key; for `query` mode it caps the number of matches returned. "
+            "Omit (or pass null) to return the full, uncapped result."
+        ),
+    )
+
+    @field_validator("key_names")
+    @classmethod
+    def _validate_key_names(cls, v: Optional[List[str]]) -> Optional[List[str]]:
+        if v is None:
+            return v
+        if not v:
+            raise ValueError("key_names must be a non-empty list (omit it or supply a `query`)")
+        for kn in v:
+            if not isinstance(kn, str) or not kn:
+                raise ValueError("each key_name must be a non-empty string")
+        return v
+
+    @field_validator("query")
+    @classmethod
+    def _validate_query(cls, v: Optional[str]) -> Optional[str]:
+        if v is None:
+            return v
+        if not isinstance(v, str) or not v.strip():
+            raise ValueError("query must be a non-empty string")
+        return v
+
+
+class KeysExtractResponse(BaseModel):
+    success: bool
+    time_ms: float
+    data: Dict[str, Any]
+    error_message: Optional[str] = None
+
+
+class KeysExtractBatchRequest(BaseModel):
+    requests: List[KeysExtractRequest] = Field(
+        ...,
+        min_length=1,
+        max_length=50,
+        description="List of extraction requests (1-50) to process concurrently",
+    )
+
+
+class KeysExtractBatchResponse(BaseModel):
+    success: bool
+    time_ms: float
+    results: List[KeysExtractResponse]
+
+
+def _error_response(start: float, message: str) -> KeysExtractResponse:
+    return KeysExtractResponse(
+        success=False,
+        time_ms=round((time.perf_counter() - start) * 1000, 3),
+        data={},
+        error_message=message,
+    )
+
+
+async def _extract_single(body: KeysExtractRequest) -> KeysExtractResponse:
+    start = time.perf_counter()
+    has_key_names = bool(body.key_names)
+    has_query = body.query is not None
+
+    if has_key_names and has_query:
+        return _error_response(start, "provide either `key_names` or `query`, not both")
+    if not has_key_names and not has_query:
+        return _error_response(start, "either `key_names` or `query` must be provided")
+    if not isinstance(body.data, (dict, list)):
+        return _error_response(start, "`data` must be a JSON object or array")
+
+    try:
+        loop = asyncio.get_running_loop()
+        results = await loop.run_in_executor(
+            thread_pool,
+            KeysExtractor(body.data, body.key_names, body.result_limit, body.query).extract,
+        )
+    except (TypeError, ValueError) as exc:
+        return _error_response(start, str(exc))
+    except Exception as exc:
+        return _error_response(start, f"extraction failed: {exc}")
+
+    elapsed = round((time.perf_counter() - start) * 1000, 3)
+    return KeysExtractResponse(
+        success=True,
+        time_ms=elapsed,
+        data=results,
+        error_message=None,
+    )
+
+
+@router.post(
+    "/keys/extract",
+    response_model=KeysExtractBatchResponse,
+    dependencies=[Depends(require_auth)],
+    tags=["Keys Extractor"],
+    summary="Extract values from nested JSON objects (batch of up to 50 requests processed concurrently)",
+)
+async def extract_keys_batch(body: KeysExtractBatchRequest):
+    """Process up to 50 extraction requests concurrently.
+
+    Each request supports two modes:
+
+    * **key_names** (default) -- recursively walk any JSON object/array and return
+      every value attached to the supplied key names, regardless of how deeply
+      nested they are. Returns ``{key_name: [values...]}``.
+
+    * **query** -- run an extended JSONPath expression and return each match with
+      its full path inside the source document. Returns
+      ``{"matches": [{"path": "users.[0].role", "value": "admin"}, ...]}``.
+    """
+    start = time.perf_counter()
+    logger.info("keys_extract_batch | start", count=len(body.requests))
+    tasks = [_extract_single(req) for req in body.requests]
+    results = await asyncio.gather(*tasks)
+    elapsed = round((time.perf_counter() - start) * 1000, 3)
+    all_ok = all(r.success for r in results)
+    logger.info(
+        "keys_extract_batch | done",
+        total=len(results),
+        succeeded=sum(1 for r in results if r.success),
+        time_ms=elapsed,
+    )
+    return KeysExtractBatchResponse(success=all_ok, time_ms=elapsed, results=results)

app/api/v1/router.py

@@ -12,6 +12,7 @@ from app.api.v1 import (
     database,
     embeddings,
     json_extract,
+    keys_extract,
     qr_decoder,
     qr_generator,
     reconcile,
@@ -49,6 +50,7 @@ api_v1_router.include_router(vector_stores.router, tags=["Vector Stores"])
 api_v1_router.include_router(webhook_socket.router, tags=["Webhook / Socket"])
 api_v1_router.include_router(csv_analysis.router, tags=["CSV Analysis"])
 api_v1_router.include_router(json_extract.router, tags=["JSON Extractor"])
+api_v1_router.include_router(keys_extract.router, prefix="/json", tags=["Keys Extractor"])
 api_v1_router.include_router(qr_decoder.router, tags=["QR Decoder"])
 api_v1_router.include_router(qr_generator.router, tags=["QR Generator"])
 api_v1_router.include_router(url_shortener.router, tags=["URL Shortener"])

app/services/keys_extractor_service.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from typing import Any, List, Optional
+
+from jsonpath_ng import parse
+from jsonpath_ng.ext import parse as ext_parse
+from jsonpath_ng.exceptions import JsonPathLexerError, JsonPathParserError
+
+from app.core.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class KeysExtractor:
+    """Nested JSON extractor using jsonpath-ng.
+
+    Supports two modes:
+      1. key_names — recursive lookup by key name at any depth
+      2. query — extended JSONPath expression with full path output
+    """
+
+    def __init__(
+        self,
+        data: dict | list,
+        key_names: Optional[List[str]] = None,
+        result_limit: Optional[int] = 50,
+        query: Optional[str] = None,
+    ):
+        if query is None and not key_names:
+            raise ValueError("either `query` or `key_names` must be provided")
+
+        if key_names is not None:
+            if not isinstance(key_names, list):
+                raise TypeError("key_names must be a list")
+            for kn in key_names:
+                if not isinstance(kn, str) or not kn:
+                    raise ValueError(f"each key_name must be a non-empty string, got {kn!r}")
+
+        if query is not None:
+            if not isinstance(query, str) or not query.strip():
+                raise ValueError("query must be a non-empty string")
+
+        if result_limit is not None and result_limit < 1:
+            raise ValueError("result_limit must be >= 1 or None")
+
+        self._data = data
+        self._key_names = key_names or []
+        self._result_limit = result_limit
+        self._query = query
+
+    def extract(self) -> dict[str, Any]:
+        if self._query is not None:
+            return {"matches": self._extract_by_query()}
+        return {k: self._extract_single(k) for k in self._key_names}
+
+    def _extract_by_query(self) -> List[dict[str, Any]]:
+        try:
+            expr = ext_parse(self._query)
+        except (JsonPathLexerError, JsonPathParserError) as exc:
+            raise ValueError(f"invalid jsonpath query {self._query!r}: {exc}") from exc
+
+        matches: List[dict[str, Any]] = [
+            {"path": str(m.full_path), "value": m.value}
+            for m in expr.find(self._data)
+        ]
+
+        if self._result_limit is not None:
+            matches = matches[: self._result_limit]
+
+        return matches
+
+    def _extract_single(self, key_name: str) -> List[Any]:
+        try:
+            expr = _build_jp_expr(key_name)
+        except ValueError:
+            return []
+        vals = [m.value for m in expr.find(self._data)]
+        if self._result_limit is not None:
+            vals = vals[: self._result_limit]
+        return vals
+
+
+def _build_jp_expr(key_name: str) -> Any:
+    if '"' in key_name and "'" in key_name:
+        raise ValueError(f"key contains both quote types: {key_name!r}")
+    if '"' in key_name:
+        expr_str = "$..['" + key_name + "']"
+    else:
+        expr_str = '$..["' + key_name + '"]'
+    try:
+        return parse(expr_str)
+    except (JsonPathLexerError, JsonPathParserError):
+        raise ValueError(f"unable to build jsonpath expression for key: {key_name!r}")
+
+
+def get_nested_values(data: dict | list, key_name: str) -> List[Any]:
+    return KeysExtractor(data, key_names=[key_name], result_limit=None).extract()[key_name]

requirements.txt

@@ -62,5 +62,8 @@ Unidecode>=1.3.8
 # JSON Schema validation
 jsonschema>=4.21.0
 
+# JSONPath extraction
+jsonpath-ng>=1.7.0
+
 # QR code generation (testing)
 qrcode[pil]>=8.0