rohitsar567 Claude Opus 4.7 (1M context) commited on
Commit
7ef3ca3
·
1 Parent(s): 9ec9eae

feat(upload-extract): switch primary LLM from NIM chain → Gemini 2.5-flash (JSON mode)

Browse files

User-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>

Files changed (1) hide show
  1. backend/uploaded_docs.py +39 -5
backend/uploaded_docs.py CHANGED
@@ -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
- llm_primary = get_brain_llm()
 
 
 
 
 
 
 
 
 
 
 
846
  llm_fallback = get_brain_llm()
847
 
848
  raw = ""
849
  policy: Optional[HealthPolicy] = None
850
- for attempt, llm in enumerate([llm_primary, llm_fallback]):
 
 
851
  try:
852
  attempt_timeout = 180 if attempt == 0 else 120
 
 
 
 
 
 
 
 
 
 
853
  res = await asyncio.wait_for(
854
- llm.chat(messages=messages, temperature=0.0, max_tokens=2048),
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