From e51f9dd4fb473f50b3a77c723a07e4eb875f9bf8 Mon Sep 17 00:00:00 2001
From: Cosmin Maria <acosmin.maria@gmail.com>
Date: Sat, 4 Apr 2026 02:55:18 +0300
Subject: [PATCH 1/2] feat: add custom_class parameter to factory functions

Allow users to pass a custom class to get_chat_model() and
get_embedding_model() to override the auto-detected vendor class.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 packages/uipath_langchain_client/CHANGELOG.md |   5 ++++
 .../uipath_langchain_client/__version__.py    |   2 +-
 .../src/uipath_langchain_client/factory.py    |  24 ++++++++++++++++
 tests/cassettes.db                            | Bin 45494272 -> 45707264 bytes
 tests/langchain/test_factory_function.py      |  26 ++++++++++++++++++
 5 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/packages/uipath_langchain_client/CHANGELOG.md b/packages/uipath_langchain_client/CHANGELOG.md
index 3674583..d30a8d7 100644
--- a/packages/uipath_langchain_client/CHANGELOG.md
+++ b/packages/uipath_langchain_client/CHANGELOG.md
@@ -2,6 +2,11 @@
 
 All notable changes to `uipath_langchain_client` will be documented in this file.
 
+## [1.8.0] - 2026-04-04
+
+### Added
+- `custom_class` parameter in `get_chat_model()` and `get_embedding_model()` factory functions to allow instantiating a user-provided class instead of the auto-detected one
+
 ## [1.7.0] - 2026-04-03
 
 ### Added
diff --git a/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py b/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py
index 66e749d..5a6230c 100644
--- a/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py
+++ b/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py
@@ -1,3 +1,3 @@
 __title__ = "UiPath LangChain Client"
 __description__ = "A Python client for interacting with UiPath's LLM services via LangChain."
-__version__ = "1.7.0"
+__version__ = "1.8.0"
diff --git a/packages/uipath_langchain_client/src/uipath_langchain_client/factory.py b/packages/uipath_langchain_client/src/uipath_langchain_client/factory.py
index 4e98d63..163df11 100644
--- a/packages/uipath_langchain_client/src/uipath_langchain_client/factory.py
+++ b/packages/uipath_langchain_client/src/uipath_langchain_client/factory.py
@@ -87,6 +87,7 @@ def get_chat_model(
     routing_mode: RoutingMode | str = RoutingMode.PASSTHROUGH,
     vendor_type: VendorType | str | None = None,
     api_flavor: ApiFlavor | str | None = None,
+    custom_class: type[UiPathBaseChatModel] | None = None,
     **model_kwargs: Any,
 ) -> UiPathBaseChatModel:
     """Factory function to create the appropriate LangChain chat model for a given model name.
@@ -106,6 +107,9 @@ def get_chat_model(
             - Bedrock Claude: Default uses UiPathChatAnthropicBedrock.
               ApiFlavor.CONVERSE uses UiPathChatBedrockConverse,
               ApiFlavor.INVOKE uses UiPathChatBedrock.
+        custom_class: A custom class to use for instantiating the chat model instead of the
+            auto-detected one. Must be a subclass of UiPathBaseChatModel. When provided,
+            the factory skips vendor detection and uses this class directly.
         **model_kwargs: Additional keyword arguments to pass to the model constructor.
 
     Returns:
@@ -128,6 +132,14 @@ def get_chat_model(
     if not is_uipath_owned:
         client_settings.validate_byo_model(model_info)
 
+    if custom_class is not None:
+        return custom_class(
+            model=model_name,
+            settings=client_settings,
+            byo_connection_id=byo_connection_id,
+            **model_kwargs,
+        )
+
     if routing_mode == RoutingMode.NORMALIZED:
         from uipath_langchain_client.clients.normalized.chat_models import (
             UiPathChat,
@@ -248,6 +260,7 @@ def get_embedding_model(
     client_settings: UiPathBaseSettings | None = None,
     routing_mode: RoutingMode | str = RoutingMode.PASSTHROUGH,
     vendor_type: VendorType | str | None = None,
+    custom_class: type[UiPathBaseEmbeddings] | None = None,
     **model_kwargs: Any,
 ) -> UiPathBaseEmbeddings:
     """Factory function to create the appropriate LangChain embeddings model.
@@ -262,6 +275,9 @@ def get_embedding_model(
             RoutingMode.PASSTHROUGH for vendor-specific APIs.
         vendor_type: Filter models by vendor type (e.g., VendorType.OPENAI).
             If not provided, auto-detected from the model discovery endpoint.
+        custom_class: A custom class to use for instantiating the embedding model instead of
+            the auto-detected one. Must be a subclass of UiPathBaseEmbeddings. When provided,
+            the factory skips vendor detection and uses this class directly.
         **model_kwargs: Additional arguments passed to the embeddings constructor.
 
     Returns:
@@ -286,6 +302,14 @@ def get_embedding_model(
     if not is_uipath_owned:
         client_settings.validate_byo_model(model_info)
 
+    if custom_class is not None:
+        return custom_class(
+            model=model_name,
+            settings=client_settings,
+            byo_connection_id=byo_connection_id,
+            **model_kwargs,
+        )
+
     if routing_mode == RoutingMode.NORMALIZED:
         from uipath_langchain_client.clients.normalized.embeddings import (
             UiPathEmbeddings,
diff --git a/tests/cassettes.db b/tests/cassettes.db
index 7caa33b2c1fb6808b8ffca3bb1dbec97faf4d980..9ffafbcd10548f169e11bddc9bdfbe8291470d0d 100644
GIT binary patch
delta 7211
zcmciHdt6jy9>?){&YZb0$S{DAipn6~#sM9!3TTM8@{)iUUJyB46$KpyHO&hn9`B$?
zddk}ab}=gzL#5ra5o@h1ax^pZQn;nrT}#`wH23?1-P#qw{bM&@ulMJrbLO0xd7kH-
zBSUe+8Q<cDOTO#<HQ%UIR;An=$G&k$vczObkR`P&iL#`TC9N#!WJxbeZn9*MC3jhB
zBTF8#WRxXOSu)9zmn?b9l8-F4l_g(U@{^?pWT~AjwU?z1vg9vI9c8JLEOnNp09i82
zQlKn#k)<G6>MBd!WU0F>^*A8)th(c-I<5**<qB-hMuVlhSskl>O`{h#t8$sU{w@7v
zy-K%L7o|O|%@N*a^%{G5yq_{zZwXT*H;cg|J!?*3(aeJMto)3^LcaHwi*?p%KNohX
z`VUuTO|(X_48HflG@<<_gC*#$C;Q)Yu>g(sQi|x{8C`Y?e%xd8exbjzh*>@r#ignM
zrRscaN1oI*inUW3^p>7&|Kr6{^YaU)%nMOYys7q6oO+8^(0;|MD$lZez1HoOCt0LU
z`(?UabYxTn2KqLCKgy^0^cwAzv3NW2${gm$LwX++VszTiG<H?HRK?+I@lh9GW8k&|
z!LRxL{Mt1a3({!6n1bgx)~B<3J!jPo7we?cURJ48H~i_4dEC@@HVaU0Yb;a6f4#ZJ
zPy3&vrCiWi0yWyN({b>m)GE9?yUE3d)<)<o_4;8roR2z4B|a_ss)IbejrktO+I-)p
zo%;&I&xRgu<Mqe&BXn=;MrxO6YBiBU4m-s%Ro7Jy>C#jxN4i)NCluur78NS-B@vN|
zG{LIl>jP5QMDF?|oB7p#<!>ou%9xL1!xg*9lE6(nVg!>;)Hzbcl4$x2U3rqRoJA<>
z1<PP=UlYpa3!*`>Yb@)OnSH_|Hg&RSxyd($_1|7Th_x3)cfL0!O^Dc()~i21Y+lPu
zN>qU57@xHzmeGSN6_?uL$+Qadu`FtODjhn84;;Q*@D|j+C<}cp%lV*fu`FDv_O%S<
z#bZr^Q4l>8=4Wy7u=0Oy7_3yUjTqH*G^673vy}1^7rcZB+9{r2&3s;H%d5kdF|AJY
zbUZ4SgwThvD~TV)X0RB3ts!5Cr1!pk^m(>jCz>2)ap~ZuIdheTvTzp8Grpblug^X3
z=w%k76Wtv7VoBp-$y2V1mg41R%W`wCgJvtU^;H+EyoF$udbxXmVYS;$w-s*f_4&Ft
zgka%}cw3y}Im)=z<4cd2Hcz?V()HAC)jpt6REyQcwfz#UAxwjRycNp#bxdIeihYbV
zKsg-LN2lCQvi9Ls9}dLZuuZZWm12L(SYAEC#M<ga!O>k@nxL>r)&WhA7rTWrFLw)$
zq#Y^H#5XP_8%^sTv8}k&`rC6cFw@h7=n7MyR=G9Sx<yo}gkYsA#hSny0;5=PtzU|D
zH`6$Jh)a`x+bLBUm~M5i9pG;<3wN&<LDVUeI$18PxRq*!U<iTU&<AV~3Skfq5fBMc
z&=>kaG{itG^oIfPAjH8yh=)Ni7>2-57zPi)!!R61z(^PckHBb1fJ8`wF^~*n;ZaC|
zaWEbxz(h!eNstDUAsr;hfK14OY{-FJm;zHF52nF%$cF-$0W-l4vtTw9LJ`b?xiAkN
zgZZ!k9*2dn2o}Q<C<X^Cg%Vf>GAxG`uo70mYFGnnp%lts9ju3P*Z>vq1Z;#&pum&x
z6l{hquobq!cHmG6Rj>n|hG*becn)^L^RNqEfEVE<*bRGNFI2-mco}M7KOBIAPz$fX
zpWx5%D!c}V;B|1qVW@*6@CLjIe}SX$7Q7Abz%h6i-h+C0AC7|yPQXd{08YV&a2h^>
zGw?B-g>!HoF2E=7DKx-E_$zz{pTi}%3}3(%_!7Q?ui+c`7QTb;;VN8%AK*v04maQ@
z_#6BkZo<#-3*3U+@GIO|b1Rjt1ByUVQ$&h}qNV64dWsvxKyjzEp?FY?6i<qY;zjYM
z_)ywXd?|jE2Po|*?I|57{*;cCPL$4+0E(FsNa;ceqI9KnqjabAp!B3zD7`3FN-!ma
z(wowUVxxpo!YJXC2udU+iqe<Tj}lFZp~O=9QwC5Tq{LAMQsOCtD1#|OC_^d3C=XE{
zrVOWypp2x9qC7$wO-Z06Qj#cRD9Mzu%!z9X)9*L>sOoXU?mo`5tN4?#oSRH3%-lv5
zBZT2m{0SJxZS|K~1V23W2Akk~7n{$k*%%&Jb(VE;!-aT5(`6Sot01~AZ=f4x3{R*y
z$3m3K&K5sBs9wNX#ULER?)ed0!5gMZOz(WqFNKZK3D(uh#g{_Pap_9FpfgOsD|T$U
zQQ&|+biji=ArkjzrE;Rx#;bK#nMJGq6_1xIziKRJ%FHRO%sdXp!vvTJsW1uBU^1kG
z1R0PCS&$7mkPA~_D&)a5m=5_+05f1F*kKmThC(QUIWQOI!DBEV7Qo}M5Ej8=SOUf1
zfTd6Z%Rq+ZumV=XDp(C`U@eqF8LWf#P!1cQ0-k`4un81+5}tz1um!flHrNgvDxnH?
zz|-&yJPXgkPIw-6!3*#ryac;p5A20%*at5|4eW;la1d(Y75EeU8D53g;1IkHPB;v8
za0K3fH{maE6yAcj;T<>z@4|af5AVZqaKQ;U2_L{I_z+IRM{ou{hO=-E&cg-x1U`iZ
zxCnoR&){>o1ef6pxB_3oSMW7_1K+}T@I73GYw!d72-o2T`~-i4zr#)V8GeCVa2tMw
zJ8ZRCC^It<Kn)^jKnpt1gBuvY9om2g7{L=v;04~`18u<<{NMp-2koH)_(MnN1f3xO
z==B4k3j{$|=my=P2lNCB^a3jcLkRSSK461T2!n8lfJlgfzR(Y%AqHZhKMa5eAr1yY
zJPd-tFa(CeFn9<ahT$**M#3m~1V%#wBtjC5fn*p9k3tHJgYhr{CPJ$4x_OcyHt0`_
zw{StRmLF=9COoI+)klh$w$|%#&tc4DEh){zrn#k4(5%v&bw*8^fJ~;OQzS|TC6kgx
z$)@B`aw$_NQz?0rX_V=dd`baj24yD2PMJlSO(~=lQRYzQQsz+}qs*r)pgc}lNLfT#
zOj$xHrZ^}|DJ7I;6q&M|vVyXbvWl{rvWBvjQc5YKtfQ=_lv6fPDkx7-Hc~cG6v~s7
zrzo2#TPRy8+bG*9oKi`tqU@kNO?ig$Eaf@MPRjF?U6dCnFH&Bj?56CY?4?vw_EBD@
z)KK<Q4hT-8=AiIX2ZLIjtTGz)`_##`UgprDYJ=n0fYKbTBe|r4UmPFG?(u`RlYa0Y
zt*$T~DCFi1r`bE)e50N@j93-Z9LM^XW@{Zut?W(TF+9?2Z``-6*-j&VV*i2B-*cRG
z<t`Rytd+*ILyeE*I_s+F5ni2L&!mdtOWkS;_V;VtG{spLcmkW0%5Kcc?KdYgqjSJv
z7hBf2b*i(ju5m~6qxI|ye)y=1c`K%ro;q$$sb`1i)$(v`I)KxhgGXpDmC9Ln(Z#|v
z+tccq#=Cb!NI~9=ypS+kSX78DJS4Pl#jL*jD?DFbR#9CRZMu6a-*^5r>*oCNvWu<J
zy5~e?*=&)yw#<GRJY$<6OjL?9f(_jCT|Mi-_Z-@fA7hVPn9CeWgGbNy>ad7tUKgCi
zlKD)-Ddt}B-hlah-S;jwi{9w;X2X)-@~{-NIxI8pKP)r<|HBexi^|3amW~?kF)S>S
z2M$=tHVeNE$(<`_nY%JAGX!Ijeu5PW)>+L)WOmCVQrPN<6y1MB=KKL85}KJ6mB|OD
z75z6OqLh`1%ar)(kv?3SUMh6qnTw0@MU8zt(`pB)49{o=)p{lNi_+k{jxU{}b)>d(
zX3`g6W15|rbY9PM9E&Z@)jASd*_*x^8`W%YJ@&T4W6h>%e#_Iepw(%5{QlFl@DG@#
z@X)NV2yQ!+!m|JURBpaztG9az?n<AJMXcF`4`|IUZZ;T8S{{tzRtLjz|G`-L2Mk6;
zRA#h|$9Y^}%`a}vMnoy=tXZb)+Zg4eR0i}4Qa1GN9ow{Gd#DUsG(Bj=`qPR{F3r<A
zmb9|w=xgEZW^0aC>@vr(A*EBbj*OP~Zv4Wy*8aMBE#s#4_3TuIE5?KKH7D7TyYFqZ
ziWh<0Y~q%;JaH>pow$|vpSV^3&xtcM-j$*vGct1d`s>qJYO|ZtJ?8B@%$su`%jJo2
zVa$D#Ib7`Qp`WDkoa3o7t}}M?aC)fQjC6n1y^rAqLoc_%`bkc&#9cy+UMWttUQry$
z)(hP9-Z`dc!}-+@WudE*IKkRM**o6q`}^0D<1}=mTU!%4=(vYcnaVCyByz9XhVd+}
z@orvP6FTO&2e*$q!QTJpe*$6}H<Z<cW*+x&Rt~9WAp#v@oil#*`^+HlfpzsvtcfeC
zbFW$OLDRPNJmbu1_N}w}q>CkbG~V`|O}G7;&@kKo?4}>p*A|_V#rIw9agS@c5E-Gc
z>rqese#cjCFS5uijL(Wr6~b94_p)T*UyWnkbbi=W21R3|^P?-BAHzzgH?DByg!due
zmVW6_yX}?xukaE&XEr#FJzP3X>&Pjo`2C12|I)^}7r(){F4DQN1}a*bu_s!Zv5j}n
z*rrxyO!<RmtY1!UZY~e)6w5{^hpU4-{QI@N`Hb;9JNFCAxlbQE^X8S#i7bfkC``xH
IHBQ<;0D$IIRsaA1

delta 4106
zcmZwKd3a1$AHeZBckV3N7ZJ5j>|=zG2$CSST51U;XoRb3DWM^@qG=K?ibg`YwO;9B
zDT-R6YtnsDN?WZ8mA1O*8f|rXtGwTzzSaKm=6OEfXPz_Xp7YF^JacCXsz0w?P<^I$
zNqyHvha)hd)Yn~IEHq3*nTBQ>PSem$!(|$7)9{#v*EE7mBiJ-TOe54Z!b~IFG$KqR
z(lnw>Bib}-nMRCh)HaP+)2L$_x0ps<)2L?}^-ZIJX*4vAMyAo&G<>EJXBtgRqp4{$
zGmYk^(ZV!Z9yVH)T@P|rJK`KgO0`-S>~B%e`?`0m*WuahNp^pr3S8H;F|J(gd%eg}
zWJOiB>TLI}{6ZyJ+Zy?2IWzS_dwi1}T5Nf-^4C|^jCP3~t=xjvZYv|k-`A?SkQ%P(
zA@<Pl(!Z`S?1Y#l+8B=>>~J^+^|Hz`1IaQq+Buw+dRTK*ti?(Zc3RsbTEOG}tT0Dk
zoON?rnpJW>c`i#;xBb^X!>#Ks|0HW$WqGciJdSmSW|gT}J2&#A>QkQN@gMhgcjh}{
z9cjwjDJa^x+I>>@xt0EtYmt`kRH|V_=kO!pT^wm)W7MV4W1*u%z6t3Tyv{q-bJCOT
za1@Ty7vEShVSRS%Bs+cNL~Wj?d+nK<QdJT|%dN^LiM&d;HT0uY8A!7BFHH()y2sv^
zlC9l1s1!TB|6VQH<NRa%BE4|o&DWD?otzMeh)!jnsQTLMiFvd0+D#armp9Qqbw29v
zOH1O;sW?rKu-liiqnvw+U3qDe%C?LPscuaVw|i{)`w6R{IH|5j4=Z%)ML{>V**f@~
zO0g4|5#jDsI7naI@L$8}IrG(Mcjv+oy{PWZ;og?fE-}p2CFqjltkt6+sl8pZ>8R$h
z_or8>PIlBwRjT9vMz-~@QuXYT9}~4G`*L-T_Sai4J8t#pk%c4l#ql>+Em>RHGQ^$2
z+rDRnHFH#;)ZRZJQpISxW}k@3RIRL&4Fa{Snr;5+<<mw5e46FU3Y^pbIj!8<{!Vsf
zlkqA<wf*a0IqEdK(zi-QT8C@<-&*SPFZH!P;tQyRPHLi~jw9cpozef)Cr1nn-yC)+
zEG2Y9$kmYP!Jh_C54sw(Ca8@!*K^d<%DvfL$7MP4odxPzdFPBkymC1lDzCkLph2dZ
zZWZi~Z){aIZRfGVcE?w_3R~z!Lo78W(ADZbA`oNyu8dGamEX$G2z+d(AIwl*_kh9}
zy(n;VPCu*C**e6!Iy|t+t~qUL&8(7sZ6j|SU}L*FE?Kpy*tPMfHqvDe=_vm#1?&AG
z_KCx1Rdaj4r$+fT=e1>5vjS*?c(g@3B%nPyAQ4IEh-7p^XLLadQjvzP=!Wh{M-TKw
zFZ4zq+={;FhyECV+b|G=Fc?FSfuR_N+c6v?kcp8Ph0(YJSr~(CjKw$@7>_$K0e4{{
zCSfwB;BHLCG~{49X5b#=VkTxG53?}`b1@I|aW58NAr|32<l}x6pb$k^j3qF!6w9z2
z4`2mWVig`lF;=4lrC5WtcnIsT9u^+P25iJ8Y{nLBg^g_}!y|YUkKu7_#}4eoE<AzV
z*n_=z68o?p2k;aQ;t&qw2+DC3PvaRpi|6n>UO)vZ@giQr%XkH^;x)XEH}EFj!ZEy!
z<2ZqLa1y8RF5biYIE@eRA<p0<oJAGR;XFRZ1yth`e2UL-5uf7=e2K4c318zIe2ee!
zJ$^t9F5^f1gr9K*zu;H=hTriY{DG_Z6W4Hk<<%^;8WK(7l;{$d#4YhiypkYEup~qh
zDhZQ>OCltZk|;^Eq?RN`Qd<%$sUx{XQdd$>QeV<Q(ooV!(pch?#7UY+no62UnoC+p
zT1r|;{F2s^fTWEiUeZ?5PLd#LFX<pjlq5+yN|Gg=B%LK)Bq@?qNt&dqq?@F>Bwf-&
z(o@n)(p%C;a;v1Tq@Sd}WPs#0$w0{<$zaJ4Nrq&oWSHc3$#7M{XNvM3N{MzH=ho3;
zq<$(QBdj=dXUNfzh~VBqIo_k5XFSc_-Cgq>$DNm*1GJs$v})!^(K;}lzs4i&gyVc&
zu`8!sQKKu4MP{mVP8Dy*m7P`oAcsRME%fO{|9lZivl7mvwss{t9L@!{THBK=b0a=f
z%k1hYhVoW)i_KKSJeqdI`eaXhm2G^Lqj`cyvrpCky)(G{vcrvgV?gKO8oTUtI`?p6
zSYUp!FH;r!Mq(63;|^qD46-p6<6vMs?!*M#g^8Gi$(Vw>F%{F0gXx%odytEnn1wvd
z#vIJWJj}<vSb&9Cg!_<>`%!>G6k#!zz{FB4!*V=;6<CQ?co4-{jS`e%4c6iztiyU(
zco-Y75u30XTd);2wxJA<;88q=$FUtduoJuR1a@N&_TowG!+spVQ#go2IE*7G$5A|u
zXYeeZ!}E9n6{y6EcnL4#6}*bq@H*bWn|KSy@HUR)1m3|(oWi?!5AWkNKEQ`KgO6|)
zRXB(9_!t*ZjZg3?KEp+PjxX>fzQQGZjc@QRzQgzU0X4XcAMq1@#ufa6U-27$$A9n#
zuHsKz!*#X7rxp7YG&rHd1vfnKA_&0<K`6oyjtE2|3el*A7}Q2A>fjdCMLpC<12jY<
zG=}UShbCx>W@wHUXo*(vqcs9(gLt$>J0zezIv^2A=!j%=LT7YA3R01VuIPsDNJkI!
zL@)G4AKZ$*=!gCofZH$-gD@CFkb$8XhTAb5Ban%a7=_We16kS&zA>6!?S0P`&ZkOX
z!t<fo+ID9-zq`h&)se1jjf|CylNgfmk~<|6BzH+BN+wArOQuNfmQ0mQljKOIOJ+#!
zk>pBdN@hv&B(o)RBy%P6B=aTrN)|{KN)}1(ljKY8mlQ|}B}J0Ok|h#TvQ)B6vRv|j
zWQAm<WR>JWNwH+Lq(o9GStD61c}TKOvR-0I9+qs7Y?N%0Y?f@1Y?auOZIUv{Ba%lY
zk4YYvY?thi?3C=1JR#XF*(2F2c~Y`ZvR`sQ@|5JD<dEdBRuSnsqWxU2LQm+YLc=Q}
J`F&fX{tE$X28I9t

diff --git a/tests/langchain/test_factory_function.py b/tests/langchain/test_factory_function.py
index 50b256c..6ccdb6e 100644
--- a/tests/langchain/test_factory_function.py
+++ b/tests/langchain/test_factory_function.py
@@ -1,4 +1,6 @@
 import pytest
+from uipath_langchain_client.clients.normalized.chat_models import UiPathChat
+from uipath_langchain_client.clients.normalized.embeddings import UiPathEmbeddings
 from uipath_langchain_client.factory import get_chat_model, get_embedding_model
 
 from tests.langchain.conftest import COMPLETION_MODEL_NAMES, EMBEDDING_MODEL_NAMES
@@ -18,3 +20,27 @@ def test_get_embedding_model(self, model_name: str, client_settings: UiPathBaseS
             model_name=model_name, client_settings=client_settings
         )
         assert embedding_model is not None
+
+    @pytest.mark.parametrize("model_name", COMPLETION_MODEL_NAMES)
+    def test_get_chat_model_custom_class(
+        self, model_name: str, client_settings: UiPathBaseSettings
+    ):
+        chat_model = get_chat_model(
+            model_name=model_name,
+            client_settings=client_settings,
+            custom_class=UiPathChat,
+        )
+        assert chat_model is not None
+        assert isinstance(chat_model, UiPathChat)
+
+    @pytest.mark.parametrize("model_name", EMBEDDING_MODEL_NAMES)
+    def test_get_embedding_model_custom_class(
+        self, model_name: str, client_settings: UiPathBaseSettings
+    ):
+        embedding_model = get_embedding_model(
+            model_name=model_name,
+            client_settings=client_settings,
+            custom_class=UiPathEmbeddings,
+        )
+        assert embedding_model is not None
+        assert isinstance(embedding_model, UiPathEmbeddings)

From ecc4bcc53e3b38d93d1d131f4ca84845a8594172 Mon Sep 17 00:00:00 2001
From: Cosmin Maria <acosmin.maria@gmail.com>
Date: Sat, 4 Apr 2026 02:57:58 +0300
Subject: [PATCH 2/2] chore: set langchain version to 1.7.1

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 packages/uipath_langchain_client/CHANGELOG.md                   | 2 +-
 .../src/uipath_langchain_client/__version__.py                  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/uipath_langchain_client/CHANGELOG.md b/packages/uipath_langchain_client/CHANGELOG.md
index d30a8d7..88ae09b 100644
--- a/packages/uipath_langchain_client/CHANGELOG.md
+++ b/packages/uipath_langchain_client/CHANGELOG.md
@@ -2,7 +2,7 @@
 
 All notable changes to `uipath_langchain_client` will be documented in this file.
 
-## [1.8.0] - 2026-04-04
+## [1.7.1] - 2026-04-04
 
 ### Added
 - `custom_class` parameter in `get_chat_model()` and `get_embedding_model()` factory functions to allow instantiating a user-provided class instead of the auto-detected one
diff --git a/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py b/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py
index 5a6230c..3b6d111 100644
--- a/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py
+++ b/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py
@@ -1,3 +1,3 @@
 __title__ = "UiPath LangChain Client"
 __description__ = "A Python client for interacting with UiPath's LLM services via LangChain."
-__version__ = "1.8.0"
+__version__ = "1.7.1"