Spaces:
Sleeping
feat(upload-extract): switch primary LLM from NIM chain → Gemini 2.5-flash (JSON mode)
Browse filesUser-flagged: NIM was being used for upload extraction inherited
from rag/extract.py's offline pipeline, even though Gemini 2.5-flash
is the steady-state primary brain (ADR-040). Switching primary to
Gemini gives us:
- Native JSON mode (response_mime_type=application/json) — the
LLM is schema-locked, no prose to clean up before parsing
- The same Tier-0 brain every live chat turn already uses, with
no cold-start regression
- Better extraction quality on non-standard PDFs
NIM stays as the Tier-1 fallback for: (a) missing GOOGLE_API_KEY
(local dev / test machine), (b) Gemini 5xx / quota path. Same
EXTRACT_SYSTEM prompt, same HealthPolicy schema, same downstream
writes — only the transport changes.
Two-tier attempt: Gemini primary (180s timeout) → NIM fallback
(120s timeout). Both surface label in the warning log so the
operator can see which path each upload took.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- backend/uploaded_docs.py +39 -5
|
@@ -814,6 +814,17 @@ async def extract_one_for_upload(
|
|
| 814 |
)
|
| 815 |
from rag.schema import HealthPolicy
|
| 816 |
from backend.providers.base import ChatMessage
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 817 |
from backend.providers.nvidia_nim_llm import get_brain_llm
|
| 818 |
|
| 819 |
_log.info(
|
|
@@ -842,16 +853,39 @@ async def extract_one_for_upload(
|
|
| 842 |
ChatMessage(role="user", content=prompt),
|
| 843 |
]
|
| 844 |
|
| 845 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 846 |
llm_fallback = get_brain_llm()
|
| 847 |
|
| 848 |
raw = ""
|
| 849 |
policy: Optional[HealthPolicy] = None
|
| 850 |
-
for attempt, llm in enumerate(
|
|
|
|
|
|
|
| 851 |
try:
|
| 852 |
attempt_timeout = 180 if attempt == 0 else 120
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 853 |
res = await asyncio.wait_for(
|
| 854 |
-
llm.chat(
|
| 855 |
timeout=attempt_timeout,
|
| 856 |
)
|
| 857 |
raw = res.text
|
|
@@ -872,8 +906,8 @@ async def extract_one_for_upload(
|
|
| 872 |
break
|
| 873 |
except Exception as e: # noqa: BLE001
|
| 874 |
_log.warning(
|
| 875 |
-
"[upload-extract] attempt %d failed for %s: %s: %s",
|
| 876 |
-
attempt + 1, policy_id, type(e).__name__, str(e)[:200],
|
| 877 |
)
|
| 878 |
continue
|
| 879 |
|
|
|
|
| 814 |
)
|
| 815 |
from rag.schema import HealthPolicy
|
| 816 |
from backend.providers.base import ChatMessage
|
| 817 |
+
|
| 818 |
+
# 2026-05-27 — switched from NIM (get_brain_llm) to Gemini
|
| 819 |
+
# 2.5-flash with native JSON-mode (response_mime_type=
|
| 820 |
+
# application/json). Gemini is the steady-state primary
|
| 821 |
+
# chat brain (ADR-040) and gives schema-locked structured
|
| 822 |
+
# output, which is exactly what the EXTRACT prompt needs.
|
| 823 |
+
# NIM stays as the fallback for the (a) missing-GOOGLE_API_KEY
|
| 824 |
+
# case or (b) Gemini 5xx/quota path. Same prompt, same schema,
|
| 825 |
+
# same downstream HealthPolicy parse + writes — only the
|
| 826 |
+
# transport changes.
|
| 827 |
+
from backend.providers.google_gemini_llm import GoogleGeminiLLM
|
| 828 |
from backend.providers.nvidia_nim_llm import get_brain_llm
|
| 829 |
|
| 830 |
_log.info(
|
|
|
|
| 853 |
ChatMessage(role="user", content=prompt),
|
| 854 |
]
|
| 855 |
|
| 856 |
+
# Two-tier attempt: Gemini primary (native JSON mode, steady-state
|
| 857 |
+
# brain) → NIM chain fallback (older path, multi-candidate failover).
|
| 858 |
+
try:
|
| 859 |
+
llm_primary = GoogleGeminiLLM(timeout=180.0)
|
| 860 |
+
primary_label = "gemini-2.5-flash"
|
| 861 |
+
except Exception as e: # noqa: BLE001 — should never happen at __init__
|
| 862 |
+
_log.warning(
|
| 863 |
+
"[upload-extract] Gemini init failed (%s) — falling back to NIM",
|
| 864 |
+
type(e).__name__,
|
| 865 |
+
)
|
| 866 |
+
llm_primary = get_brain_llm()
|
| 867 |
+
primary_label = "nim-chain"
|
| 868 |
llm_fallback = get_brain_llm()
|
| 869 |
|
| 870 |
raw = ""
|
| 871 |
policy: Optional[HealthPolicy] = None
|
| 872 |
+
for attempt, (llm, label) in enumerate(
|
| 873 |
+
[(llm_primary, primary_label), (llm_fallback, "nim-chain")]
|
| 874 |
+
):
|
| 875 |
try:
|
| 876 |
attempt_timeout = 180 if attempt == 0 else 120
|
| 877 |
+
chat_kwargs = {
|
| 878 |
+
"messages": messages,
|
| 879 |
+
"temperature": 0.0,
|
| 880 |
+
"max_tokens": 2048,
|
| 881 |
+
}
|
| 882 |
+
# Gemini supports native JSON mode via response_format —
|
| 883 |
+
# forces the model to emit a JSON object the schema parser
|
| 884 |
+
# can validate without prose-cleanup. Harmless on the NIM
|
| 885 |
+
# path (the kwarg is absorbed by **kwargs).
|
| 886 |
+
chat_kwargs["response_format"] = {"type": "json_object"}
|
| 887 |
res = await asyncio.wait_for(
|
| 888 |
+
llm.chat(**chat_kwargs),
|
| 889 |
timeout=attempt_timeout,
|
| 890 |
)
|
| 891 |
raw = res.text
|
|
|
|
| 906 |
break
|
| 907 |
except Exception as e: # noqa: BLE001
|
| 908 |
_log.warning(
|
| 909 |
+
"[upload-extract] attempt %d (%s) failed for %s: %s: %s",
|
| 910 |
+
attempt + 1, label, policy_id, type(e).__name__, str(e)[:200],
|
| 911 |
)
|
| 912 |
continue
|
| 913 |
|