From 47a45631b6575347d62793eee2b6aeed79a697b4 Mon Sep 17 00:00:00 2001 From: Jaseem Jas Date: Thu, 2 Apr 2026 13:22:15 +0530 Subject: [PATCH 1/8] UNS-480 [FEAT] Add Gemini LLM adapter for Google AI Studio Add a new LLM adapter for Google's Gemini models using LiteLLM's gemini/ provider prefix. The adapter follows the established SDK adapter pattern and is auto-discovered by register_adapters(). --- .../public/icons/adapter-icons/Gemini.png | Bin 0 -> 8231 bytes .../sdk1/src/unstract/sdk1/adapters/base1.py | 19 ++++++ .../src/unstract/sdk1/adapters/llm1/gemini.py | 40 +++++++++++++ .../sdk1/adapters/llm1/static/gemini.json | 56 ++++++++++++++++++ 4 files changed, 115 insertions(+) create mode 100644 frontend/public/icons/adapter-icons/Gemini.png create mode 100644 unstract/sdk1/src/unstract/sdk1/adapters/llm1/gemini.py create mode 100644 unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json diff --git a/frontend/public/icons/adapter-icons/Gemini.png b/frontend/public/icons/adapter-icons/Gemini.png new file mode 100644 index 0000000000000000000000000000000000000000..65d2bb32cb59d622c6ef5dd3572fd96e6e3caf27 GIT binary patch literal 8231 zcmV+?AlToDP)wWvYYk&VaKPTrL;9UOy z`=)H4?>*02d+p2jJnwqf{`Rkmzp|~V0z$8BYpTF22CS+4M6Vd|iUDh^z#3nEfoFg} z_y>U9_SrR-aL+Ekzy_3$UX?$3b-neo!{7YtdnUqK9=)gr*sZs$l&xHQ&C2Ei@kKSD zzR&=B@WBNVedVn#qygKd|M-3Ep@sBdAq^N6UyjT~K)LO9da2tSjax_q?ABWuzhGFn zZ6XR*9)5%ta9cbB?AFgZ69BcH7%(D7gSkHZF#cUEES>>TZbI>p-dT*_>r@~y{dK{i>v``oiDmD_Kr zoGNz)5JPzKhtcFO*BWh)>l8Scl z0O}V3=?9rS^X)(V;bTv%o{D>LrNuCy-hKx`+5h$3 z-M!}O9k&1m~i{xjspzw@PUuAY#)d!>0d0O}h< zddS^k<74TsoT_$(Av)9tx++)_fK&z0?>+wXWA{9?yUMK6adT~eJ^avc*BvCur87}e zbxm~yB-SQ<>N6JMG^wgFrVTTZFoX5cd;ak0$G^WyTJF|m=Gp+f-**VgY%GS{7cEWM zBB>~r+%!b&vi5_f(#lSOE9159sz^kR z@={dFnB%Y{2qpy7hwdYS57Gy?|H94l`;K}}4aj%iP7LXpSe+5YnJmDLD#WNrQWa4R z1xJ!G)?O^aec@Obs=$CgbAxnQy1Q>#qn8seHYeuMfO7X8!{_cYRa3lgKQF zPHMPu%g$8pyo!OU{u2!{=|G{)f9<)ddSLs(%BP zftUeZg$L4Cjp|<_PKM}Nmt`_Yp8<{-9E7_P=qrsR6-t>kbNa(4>lZgazjHxWd(XTW zkUxLh0QE(vy>tr?G`V8=1yW_8ud=SlBs0(>R$W(hf*t?{AQql9^r1d=H9SbgwUQPJ z^@s34zp#1wo13Srt@KW=Cp+n1)$aP}LrcGLhl@0*36NYMeOG0MB3+hALzlT9xGORl z5uMx@o)lS^5veZ#25>^+L;dPnT-prKOga$%`sU%$r|PPH$(>wm?F}g3dgw2I=KD%zD;Y*S7kT~2*TldOHwbvYhGCn2|@2V_!oGRS1$J_7=)U{cByxnIKu zEu3-yHb`MKA+lA%4i93Wv;|G{Ca-) zDPKLYcXXvSHel&L-AM_Ln{1)JmQG5N(rf87r>%e`DW^2VGf$A*$e@}~qd}^Z5d*i~S9c`K>h-Oe0rjy5pZk~B=5IgjSak&zU#fCdiK!r3 zG!@kxfbQK?`=KTNpj~c!(0(fD?;=XeddvWU^j+puIO4#m&4bj6PpTw^q{0IhIAJ8_ z{NR&x?cMdkwVMZPV?cfE54Qa4tLx(rBmG%B(TT_;iIS!&eMXg+iR(Jjb)MGMw7#SR z2R&e_`%PuPCE0HauiGN)w&-Mwu6tf5&&%XF>7SQwi}X{P3_cmCUjk&Zg(g#4mwi&A zGZl)}xVUC4%ayCK8$b92UvsBF_|)oU{mLuJ%Es#O`tASu+`qXhN=@PxQFcH$CqPhn zib_!}siBdaq928&fGB5`EL=phfXFa5Nl;QM5dh+-Fiy^d;6(BN2Z1 z)1WCg1eb&&lO<(TE#6jP3P~k4M!Ba-DEA}W+>JtRL~{b%(FuEEz|Hp(u*wdsssYpg zeDjw7@)_&JI%yto=>&iQvGDd?=~C(8!Ds+P2l|`>L+S>7AwqXZ4-dL5U11H76qd2j z7w1Yd7)7QVh3JS#e5h~a6+I)GAD_OLfYtP1^$d9a-~FR4zjw1IG9_!V#;VIDamAC5 zqLS!ii7c90~NPvoWa3q!3nhj7q zbPu`Fv%?0)UNWL;C`ax@zq;Aq@^Sy>ldIvuJ~v?a%8k!``tsq)Z-nAw z?r%;1#8_0JHDXmARECI6dX3tWQDn3mWQHu=4BQ9dE)yPXx`pb&M4l zC@`So(n-N@F08{z1^_lB9c1|GN*4V+3eL;b>%2{V9U zLnSH{>u_!5-b3zK24*DufF(Ny%=p!>(Z2RzUl~xJ{2xF3Z`?d=8Y zFbEdj+Tq~`k3awv=OO8|)1@R!olX--x0E`NA*ChNB~81gP9sn;!CLYf!6_|#M+Wl=Pj zQfkntzylqW3$!LN*2<_kie^QghH_S526JHv5EtSafFD#nMpVbU2epVG;_ZLX!OkOa z_gj_UA-Qq0Uw)@!LSW=x#9lMN{_N{Zzxm$bflr{^Q)htpX1&)wgdRGn31(1EJajsA z0^KvzPd-h&CCQfPOgW{jfDV$TiW1bIX%O_t0QH3{tC3cR2LQ}qWHY;E_e@k617~*9HP(Cbo~;kLYv;8icmN>TYvNXLxS-l3Mztv_ummJ4L*J>&u#Dg$ zc%U2@yEro!CS-SIEhJUA(lS_T2nu~5SQK=J1K@}eKm;qk;5tK0-W{`|Yka(IHg=I+4J(7J((k+VIZ%eM| zlE{)uL{Y4?g>_)f3^N4JIhmu{R1WnE3#$28${h(zWC3@N84x^Z6ryI$j>*flX0O<= zn+EvP_th`|bp7&0_RM|IJ%EbeJP6*1%}Njok=5QxB1{6Hc3fNfXVmm0IcZizSxP$9 z8Bl@-(EuaLz=Bm_Dl!N+E~EO;9Zr-s6^wW$t~Y%Vn-kFqPUO5R?Dh=VLmOV60sg&v z%YVCY`uUC1|N5@_7x!8+>%76D^e3YR{7M}W1CrHlgx5s1Nk7IgK)5vI5EuF?2mt9J zF#w5y$_kGDES!Zh6OPKvZI>VxM!9drnW%}b;Lcla?jQgfTB-{Byl`+iesq*C^=lQzK5TOyz+ z%s^8sCIK@R&rEKTDr-TLp)6GvV~I$LScEeZViCy#+c1C%QxzT&Gk~EoYCYm8bmY7f znynL6Xy+hQtCfoG7KaVMxIDZ>Cw66kKXY%n|HAy+hnBCu3%a|oDS0ASXcrMu)2Nlju&t%D6xJF!j*Vp4FB5ksQfFmyC1z|%xiCP<@e#e?Kp?mGZ36RnY)E|-ulyhSEF;y54 zXdYm4Kz+-C=vji~aKZ{3W=5sZ+K~QaPK~!nm;(SA#l{|4kBA22$QpToNEPNlAj9qFhjU z+9b$I4V>}TK@cem0HF~Ha%XWwpo#;iJe5(Tp>n_}xRXUm??dW$000lrNkl$5?sDOXhnz5bD0*6ngOa%u39zZd$9Bl1ffFr@{MbeXWyP z5&N9`0+dND6SH;36Ovx36YIEkTsq3?Ma^0gy_#Zhsr zV=%KJqK6%3fCH7iGF9=~`q4zE6Y@RM7e^Teqlp6v8{-3MV+_ZbP@K)l#(o?VWZMTF zBU;xA+S;oTu6Wz_SCQ>H_~5tpKmJTQ{(2TA(TLolsEc?QBS8dvg`x@}29L!61F$wA zHUnZWnj<9SHsQKecX4m?z!U~hix&C{VnYzOh=zV=$ZC)qaPl#NI6H;nh+Bd01ioEF zb1kvySbE_4mN7593wdF4*3#pyot*!4dhK-pbo80fO>0#O^&@(?qOJ|17Q?{BAi97= z5B84ZO}Ppipde&Qk?D$L_wk29m1+PZ8c2ooQY%;VH<+^#(w~fw=7EYHZ%vBm6lPGw zAgzf*g40@U(EX-M7}6u|@E&A`y<=xO<{DtXxBhu|>~-C7*I5XS^j8ffVOL48^17dm}Ev z$4-aPC77V9;Gk;as^G zL48aD#HKr%0k~qsMq(ZkdlR}vcHvmi@GUrq^VW?B9kdTNXpKxC^Bw%9T;U}du=PZf zBd-DuKKGB;9rX*=q=&eb)kG^x5h}u%0T$BhNYGQ15naSNRB?C^x>FDd(Ot~dJmjtv z24tfG;*sf)KAOgu1U0lqCp7{Cz`1d0HV~x8&DzWc29Ik-GdP+J%y6Kwy%V7I%hi4B zgv^w_J8{IE$fRg=hc!-P<&5*}1kAqno?NxqaZiNz7P~mg3x#b zVl<$&H{CsKfcn_Q&=;VcY`fUClOuNC4aK250J8H@FLjT+)PXNH?%+c&Kj96(|LQ}p zaE1r!oSYFQ#A^Zv!PUegAE8%(2)dIgRiO%P1@e*!8MATpRF2YLhzm;rCo0j8$c18n z5jkOt5R=2%h?_BEe|TnX3{L1wE5P8m{tP(B%AGslF@mmq8}BXY+xZBRjAZY`4I6*% zxYvGakRY^ ztV^7hyxT&GKZs4+AI%2UX_7$!IET|faEOehLfhZRv{Kkij>!?ISl&qr(TGX}6z1As z%m6XduD4Ccg!F&)r|hnedeOb|q9b1JfLEHgk`S7)tBt|Y;3f<>YFm>DdG5pz+lU=`%>wO$fX~qSgJT%u3;@OUBKqPaF=}U@ zo8C*nUa2|TT>{^FyPmz%JYvK9&pYKS$8CInNF35L6cA=%ip2oBqccG9vsy?6DjVch zEZZEbhKPYHBUYx0CPZ&Ggy=y8HUKqj5Oh}JN@UAu(uW6xYcn|xr`VjpN5bI2mU~A2 zf(X$cxPcH~(7X2<`^tdv4v*RJ4~~D$`51Wl6H}t|^8mwI->hfkvv}qCk@&O5 zzWHb08i0E~cEjI4>)3yVuP+=zAOg%LjMPS~!VMCl3uomlg16BHe71{=wlTzoB{Qly z7mAnDh?OC?G~Ef!LWop}!Y#}UB3zUxE=_H7Vk36+H=xgH)&sr`NO{zpyAAnxW6_7~(s*qmdPh=o?ViLw|D2O7M5Sh79pte|8Dlbn3<_x|( z0K+QeW{#^6hx#LD!n7c$43Q%;@rdXHcgBs8U->q==@MGCZS@SmS9;xsi%;5kiLeWk zhyz#y03@u6Dc6K3h!tF=ih!7K6*h1Vw@QnZ7zC&+m87CGVi2fo#2A1G?kq@cfS@r) zz=uW>k-XV34)80^r&TCBv|?LT1906FHe7rPJQ&p%*JXqUY6?V8kfu94ur|iFNi0SX zT0&w5XY3laMi2Zz6Pf}+2=f3E6LD<#2PA-gg8caQ%s~*CZoEXUcq`NDZn%Zjy&%zv z8!masq@lgEaQxj~azxI?gE0n79m8aB z^lKrk9AXj<;k04%n~srBy(gWs7V1ONwJ`t}e$9Sw`0*qE*%AAlu0$dI=mfzGhb7Wd zo7tFKeUEvNJDw&PV1V+Jc}l$G3|XOmhPVa;q~NWCa|S|vnnvV^9tI4=83th4XlKUB z{hVXu-@ZrAUL*CPgW$7;5n?eGR3T@iDkGw@!-gm|&Q1-eLd?W1 zk%?=A;T5UJ%OAWO3XnL7#i_-a0Ti5w#DtSDqv+Ak{;1HJZfkt`jr*PU_GA7DCIPaZ zLvEBBQ+5?Bzj!c3EQ0<(sLWV<1%5zQqnIDY;uJ9A5+Do~%4hT0Bm(U}Xq&k;~!7VtZ z0;&)?qa0yC?8N{=|1y|sL%HlN-PuR2{nuanJAem=tv~IojeoZp02{`%Z8iwRDuQ2R zAW$~ePskdtmr>;jLTmE?B81m4xjDfc830Az`PyGPYX4J5586*BC$a{o zc)W?)0x5mwN4Vlt3E=94y_K*lGCjCc#Yk@kpr5ge>!q6x#)akw^JM^ZdGlfKMhtj_ z_Kbq50FEb1th6=!fZR<@Jaj_YF%rWnC~hEl9z_f^Z`v3XAN%CrS-S^*e1jcf?hH8U zRi~VH#Jd=aag-t=U;Zbi0vZ-w(F+2E;tn$+*z#Wh0i$C8XaF>N--$O=cVigf$a8adLfQi(ac6g$<#<#WivPBjv03~$ zs1_?TmVxleV-2X6pE-BM2VWDIQv+~km!Eunh|VZIzU+h>0k#}L3^o{!`cny)O^q#S;;1AS=`iV1o^nY9)9OFs{bMB!})YT9PXd zA3FVXoMfIc?*>4-cgG}v)>bJhPa$xmf32I%ZrcsDRUbE z-9akM;N^0JT~wV7HsJS~Q8lg~KK0EwaQ-o$25hg}83&$v`a!2u8lQqJ9)ZZh1A^#b z115n{F~$lER5qM01M`#(FB<}c66CO?gcG?^t~@P(45rhM zoOT`#SqLnQ0cRa}@*594+1&AF0?U9Cew5+O#8zN2AYydJoVe;KxKsE_+xBm;1}=Nx=;CHx40(2iFW4Q>LC%0p;=^rXLmVGD!>GvI^Ay%Uqbh>baW zDUthFjjeL?fWw9bQXjXvUY$)xNqc!6W+0K>f?qN(g57|j~;UJxB!gtx)Uxy=OSWJ z4fyRh-7reuB3*yNg^Q>@?s-uSz^(t_?YG=;!n=Ow+>iavx!0k&pjco7pg0UT|BzDv z@CC(U8?d0Fue|M73|Ohg8t_*PSOd{lUh@BY16GQ$7X1GK00960I4g~~00006Nkl str: return f"anthropic/{model}" +class GeminiLLMParameters(BaseChatCompletionParameters): + """See https://docs.litellm.ai/docs/providers/gemini.""" + + api_key: str + + @staticmethod + def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: + adapter_metadata["model"] = GeminiLLMParameters.validate_model(adapter_metadata) + return GeminiLLMParameters(**adapter_metadata).model_dump() + + @staticmethod + def validate_model(adapter_metadata: dict[str, "Any"]) -> str: + model = adapter_metadata.get("model", "") + if model.startswith("gemini/"): + return model + else: + return f"gemini/{model}" + + class AnyscaleLLMParameters(BaseChatCompletionParameters): """See https://docs.litellm.ai/docs/providers/anyscale.""" diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/gemini.py b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/gemini.py new file mode 100644 index 0000000000..395ddccd01 --- /dev/null +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/gemini.py @@ -0,0 +1,40 @@ +from typing import Any + +from unstract.sdk1.adapters.base1 import BaseAdapter, GeminiLLMParameters +from unstract.sdk1.adapters.enums import AdapterTypes + + +class GeminiLLMAdapter(GeminiLLMParameters, BaseAdapter): + @staticmethod + def get_id() -> str: + return "gemini|085f6c03-b57e-4594-85bb-40e2616c2736" + + @staticmethod + def get_metadata() -> dict[str, Any]: + return { + "name": "Gemini", + "version": "1.0.0", + "adapter": GeminiLLMAdapter, + "description": "Google Gemini LLM adapter via Google AI Studio", + "is_active": True, + } + + @staticmethod + def get_name() -> str: + return "Gemini" + + @staticmethod + def get_description() -> str: + return "Google Gemini LLM adapter via Google AI Studio" + + @staticmethod + def get_provider() -> str: + return "gemini" + + @staticmethod + def get_icon() -> str: + return "/icons/adapter-icons/Gemini.png" + + @staticmethod + def get_adapter_type() -> AdapterTypes: + return AdapterTypes.LLM diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json new file mode 100644 index 0000000000..90a171d892 --- /dev/null +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json @@ -0,0 +1,56 @@ +{ + "title": "Gemini LLM", + "type": "object", + "required": ["adapter_name", "api_key", "model"], + "properties": { + "adapter_name": { + "type": "string", + "title": "Name", + "default": "", + "description": "Provide a unique name for this adapter instance. Example: gemini-group-1" + }, + "api_key": { + "type": "string", + "title": "API Key", + "default": "", + "description": "Google AI Studio API key", + "format": "password" + }, + "model": { + "type": "string", + "title": "Model", + "default": "gemini/gemini-1.5-flash", + "description": "LiteLLM model string. Supported: gemini-1.5-flash, gemini-1.5-pro, gemini-2.0-flash, gemini-2.5-pro. The gemini/ prefix will be added automatically if omitted." + }, + "temperature": { + "type": "number", + "minimum": 0, + "maximum": 2, + "title": "Temperature", + "default": 0.1, + "description": "Sampling temperature between 0 and 2" + }, + "max_tokens": { + "type": "number", + "minimum": 0, + "multipleOf": 1, + "title": "Maximum Output Tokens", + "description": "Maximum number of output tokens to limit LLM replies, the maximum possible differs from model to model." + }, + "timeout": { + "type": "number", + "minimum": 0, + "multipleOf": 1, + "title": "Timeout", + "default": 600, + "description": "Timeout in seconds" + }, + "max_retries": { + "type": "number", + "minimum": 0, + "multipleOf": 1, + "title": "Max Retries", + "description": "Maximum number of retries" + } + } +} From 862e0aece780bbbfa921c0b60b170792370dd453 Mon Sep 17 00:00:00 2001 From: Jaseem Jas Date: Thu, 2 Apr 2026 13:53:17 +0530 Subject: [PATCH 2/8] UNS-480 [FEAT] Add defaults for max_tokens and max_retries in Gemini JSON schema --- .../sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json | 2 ++ 1 file changed, 2 insertions(+) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json index 90a171d892..73f5efec86 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json @@ -34,6 +34,7 @@ "type": "number", "minimum": 0, "multipleOf": 1, + "default": 8192, "title": "Maximum Output Tokens", "description": "Maximum number of output tokens to limit LLM replies, the maximum possible differs from model to model." }, @@ -49,6 +50,7 @@ "type": "number", "minimum": 0, "multipleOf": 1, + "default": 3, "title": "Max Retries", "description": "Maximum number of retries" } From 3e7e715ccbd9b1908d8fb939f06f2c8154866dc0 Mon Sep 17 00:00:00 2001 From: Jaseem Jas Date: Thu, 2 Apr 2026 15:55:16 +0530 Subject: [PATCH 3/8] UNS-480 [FIX] Address PR review: avoid dict mutation, validate blank model, update default model --- unstract/sdk1/src/unstract/sdk1/adapters/base1.py | 9 ++++++--- .../src/unstract/sdk1/adapters/llm1/static/gemini.json | 4 ++-- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py index 0ee004ac0e..444c2f9090 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py @@ -646,12 +646,15 @@ class GeminiLLMParameters(BaseChatCompletionParameters): @staticmethod def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: - adapter_metadata["model"] = GeminiLLMParameters.validate_model(adapter_metadata) - return GeminiLLMParameters(**adapter_metadata).model_dump() + result_metadata = adapter_metadata.copy() + result_metadata["model"] = GeminiLLMParameters.validate_model(adapter_metadata) + return GeminiLLMParameters(**result_metadata).model_dump() @staticmethod def validate_model(adapter_metadata: dict[str, "Any"]) -> str: - model = adapter_metadata.get("model", "") + model = str(adapter_metadata.get("model", "")).strip() + if not model: + raise ValueError("model is required") if model.startswith("gemini/"): return model else: diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json index 73f5efec86..31df5972b1 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json @@ -19,8 +19,8 @@ "model": { "type": "string", "title": "Model", - "default": "gemini/gemini-1.5-flash", - "description": "LiteLLM model string. Supported: gemini-1.5-flash, gemini-1.5-pro, gemini-2.0-flash, gemini-2.5-pro. The gemini/ prefix will be added automatically if omitted." + "default": "gemini-2.0-flash", + "description": "LiteLLM model string. Supported: gemini-2.0-flash, gemini-2.5-pro, gemini-1.5-pro, gemini-1.5-flash. The gemini/ prefix will be added automatically if omitted." }, "temperature": { "type": "number", From f5fdb0d6a7f7a739881cc5e071d7c7dc2a8165ad Mon Sep 17 00:00:00 2001 From: Jaseem Jas Date: Thu, 9 Apr 2026 10:27:36 +0530 Subject: [PATCH 4/8] UNS-482 [FEAT] Add Gemini thinking mode support with tests Extends the Gemini LLM adapter with optional thinking mode, mirroring the Anthropic/Bedrock pattern. enable_thinking and budget_tokens are consumed from adapter metadata (not Pydantic fields); when enabled, temperature is forced to 1. Schema gains an allOf/if-then conditional so budget_tokens (min 1024) is only required when thinking is on. --- .../sdk1/src/unstract/sdk1/adapters/base1.py | 38 ++++- .../sdk1/adapters/llm1/static/gemini.json | 30 +++- unstract/sdk1/tests/test_gemini_adapter.py | 154 ++++++++++++++++++ 3 files changed, 220 insertions(+), 2 deletions(-) create mode 100644 unstract/sdk1/tests/test_gemini_adapter.py diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py index 444c2f9090..7a37f9711c 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py @@ -648,7 +648,43 @@ class GeminiLLMParameters(BaseChatCompletionParameters): def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: result_metadata = adapter_metadata.copy() result_metadata["model"] = GeminiLLMParameters.validate_model(adapter_metadata) - return GeminiLLMParameters(**result_metadata).model_dump() + + # Handle Gemini thinking configuration + enable_thinking = adapter_metadata.get("enable_thinking", False) + + # If enable_thinking is not explicitly provided but thinking config is present, + # assume thinking was enabled in a previous validation + has_thinking_config = ( + "thinking" in adapter_metadata + and adapter_metadata.get("thinking") is not None + ) + if not enable_thinking and has_thinking_config: + enable_thinking = True + + if enable_thinking: + if has_thinking_config: + result_metadata["thinking"] = adapter_metadata["thinking"] + else: + thinking_config: dict[str, "Any"] = {"type": "enabled"} + budget_tokens = adapter_metadata.get("budget_tokens") + if budget_tokens is not None: + thinking_config["budget_tokens"] = budget_tokens + result_metadata["thinking"] = thinking_config + # Gemini thinking mode requires temperature=1 + result_metadata["temperature"] = 1 + + # Exclude control fields from pydantic validation + exclude_fields = ("enable_thinking", "budget_tokens", "thinking") + validation_metadata = { + k: v for k, v in result_metadata.items() if k not in exclude_fields + } + + validated = GeminiLLMParameters(**validation_metadata).model_dump() + + if enable_thinking and "thinking" in result_metadata: + validated["thinking"] = result_metadata["thinking"] + + return validated @staticmethod def validate_model(adapter_metadata: dict[str, "Any"]) -> str: diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json index 31df5972b1..88235996bd 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json @@ -53,6 +53,34 @@ "default": 3, "title": "Max Retries", "description": "Maximum number of retries" + }, + "enable_thinking": { + "type": "boolean", + "title": "Enable Thinking Mode", + "default": false, + "description": "Enable extended thinking for supported models. Thinking mode is only supported on: gemini-2.5-pro, gemini-2.5-flash, gemini-2.0-flash-thinking-exp. When enabled, temperature is forced to 1." + } + }, + "allOf": [ + { + "if": { + "properties": { + "enable_thinking": { "const": true } + }, + "required": ["enable_thinking"] + }, + "then": { + "required": ["budget_tokens"], + "properties": { + "budget_tokens": { + "type": "integer", + "minimum": 1024, + "default": 1024, + "title": "Budget Tokens", + "description": "Number of tokens allocated for the thinking process. Minimum: 1024." + } + } + } } - } + ] } diff --git a/unstract/sdk1/tests/test_gemini_adapter.py b/unstract/sdk1/tests/test_gemini_adapter.py new file mode 100644 index 0000000000..cccfc352ce --- /dev/null +++ b/unstract/sdk1/tests/test_gemini_adapter.py @@ -0,0 +1,154 @@ +"""Unit tests for the Gemini LLM adapter (UNS-480 / UNS-482).""" + +import json +from pathlib import Path + +import pytest +from unstract.sdk1.adapters.base1 import GeminiLLMParameters +from unstract.sdk1.adapters.llm1.gemini import GeminiLLMAdapter + +BASE_METADATA = {"api_key": "test-key", "model": "gemini-2.5-flash"} + + +# ── validate_model ─────────────────────────────────────────────────────────── + + +def test_validate_model_prefixes_when_missing() -> None: + assert ( + GeminiLLMParameters.validate_model({"model": "gemini-2.5-flash"}) + == "gemini/gemini-2.5-flash" + ) + + +def test_validate_model_does_not_double_prefix() -> None: + assert ( + GeminiLLMParameters.validate_model({"model": "gemini/gemini-2.5-pro"}) + == "gemini/gemini-2.5-pro" + ) + + +def test_validate_model_blank_raises() -> None: + with pytest.raises(ValueError, match="model is required"): + GeminiLLMParameters.validate_model({"model": " "}) + + +# ── validate: thinking disabled ────────────────────────────────────────────── + + +def test_validate_thinking_disabled_by_default() -> None: + result = GeminiLLMParameters.validate({**BASE_METADATA, "temperature": 0.3}) + assert result["model"] == "gemini/gemini-2.5-flash" + assert "thinking" not in result + assert result["temperature"] == 0.3 + + +def test_validate_excludes_control_fields_from_model() -> None: + result = GeminiLLMParameters.validate(BASE_METADATA.copy()) + assert "enable_thinking" not in result + assert "budget_tokens" not in result + + +# ── validate: thinking enabled ─────────────────────────────────────────────── + + +def test_validate_thinking_enabled_with_budget() -> None: + result = GeminiLLMParameters.validate( + {**BASE_METADATA, "enable_thinking": True, "budget_tokens": 2048} + ) + assert result["thinking"] == {"type": "enabled", "budget_tokens": 2048} + assert result["temperature"] == 1 + + +def test_validate_thinking_overrides_user_temperature() -> None: + result = GeminiLLMParameters.validate( + { + **BASE_METADATA, + "temperature": 0.7, + "enable_thinking": True, + "budget_tokens": 1024, + } + ) + assert result["temperature"] == 1 + + +def test_validate_thinking_enabled_without_budget() -> None: + result = GeminiLLMParameters.validate({**BASE_METADATA, "enable_thinking": True}) + assert result["thinking"] == {"type": "enabled"} + assert result["temperature"] == 1 + + +def test_validate_preserves_existing_thinking_config() -> None: + existing = {"type": "enabled", "budget_tokens": 4096} + result = GeminiLLMParameters.validate({**BASE_METADATA, "thinking": existing}) + assert result["thinking"] == existing + assert result["temperature"] == 1 + + +def test_validate_does_not_mutate_input() -> None: + metadata = {**BASE_METADATA, "enable_thinking": True, "budget_tokens": 2048} + snapshot = metadata.copy() + GeminiLLMParameters.validate(metadata) + assert metadata == snapshot + + +# ── Pydantic field surface ─────────────────────────────────────────────────── + + +def test_thinking_controls_not_pydantic_fields() -> None: + fields = GeminiLLMParameters.model_fields + assert "enable_thinking" not in fields + assert "budget_tokens" not in fields + assert "thinking" not in fields + assert "api_key" in fields + + +def test_api_key_is_required() -> None: + from pydantic import ValidationError + + with pytest.raises(ValidationError): + GeminiLLMParameters(model="gemini/gemini-2.5-flash") + + +# ── Adapter identity ───────────────────────────────────────────────────────── + + +def test_adapter_identity() -> None: + assert GeminiLLMAdapter.get_name() == "Gemini" + assert GeminiLLMAdapter.get_provider() == "gemini" + assert GeminiLLMAdapter.get_id().startswith("gemini|") + metadata = GeminiLLMAdapter.get_metadata() + assert metadata["is_active"] is True + assert metadata["name"] == "Gemini" + + +# ── JSON schema ────────────────────────────────────────────────────────────── + + +@pytest.fixture +def gemini_schema() -> dict: + schema_path = ( + Path(__file__).parent.parent + / "src/unstract/sdk1/adapters/llm1/static/gemini.json" + ) + return json.loads(schema_path.read_text()) + + +def test_schema_required_fields(gemini_schema: dict) -> None: + assert set(gemini_schema["required"]) >= {"adapter_name", "api_key", "model"} + + +def test_schema_enable_thinking_default_false(gemini_schema: dict) -> None: + assert gemini_schema["properties"]["enable_thinking"]["default"] is False + + +def test_schema_budget_tokens_conditional(gemini_schema: dict) -> None: + all_of = gemini_schema["allOf"] + assert len(all_of) == 1 + conditional = all_of[0] + assert conditional["if"]["properties"]["enable_thinking"]["const"] is True + then_block = conditional["then"] + assert "budget_tokens" in then_block["required"] + budget = then_block["properties"]["budget_tokens"] + assert budget["minimum"] == 1024 + assert budget["default"] == 1024 + assert "maximum" not in budget From 1b34aadccab692be307b7443e1b0525ff81b51d1 Mon Sep 17 00:00:00 2001 From: Jaseem Jas Date: Thu, 9 Apr 2026 15:52:56 +0530 Subject: [PATCH 5/8] UNS-482 [FIX] Use pytest.approx for temperature float comparison --- unstract/sdk1/tests/test_gemini_adapter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unstract/sdk1/tests/test_gemini_adapter.py b/unstract/sdk1/tests/test_gemini_adapter.py index cccfc352ce..58c82d54ea 100644 --- a/unstract/sdk1/tests/test_gemini_adapter.py +++ b/unstract/sdk1/tests/test_gemini_adapter.py @@ -39,7 +39,7 @@ def test_validate_thinking_disabled_by_default() -> None: result = GeminiLLMParameters.validate({**BASE_METADATA, "temperature": 0.3}) assert result["model"] == "gemini/gemini-2.5-flash" assert "thinking" not in result - assert result["temperature"] == 0.3 + assert result["temperature"] == pytest.approx(0.3) def test_validate_excludes_control_fields_from_model() -> None: From 8a4fa49056e6fcdd8b052c03a0cb7b455a9ee80b Mon Sep 17 00:00:00 2001 From: Jaseem Jas Date: Thu, 9 Apr 2026 16:03:49 +0530 Subject: [PATCH 6/8] UNS-482 [FIX] Validate budget_tokens when Gemini thinking mode is enabled Raise ValueError if budget_tokens is missing, not an integer, or below 1024 when enable_thinking=True. Previously these cases silently produced incomplete or invalid thinking configs that would fail at the API level. --- .../sdk1/src/unstract/sdk1/adapters/base1.py | 16 ++++++++++---- unstract/sdk1/tests/test_gemini_adapter.py | 21 +++++++++++++++---- 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py index 7a37f9711c..f7d98f0fa1 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py @@ -665,11 +665,19 @@ def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: if has_thinking_config: result_metadata["thinking"] = adapter_metadata["thinking"] else: - thinking_config: dict[str, "Any"] = {"type": "enabled"} budget_tokens = adapter_metadata.get("budget_tokens") - if budget_tokens is not None: - thinking_config["budget_tokens"] = budget_tokens - result_metadata["thinking"] = thinking_config + if budget_tokens is None: + raise ValueError( + "budget_tokens is required when thinking mode is enabled" + ) + if not isinstance(budget_tokens, int) or budget_tokens < 1024: + raise ValueError( + f"budget_tokens must be an integer >= 1024, got {budget_tokens}" + ) + result_metadata["thinking"] = { + "type": "enabled", + "budget_tokens": budget_tokens, + } # Gemini thinking mode requires temperature=1 result_metadata["temperature"] = 1 diff --git a/unstract/sdk1/tests/test_gemini_adapter.py b/unstract/sdk1/tests/test_gemini_adapter.py index 58c82d54ea..fb59bcf611 100644 --- a/unstract/sdk1/tests/test_gemini_adapter.py +++ b/unstract/sdk1/tests/test_gemini_adapter.py @@ -71,10 +71,23 @@ def test_validate_thinking_overrides_user_temperature() -> None: assert result["temperature"] == 1 -def test_validate_thinking_enabled_without_budget() -> None: - result = GeminiLLMParameters.validate({**BASE_METADATA, "enable_thinking": True}) - assert result["thinking"] == {"type": "enabled"} - assert result["temperature"] == 1 +def test_validate_thinking_enabled_without_budget_raises() -> None: + with pytest.raises(ValueError, match="budget_tokens is required"): + GeminiLLMParameters.validate({**BASE_METADATA, "enable_thinking": True}) + + +def test_validate_thinking_budget_tokens_invalid_type_raises() -> None: + with pytest.raises(ValueError, match="budget_tokens must be an integer >= 1024"): + GeminiLLMParameters.validate( + {**BASE_METADATA, "enable_thinking": True, "budget_tokens": "hello"} + ) + + +def test_validate_thinking_budget_tokens_too_small_raises() -> None: + with pytest.raises(ValueError, match="budget_tokens must be an integer >= 1024"): + GeminiLLMParameters.validate( + {**BASE_METADATA, "enable_thinking": True, "budget_tokens": 512} + ) def test_validate_preserves_existing_thinking_config() -> None: From 54b0ad31cb682ad2e6fe56e99948d11db7b9f345 Mon Sep 17 00:00:00 2001 From: Jaseem Jas Date: Thu, 9 Apr 2026 16:06:08 +0530 Subject: [PATCH 7/8] UNS-480 [FIX] Add gemini-2.5-flash to model description in JSON schema --- .../sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json index 88235996bd..569948031e 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json @@ -20,7 +20,7 @@ "type": "string", "title": "Model", "default": "gemini-2.0-flash", - "description": "LiteLLM model string. Supported: gemini-2.0-flash, gemini-2.5-pro, gemini-1.5-pro, gemini-1.5-flash. The gemini/ prefix will be added automatically if omitted." + "description": "LiteLLM model string. Supported: gemini-2.0-flash, gemini-2.5-pro, gemini-2.5-flash, gemini-1.5-pro, gemini-1.5-flash. The gemini/ prefix will be added automatically if omitted." }, "temperature": { "type": "number", From 1f356476f2726f818b3279468534e0631a20add6 Mon Sep 17 00:00:00 2001 From: Jaseem Jas Date: Fri, 10 Apr 2026 09:42:38 +0530 Subject: [PATCH 8/8] UNS-480 [FIX] Clean up Gemini JSON schema descriptions per PR review Remove internal implementation detail (LiteLLM) from model description and remove experimental model from thinking mode supported models list. --- .../sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json index 569948031e..b03dd9bdf4 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json @@ -20,7 +20,7 @@ "type": "string", "title": "Model", "default": "gemini-2.0-flash", - "description": "LiteLLM model string. Supported: gemini-2.0-flash, gemini-2.5-pro, gemini-2.5-flash, gemini-1.5-pro, gemini-1.5-flash. The gemini/ prefix will be added automatically if omitted." + "description": "Supported: gemini-2.0-flash, gemini-2.5-pro, gemini-2.5-flash, gemini-1.5-pro, gemini-1.5-flash. The gemini/ prefix will be added automatically if omitted." }, "temperature": { "type": "number", @@ -58,7 +58,7 @@ "type": "boolean", "title": "Enable Thinking Mode", "default": false, - "description": "Enable extended thinking for supported models. Thinking mode is only supported on: gemini-2.5-pro, gemini-2.5-flash, gemini-2.0-flash-thinking-exp. When enabled, temperature is forced to 1." + "description": "Enable extended thinking for supported models. Thinking mode is only supported on: gemini-2.5-pro, gemini-2.5-flash. When enabled, temperature is forced to 1." } }, "allOf": [