From 2e0750e4654835aebc6b4ca16e7de58716bf8103 Mon Sep 17 00:00:00 2001 From: Kangle <147648913+RCAKangle@users.noreply.github.com> Date: Fri, 28 Aug 2026 12:55:15 +0800 Subject: [PATCH] Add local OpenAI-compatible server support --- README.md | 35 +-- __pycache__/__init__.cpython-312.pyc | Bin 0 -> 384 bytes .../chat_optimize_nodes.cpython-312.pyc | Bin 0 -> 18566 bytes chat_optimize_nodes.py | 121 +++++++---- tests/test_openai_compatible_provider.py | 204 ++++++++++++++++++ 5 files changed, 303 insertions(+), 57 deletions(-) create mode 100644 __pycache__/__init__.cpython-312.pyc create mode 100644 __pycache__/chat_optimize_nodes.cpython-312.pyc create mode 100644 tests/test_openai_compatible_provider.py diff --git a/README.md b/README.md index a798454..6fcdabf 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ Local LLM chat nodes for ComfyUI, with a clean handoff path to downstream prompt ## Features - Multi-turn chat with Ollama `/api/chat` - Optional Hugging Face Inference API support via LLM Config -- OpenAI-compatible chat support (OpenAI / DeepSeek / Qwen) via custom `base_url` +- OpenAI-compatible chat support (local servers / OpenAI / DeepSeek / Qwen) via custom `base_url` - Anthropic Claude support via custom `base_url` - Session memory by `session_id` (in-memory only) - Optional `system_prompt` @@ -15,10 +15,10 @@ Local LLM chat nodes for ComfyUI, with a clean handoff path to downstream prompt ## Nodes -### Chat (Ollama) +### Chat (LLM) Inputs: -- `model_name`: model id (Ollama or HF, depending on provider) -- `base_url`: Ollama base URL (ignored when provider is `huggingface`) +- `model_name`: model name or model id for the selected provider +- `base_url`: provider base URL (ignored when provider is `huggingface`) - `user_message`: user input text - `action`: `send` / `regenerate` / `clear` / `deliver_to_optimizer` - `session_id`: conversation id @@ -41,7 +41,7 @@ Behavior: Outputs a `LLM_CONFIG` struct that can be shared across Chat nodes. Fields: -- `provider`: `ollama`, `huggingface`, `openai`, `deepseek`, `qwen`, `claude` +- `provider`: `ollama`, `openai_compatible`, `huggingface`, `openai`, `deepseek`, `qwen`, `claude` - `base_url`: provider base URL (see sections below) - `model_name`: model name or model id (provider specific) - `temperature` @@ -62,15 +62,26 @@ Notes: - Many public models still require an HF token. - Chat history is converted into a single prompt before the API call. -## OpenAI-Compatible API (OpenAI / DeepSeek / Qwen) -When `provider=openai|deepseek|qwen`, the Chat node will call: +## OpenAI-Compatible API +When `provider=openai_compatible|openai|deepseek|qwen`, the Chat node will call: `{base_url}/chat/completions` Notes: -- Use an OpenAI-compatible base URL for your provider. -- Set `api_key` in LLM Config. +- Use an OpenAI-compatible base URL that includes `/v1`, or enter the full `/v1/chat/completions` URL. +- Set `api_key` when required by the server; it can remain empty for local servers without authentication. - `model_name` should be the provider's model id. +### Local OpenAI-Compatible Servers +Select `provider=openai_compatible` in LLM Config and use the base URL exposed by the local server: + +| Server | Example `base_url` | +| --- | --- | +| LM Studio | `http://127.0.0.1:1234/v1` | +| llama.cpp (`llama-server`) | `http://127.0.0.1:8080/v1` | +| LMDeploy | `http://127.0.0.1:23333/v1` | + +Use the exact model id reported by the server as `model_name`. Ports may differ if the server was started with custom settings. + ## Anthropic Claude API When `provider=claude`, the Chat node will call: `{base_url}/v1/messages` @@ -83,12 +94,12 @@ Notes: 1) Place this folder under `ComfyUI/custom_nodes/` 2) Restart ComfyUI 3) Add nodes: - - `Chat (Ollama)` + - `Chat (LLM)` - `LLM Config` - `Chat History Viewer` ## Usage -1) Create `Chat (Ollama)` and set `base_url` / `model_name` +1) Create `Chat (LLM)` and set `base_url` / `model_name` 2) (Optional) Add `LLM Config` to select provider and set `api_key` 3) Set `action=send`, click Execute to chat 4) When satisfied, set `action=deliver_to_optimizer` and Execute @@ -97,7 +108,7 @@ Notes: ## Screenshot ![ComfyUI node graph example](LLM_embeder_nodes.png) -Example graph: `LLM Config` drives `Chat (Ollama)` via `llm_config`, and outputs route to a `text` node plus `Chat History Viewer` for full history review. +Example graph: `LLM Config` drives `Chat (LLM)` via `llm_config`, and outputs route to a `text` node plus `Chat History Viewer` for full history review. ## Example Workflow ![Example ComfyUI workflow](Example%20workflow.png) diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a39b7f4edcac5b557934f1587aa2622e734bd617 GIT binary patch literal 384 zcmX@j%ge<81o39InZJSbV-N=hn4pZ$3P8qmh7^Vr#vF!R#wf;IrYI&xh7_hK=5&S> z=0%KAER`&ptS=dX$~2j73H$lGxW+sCI0gsD`#J^$c>1{q-;x3gx_AZ$_&7$!`#Ji$ z!sV)1_4LbAlQfx&n1Kcqu>gr%yy31+@h+Z0uFfI;L6Lr%tha=dGZIVU^9xEcb2F<_ zJjd#sWN=->EiiwZU%*!l^kJl@xyv1RY to1apelWJGQ4Rin~K#E0x#0O?ZM#lRL0Uy{H1SJ}HA8`dVvKMgzB>>WyA{~ng;1i5;Y`~=weF6+W~RDMJCQRt8FzA>TL|=Vy;i15>pz;Ac802albqX` z_V@c13y>6OIi223KM~)4kKgy)-}ifdt6$k{77D_sc5&>zZi@OZm{Ef}pLn!OOHoVI z1&XCudYl?mp7bD1@|r;nJT>Fm3EiM>LO-aVFbo=Kgwc*0CrpDTlGcrzCoF>&nu<_^ z45#C)oSwwiVNC@LgFhw3*;xJ4kYf#lcGkEaW>~|T-qsB|STkpZ8tbr*wLsp<(k#u{ zhqWvN&`D>c*!imiW}VjDixz>9(~SVhVjp{%J;R!z#9 zp{%7)R>QUuOsQ3H2v*YuJ=SrSvFAuf8=|qJ8?q!b*7JB78~l(~O4)61Qy=C51{=95 zu9~Z1x8KwbHnBU{ooqWvH?tjsEh7}W>nHSJ>oq#K`?uJlFQ}Du-I1ZVtm}=$;(z>vho)ok$%#;O zlI3D*VW_wFOsHq#6324<%VbTHu=I(@EOd)-?Nirav6rSJ<7{Y@3jgU51N7w4|Vhw`rl3Y0A%+prHxmneZ6TVG2+sKU!jU_C-_3M{R@ zYW@m9=r>hEb{zx<5XQ$LR0K)T0=x=5;7=Gy4{%ymS4m+Dz6x@B(&{Vb&(fFMD2m6H zI#!P@S;Lt5@wo_v^*m`e0AY-?X-fb?f1C$A9lsVB8t1~%Kzx)7z$*+)PRFOFwc;A6b<$;`p@dX^lEj-0c8UN14uWZA;ORi-w})TX9am}bdSk*ZHM zBu~!wdWA3za7EG7IF)lRCkIVWA*x+&5Gyy~-Jc4VDaU3h#CgCmdAv|_KNn|6OJ`Pqk zWPsJE0TG!TpV6;4K1JkbHM7tlVGw9RgUEYRn=nSGgh|l8r5ChOqS$`ATXG$=RkL7}b#k{cYoiI;A>c z7R+NvWi|!1PneZ?tk7=}NUcv;1j|@iv0P2DrU<2SqmuPz)}q$Kkz@K~x?|Sd5uS>4 zAZa{)OUu_lUwkb@vH?(<3s3NMm@(#G-z}|d%#%r3$A$sh>q&L~oy*1{h33mt&ew1n>e`&rVNF#bhTRjsP|dPx2u+N%85JtR3Ov zL58me(Bj<4ocI8ik>Rz7e#WWrweiU?%eRnj;W$t6K}g1MBayP?L*bl-$HzmH#%@6b=gdhl?&}>39jV5gr8Zrawd|NXCi!bJ2k%U0?LBiR9@w38C*S~z zMxRt!J8w)H795`%J(9_B-M(N?CRR+%zcJdSnkJ|+XHAu#LwCVntU4C+zy6g+uQ?1;vm&xNtCeGR5bZhE=8| z%e17gEjQ$t{i1F^IfWZ|5h5-*OlfkMMv$^o-P9{7I5HYRQ;HTN+g1@Hq1^%o4q5?5 zCumtc;R59XUFwLj20`~D>c<)t{b}mGE@-3q$~MNw+Pnm9Okdf?{8*dm(lZjYr8ug3 zq<%!P7SfY;-Hd;zqe}P1GDRrXT8a;^<9!8^0q-Rr)KTFO+WKT2KRM3Hy6G6l^VdISrPjsPR4{!$=W3fb0$wK% z-eamHYuWYQh2GS`^ojH#v0~>Z)^^EOmVDuRuc|c%(yi$Vv10ot)*V>$(mh*1t$Hdw zla7cLyFRh*evt2{`~9;w&gQCjidF5IOUu4xvsiKX6YEoJ<&@LE=B4zOcY5FIO;)ey zYV(IIc_GJCtupmlras3sin>OHv+3iLm$|6Sz)P6oc%Vi+Cqvm(XiBzEgs+C8+)Rjo z08@iaMIJb0f;~_}fvNor7(B>?hO2Z;4{9i!Mm!k_Qyk8x5>$Ku#z6BXRwK{|A(&Y$ zVM+yAQ7D|JgasHLlV^Agt3xj6x%>{;M6Kqdlc($)(c=mXK-3j)%D?gnAv*Wi5 z_A$idO##?levkRexKFKo60F+h+q7U$I5v#zSNl*YA#W7yuzz-z!PAs*u-1fgT?_)90{l397Be^DG0??+XZJFc^V1Da+U(A9a|&& z+!Rz?1lnZpo82)Qk59!8cc754GZG!sw=b`Z%g*72qBC>Q2A zJ|Nq=rsJcNd}NmBdwe&H!J~4^ABRXbM4}+f$7IV?m=8~Iaa6imu|f|NaFQn?wn7B$ zBU}{E9><1I=L(h)p@GmGSsRJQ`3~%K7e>1=+Jn(vh-6D72C4uM$sr)EET|L;9kiaG zj-mwGg9FLBany#jAdB!`Y@{9Mq6$SUQ>2Y!*dL*d8<)VJzzWfswqe`>5-px|C>|)>eQW?+`bdyH!iL)FK#K$d~PgDHr>=I z3QOAuZJExMhQ0Ugm8pI2AG`rPAXl{`XWuEBc0K?lqx*gV`hzA=Mv`3%GxNF!<-WPT zVvST$^S=LvUu@X3d^A_lxnRF z$YiWy#eS$|Dpr}gEK`?b>P21s5`IBLBfDD z45kv)zfqlm^QP+~SkE^?cD9Vj2MOLf9S-A=*)ZLGdqP*YJy>RDiY(&-eOLnpLJD%QU3-=9pGRr%b$+5X;+Ej6vA* zc?S5T=xWX~&7!VZ;f;jhUFNQpW_aU-;T74NK&d1Zm>zB=2`Zrhrl$d>2fAFI>1hZv zqe1jjbUP5t2xGGp7@J&dU@gfc%|3y78*5jY-58ocHU$-X`l)O7premJ3HAIbh>r73m@Lq1*@D8FBB&8A z!=HhA&{Rh!BSTzFVK&>a$##rT_72*~IZ=4cSuEX(&G6^2_8S=Cf#T0&)Q=I7s~zhU zSEL>#S;l`wzk-K5y<=*8PlKX?|*+y*|1yn!L7RufOY=JAv$>_2#}6W8i~t zN-i*XEpA(CUu;iTWRB)s9iV+Uyw_h}cs+F}eJSS%t~xrhj*gsTw`kh^fU!N!1Pa3Z z?_n#KC{(K9ry+_ppdKjA+LG*`U;@pnz73_B!WOO9sod%r21WW1DCh+$3e(94p}<%G z=CqDt^`#VEf$pZR>xMKV)R5*CSXDwN=wkGHpsp`!-ZSf25;^weez(;Fh@ z(mYg1wI=ez0DcR8u^EVfbJYgsdZpIj+$lxbWbD^F7dn#%-aV2Uy~ngjR{yHCE^Dnz zYkz9|z?d;F@5u!Yd}8fCccjvj9s0kkGAO`F3tb`DZHIOKu|E1wouih(#}+UQ*m9vy44{UTRQ&`{r35EyyQvwWS^Ld` zSOvKK9GYbWTO55cBoNgaMKXtaN?0e@gv;`(Z5dleBnm}}C}UkDrw5rrFO;!v%*Bzo zk}!s*_Dw-B3zoQBEgVC;%BCQcsUw4w;3-H6F2OZcP&ML z5(IbraV2oO;4b+N%GKF`lu%F*p9HJ60V(0DM=jf;r-D(8ND0n$y*T4Y3rJAkJKrfd zK}vA4UZOaA)}e(;y|#LEn1u+S+ktcZW(`X0?z5OL!q~ zzPb|?UXXBSEqiwD*}HSsft|Z{&z6DBZIqv!iVW=pFDbf46?*s-46D#X#GarAJcQu$ zBu<`NnxZl7w-m5opp*q8JO3&+B-FGZuM}kzBAS58TM+scfdmQRI%4~E<&#aKN%_|n zO$KqiqLyn@2|Dq)@hsHJF0j~7gyX2Dt6NCOdvWQT#2Q&Y90Bh@45g$^GL4c|BpbgX z*%WrY*n1Qj{%iQf{*p*FpcUK7x6)*iudW!IB)e-tE17NAYZhw2U1qLRLd6VX~%bsj*}3gb2(;7x+AIVEJk;@LX0`Gj}F=SlY8s)OkO1 zRD2C%r>Sz)RF4XGdBw8ft})yCyx4LwYd!V1T8-JVW~7{z;BY{RtL478BGrDcvLkaW z=iQZ;y&i(>RW4NrVM?Z~sUok%Zz+iV{VL{PPv&y^QuLx_UVon{PjWX0R+t9K?O!^y zcqsJ(7=YbvtM2x!yFKULC0ci>A{RJNl^X}42+LoF*E^Kw5|E@QG)1G_^9ao$L?TlINlFzr!zv@;ZnpU~mD0iRuwbo0eb6`i`xbx~2NIxBHTZvuJ^K z&GjkLCNwDomBT!mf}$mgqwto~poY`3^twABI#x5NXSIU{RySy5^@Aq3tzkZXt?KsDP2CLXAP|JCPG8O1Adb7eUwY9#3FgV&dpg^9!N1*yckkY} zci(`b$7y8q1VBg7&(5RV7}TLx8vP(Vrq;vN zF>!{XV6IdQFIIHdsPsrSgon^XBO8yNJ=@#U)pz_BEnD#F$+5G2$4{OBVYrA(4xSqZ zyctz**T~NBGz>{50d_yIjLH^t|AQ-aI5Gmvg!oM(JowU0iW@dOPOi?t7}}xnSPUp2 zmB&ctWZ$_91EGNz&-L`rdepbrzOM9d0XQ}f5ttH9RG=*rT)rp#b)-Q^|BVz4yQp2`8 zO&{;c?&uTi&q_@@m(6$fiyd8J<58)(eYxUJx45fIY&t5{?tI7?D(nv(cBf^HvYRag z`{CeSfjN=rmPU3JJY0Cfrr_QH2fksroP=%qlfC1sW00W+&l4BXxd2N|3V2tZvtXsf<$Eh%ujfKkz zSMYX^nKcv9%nqJWwZuY794w5Mhf(|k6qO;$#MM96&4L=4@V?gtN~K5eqQL{^R3W(dAH!Y|<>U#gc9Grtu9*TAx`w2Eag=LG z07aF7cVrhFf>W^OJqOh7+7I(MmZ*qcfOd6on%wI=lpxp$HnrXDz&hw>0lJ=yvR*md z8omYm#Fc^_lwG?}$(o7SWv7Ivj1X8VMS zjhHRCiYHu9E>i!0^q#@LPTq4P6|c^}CaQ_z9f>NoNvI;e`HC0e>Ug5+jjH#CmFEmS zLtWLrNX^i}W)Me|dqYIXkWJwBgBz+KD->9dYyuaaa(%!8sr>B#{$&`gXqYCxP(A@= z!2QtdtVGBd6uk-G2a!K`qFs+-#-R99x00_uLWslj8VTr_lB%4@{n}Umcmu20zX1FLw#YIi^I1XSK8V^sioXkY0 zC*V#9hrZSz%?F@v*7g#z=m7ERz7hyJ6txwlYa%NbE#*Y)mi2r%I>PZ6uv-v%P{rn0 z{y79kFDNl(^QifeyIQij=jsqgt`!EWQLGyuF1~@U!Tk9@!k+#ZqaR>Y>g{|BE6{cU z-e&fy>=*}EbSz#FN)#pcIszTWh`a=wDzXrlwQ?Oz%UvCcRRf5k-Ik7tQo$zyYSvZ| z-EaeKAZWz?5-2~&Y<|j9J#UiS-Y}^v2>F z>9OUuobRdm=cS6q_iJy|W~|F&xr(m&Q;*db;2Uzj{p;Z`)0mMwqbHwKp3c@k zleGVasZ5Ri%$}`%_P)0%J-A$x^B!3^@wvBY)!UZ!wq+`p&*!`c=TAIf%97u_$5a=g zU`Th>Zy5il?#k3iI(%d7{pgKo=7n7CzB{ISwa?wJ#Ky1WDvu?NpSyi2!@FGG)hiCUX9Io@APH^CqHqXl5FmzaL?9w-&cLJKkfg~E17R(YucAJ z%SZpl*O}D9sHtO%&#bzevhJqz3pw|WWhUo-N~&wk=rZoq8_DNC^;A4GQvQP=PTM`A zZ|6PRPB2c-pS|y^OPl`27sNiiOP!0IsRQ3XlD?dC@0MDEKOOmCBon#g%C$V5JPSR5 z1;D=z1@Z>bTK$=?4jTC?mu41cQdiQ~a=zWGzNfRkr*pn%lRBxg=Ka3DiLb2KxPi7dy(rl+$`_zPf{0qC;4No?LPx(|rf0}uR_^R`bJcgmD@-LS6u z+p_+)j5p`swS4@Z|L~e-o9zIuH*wFs<9^ctc$1PpfPYOLa6>?9-4Cw+dx1R}_6vup z*}HgV%|&@@#rlD4`M~|A=Ljscs9WD+J!IN{uRM6a>FKWo^iXy6$rF;VLG-mtyPr-O z)3&U?UGi0nzBXz5F31~U(vq)P^zD{*AIaC>Z`!{+de`#t^CEkBrD;5QMt#jycWbhN zvn$SXun{fo@7a^HSyRJHovZGHS@*%5`|y0XML8(bW9=FMO1z=6_H$ z>}Es9vNkg<2KN84?SHOmq2$XLifbQN9@U@TL;Y%xx!0oo)l>Rjz4ov4dPoNi3dKid zmGFb(1ARe13f3so_|qh&R!Tm^|9}4cMm^+J?Jh? zPLA_vc;Qj-B*cm6o21CNdS z8HgO^bI(hT8i-uJxznIY`f8-gMyaMw3N%ag%~EslVfpK{zU6_(x2DGgcq853H6uxy zC|~88nIvJ`J!=d}TB&mXnvEpwl)q}tL6T0YvUaVEBwbWh-I|*uJydy3YGASN3o?kZ z6kNa%{0bqjpiz(ZK+zv!gL#`8G=mkzB-_9M0N=oX691Sf1^r04p^uQCqpk_Y5)Bwj zz~M*CA7IvD4TY4Eq)g~q0b>b#T6R+hCKMA3n%TO%gLSgt+<>1_kA1iceQZv7Zh~0` zj4S1RJX!;SM*jQo;BR1*!sv$>{RpF*82uO{f|Uf0XzSOo>?%aE{&?@%u7MFqLU?c_ zpP;}`O2XQak=H$cdg3>;Zyx~}1bcyNrjaN_;j6&8G-*+09lNn8@~vd(zXrUJ#bCJWgGZVLWg>} zV8mOUEPc~FjQF7G=s&9Gb$DWtk>Qbn!Quq&UfYJjENHZQ#o=J!*aS_Y$A82w9(koNJYj`O!M`V)+h zEb%yv&D879u?(Z{;Y>Opk*8j^0`-E_S}o;iMC#Rm)XSHrUXOH8rH~Kop)~n;8fyL! z@-b*;je}sYh0jdsLFSq{$jD_V4L+|}pkIh^GaSFcE(_)=g_!usP>Nj^Kh0tdXoro^ zJV>rdAvOiIE%A@2hIqn!06-*LEM22Mt`J|gi?ASfs|69M_@q$Jq)I0 zVjczqz7spr$eQU0exw)!ldpom8U=rGfc1Ce_@6_;toQLH;T(RBYw-dA81Q{HX=Cb| z=nXE9EDvTK&xxkz$m{NdS>x^nty?0XY*UP`C;LOaC;JCV;^Ut|7i2D*k@DYT1x9~{ zlQ|5DJW}clNC_$6o8^^rXC!wmVrJb&zUrMj{e_O!cYz&F?Tp++-GDiEs3Mzt5yn}f ziZ}M~^kBe`0!!XC+2ZrWqJd|N4;X>LsivGQAOLHh0AXZJPk>l=2Vsj*;1ayr$2xOg z=~V&4Rbf?w+!nps)%x3x4^hj7K>Qa9zd}x_kWs=o=P-*{q}Zs@(V!Tr$$7ktS+rIk zC!{@S`Z~HTG67()D9b3K+CEGzW2q*851D-Rin5LB#Ih=%`YH|#{+Cem3q(4k`E%5L z#{KSLJ-=lM|n9~%EOs}L_SjF`3fWhZ= zI14|}ENQ>1*8}DkZ>o;zpr=hiZI{r1;ZIZQ6Cjp~=Wzf&&7f7oZDAa(G7h<0qfO|Z z@L2^gYhclZTQL1pC-q)G{OZAog{PPVFkb!pR)c5c8M* z4Luk`sRTTW=C41|f(cCD$Nys(k=&$=d;vQQACSZMg80Q=&`fUP!zVGZYtf;|qzqU1 z;~}u+MB}0H$;m0%!B0oQ%@yLVazo(v;t9&}$irkCehXMMu6%$)3_uh#$j1D)gXA6o z+}nqq_GQoo<09(byj zdKY`Y-?wUSfD40+tMCobT=$yBtUpRi`wy+|KbhTs^5dqJb1#dpa=HD(s~y8x(+<%D zABFF@?{KX;0$E4kLCZm@Y2QP=)@FaC)gzzNgIC{7IU0Wtg5<6d`tQE50~&))2)!0s zrCV?B&BQ)DDAw=0)4$^A`oz?QN+N*#uh==5<+nXYTj(FqUE6d^en>4>cC}LyOrIf8nPg{fi1c)c&GYkLgA&q``PcNCT01kiE=W zkOv0ZKr#!?qQ5er=qs~sQl1w%0*cHrt8eeX|0_T%yHP<59qaAt?+=~nI(H5{v;A@< zDeOMkf3COd#Ski)#X8$Jdya;>PoD2NHUN5Nu#Nu@Fu3x?21;DA0j|ry=Ol`biHD7- zoPa}-N%Ze3*kJ&NE0_d%JlZ^-EJa0ighJtHbP|vczb^(o*m1t+8y9-|!3%owOb=+Z z%B>l+Yx8F@I*8Fxh~N?md@veUY>O{p>1!AbV-&@R$7th@fnD^{7zbaI{1WP9Q{fX6 z0}tOc!)@YS*aT6H7ct6WbPpog2-Ad5J3^aSy=;0~QRkoKe+^)OpHwUgWC+lJrtees z&nV_I%JHA5)-2WfYs&X)st)jJ@BTS^QkU#cHKpR|V;Oh)`5SL!E8E5L_N=obN9}vy z^2|BbbTqvyrAzgv8&el=*dI}lT{GBdXX@A*1&{QxNBCH40lipiY6H>sq0L7hknCk^ zT6jKS>}z^>!f^1}#lDm_)t@pgKAW{StQoP`1YgG{_b%*QGh@yIMvzqZ;TxCz!(_^K(G3zz zLJLna8F-eO3>KRxS7oYsv2M+bIru={lk8bIvc_Nz3V#wqa&0 zwy1MtFlSZgXu})`a>^VXm{VqjDP>leI!~ME9y&dcY0AWwk7cgjYR@(wdPHF%=q4Lx zcQ8LYc#_${Ge0{_m6#nC7iWjL;_NV2oE_$r*x1gQZCW7EDaShLlxE tmDfuZ?NC$;MU}9~GXI?IKU>bdWh6qy{{blz9drNy literal 0 HcmV?d00001 diff --git a/chat_optimize_nodes.py b/chat_optimize_nodes.py index 594b858..4a70467 100644 --- a/chat_optimize_nodes.py +++ b/chat_optimize_nodes.py @@ -1,15 +1,24 @@ import asyncio import json import requests +import threading from typing import Dict, List, Optional from aiohttp import web from server import PromptServer # Simple in-memory stores for chat state. _chat_sessions: Dict[str, List[Dict[str, str]]] = {} +_chat_session_locks = {} +_chat_session_locks_guard = threading.Lock() REQUEST_TIMEOUT = 60 +def _get_chat_session_lock(session_id: str): + with _chat_session_locks_guard: + if session_id not in _chat_session_locks: + _chat_session_locks[session_id] = threading.Lock() + return _chat_session_locks[session_id] + def _build_headers() -> Dict[str, str]: return {"Content-Type": "application/json"} @@ -277,6 +286,31 @@ class ChatNode: refresh_session: bool = False, auto_clear_input: bool = True, llm_config: Optional[Dict] = None, + ): + with _get_chat_session_lock(session_id): + return self._chat_locked( + model_name=model_name, + base_url=base_url, + user_message=user_message, + action=action, + session_id=session_id, + system_prompt=system_prompt, + refresh_session=refresh_session, + auto_clear_input=auto_clear_input, + llm_config=llm_config, + ) + + def _chat_locked( + self, + model_name: str, + base_url: str, + user_message: str, + action: str = "send", + session_id: str = "default", + system_prompt: str = "", + refresh_session: bool = False, + auto_clear_input: bool = True, + llm_config: Optional[Dict] = None, ): provider = "ollama" hf_token = "" @@ -337,49 +371,46 @@ class ChatNode: response_text = "" if messages and messages[-1].get("role") == "user": - try: - if provider == "huggingface": - response_text = _call_hf_inference( - model_name, - hf_token, - messages, - temperature=temperature, - top_p=top_p, - max_new_tokens=max_new_tokens, - api_url=hf_api_url or None, - ) - elif provider in ("openai", "deepseek", "qwen"): - response_text = _call_openai_compatible_chat( - base_url, - model_name, - api_key, - messages, - temperature=temperature, - top_p=top_p, - max_new_tokens=max_new_tokens, - ) - elif provider == "claude": - response_text = _call_anthropic_messages( - base_url, - model_name, - api_key, - messages, - temperature=temperature, - top_p=top_p, - max_new_tokens=max_new_tokens, - ) - else: - options = {} - if temperature is not None: - options["temperature"] = temperature - if top_p is not None: - options["top_p"] = top_p - if max_new_tokens is not None: - options["num_predict"] = max_new_tokens - response_text = _call_ollama_chat(base_url, model_name, messages, other_options=options) - messages.append({"role": "assistant", "content": response_text}) - except Exception as exc: - response_text = f"[chat error] {exc}" + if provider == "huggingface": + response_text = _call_hf_inference( + model_name, + hf_token, + messages, + temperature=temperature, + top_p=top_p, + max_new_tokens=max_new_tokens, + api_url=hf_api_url or None, + ) + elif provider in ("openai_compatible", "openai", "deepseek", "qwen"): + response_text = _call_openai_compatible_chat( + base_url, + model_name, + api_key, + messages, + temperature=temperature, + top_p=top_p, + max_new_tokens=max_new_tokens, + ) + elif provider == "claude": + response_text = _call_anthropic_messages( + base_url, + model_name, + api_key, + messages, + temperature=temperature, + top_p=top_p, + max_new_tokens=max_new_tokens, + ) + else: + options = {} + if temperature is not None: + options["temperature"] = temperature + if top_p is not None: + options["top_p"] = top_p + if max_new_tokens is not None: + options["num_predict"] = max_new_tokens + response_text = _call_ollama_chat(base_url, model_name, messages, other_options=options) + messages.append({"role": "assistant", "content": response_text}) _chat_sessions[session_id] = messages @@ -396,7 +427,7 @@ class LLMConfigNode: def INPUT_TYPES(cls): return { "required": { - "provider": (["ollama", "huggingface", "openai", "deepseek", "qwen", "claude"], {"default": "ollama"}), + "provider": (["ollama", "openai_compatible", "huggingface", "openai", "deepseek", "qwen", "claude"], {"default": "ollama"}), "base_url": ("STRING", {"default": "http://127.0.0.1:11434"}), "model_name": ("STRING", {"default": "llama3"}), "temperature": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), @@ -482,7 +513,7 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { - "ChatNode": "Chat (Ollama)", + "ChatNode": "Chat (LLM)", "LLMConfigNode": "LLM Config", "ChatHistoryViewer": "Chat History Viewer", } diff --git a/tests/test_openai_compatible_provider.py b/tests/test_openai_compatible_provider.py new file mode 100644 index 0000000..e49a370 --- /dev/null +++ b/tests/test_openai_compatible_provider.py @@ -0,0 +1,204 @@ +import importlib.util +import sys +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] + + +class _DummyRoutes: + def post(self, _path): + return lambda handler: handler + + +def _load_chat_nodes(): + fake_aiohttp = types.ModuleType("aiohttp") + fake_aiohttp.web = types.SimpleNamespace( + json_response=lambda payload, status=200: (payload, status) + ) + + fake_server = types.ModuleType("server") + fake_server.PromptServer = types.SimpleNamespace( + instance=types.SimpleNamespace(routes=_DummyRoutes()) + ) + + module_path = PROJECT_ROOT / "chat_optimize_nodes.py" + spec = importlib.util.spec_from_file_location( + "chat_optimize_nodes_under_test", module_path + ) + module = importlib.util.module_from_spec(spec) + with patch.dict(sys.modules, {"aiohttp": fake_aiohttp, "server": fake_server}): + spec.loader.exec_module(module) + return module + + +chat_nodes = _load_chat_nodes() + + +class OpenAICompatibleProviderTests(unittest.TestCase): + def setUp(self): + chat_nodes._chat_sessions.clear() + chat_nodes._chat_session_locks.clear() + + def test_module_loader_restores_dependency_modules(self): + original_modules = { + name: sys.modules.get(name) for name in ("aiohttp", "server") + } + + _load_chat_nodes() + + for name, original_module in original_modules.items(): + with self.subTest(name=name): + self.assertIs(sys.modules.get(name), original_module) + + def test_llm_config_exposes_openai_compatible_provider(self): + providers = chat_nodes.LLMConfigNode.INPUT_TYPES()["required"]["provider"][0] + + self.assertIn("openai_compatible", providers) + + def test_openai_compatible_provider_reuses_existing_chat_client(self): + config = { + "provider": "openai_compatible", + "base_url": "http://127.0.0.1:3762/v1", + "model_name": "local-model", + "api_key": "", + } + + with patch.object( + chat_nodes, "_call_openai_compatible_chat", return_value="local response" + ) as compatible_call, patch.object(chat_nodes, "_call_ollama_chat") as ollama_call: + _, history = chat_nodes.ChatNode().chat( + model_name="unused", + base_url="http://127.0.0.1:11434", + user_message="hello", + session_id="openai-compatible-test", + refresh_session=True, + llm_config=config, + ) + + compatible_call.assert_called_once() + ollama_call.assert_not_called() + self.assertEqual(compatible_call.call_args.args[:3], ( + "http://127.0.0.1:3762/v1", + "local-model", + "", + )) + self.assertIn("Assistant: local response", history) + + def test_provider_error_propagates_without_storing_failed_turn(self): + session_id = "provider-error-test" + config = { + "provider": "openai_compatible", + "base_url": "http://127.0.0.1:3762/v1", + "model_name": "local-model", + "api_key": "", + } + + with patch.object( + chat_nodes, + "_call_openai_compatible_chat", + side_effect=RuntimeError("connection refused"), + ): + with self.assertRaisesRegex(RuntimeError, "connection refused"): + chat_nodes.ChatNode().chat( + model_name="unused", + base_url="http://127.0.0.1:11434", + user_message="hello", + session_id=session_id, + refresh_session=True, + llm_config=config, + ) + + self.assertEqual(chat_nodes._chat_sessions[session_id], []) + + def test_same_session_requests_are_serialized(self): + session_id = "concurrent-session-test" + config = { + "provider": "openai_compatible", + "base_url": "http://127.0.0.1:3762/v1", + "model_name": "local-model", + "api_key": "", + } + first_started = threading.Event() + second_started = threading.Event() + release_first = threading.Event() + call_count = 0 + call_count_lock = threading.Lock() + errors = [] + + def fake_chat(_base_url, _model_name, _api_key, messages, **_kwargs): + nonlocal call_count + with call_count_lock: + call_count += 1 + current_call = call_count + if current_call == 1: + first_started.set() + if not release_first.wait(timeout=2): + raise TimeoutError("first request was not released") + else: + second_started.set() + return f"{messages[-1]['content']}-reply" + + def run_chat(message): + try: + chat_nodes.ChatNode().chat( + model_name="unused", + base_url="http://127.0.0.1:11434", + user_message=message, + session_id=session_id, + llm_config=config, + ) + except Exception as exc: + errors.append(exc) + + with patch.object(chat_nodes, "_call_openai_compatible_chat", side_effect=fake_chat): + first_thread = threading.Thread(target=run_chat, args=("first",)) + second_thread = threading.Thread(target=run_chat, args=("second",)) + first_thread.start() + self.assertTrue(first_started.wait(timeout=1)) + second_thread.start() + second_started_before_release = second_started.wait(timeout=0.2) + release_first.set() + first_thread.join(timeout=2) + second_thread.join(timeout=2) + + self.assertFalse(first_thread.is_alive()) + self.assertFalse(second_thread.is_alive()) + self.assertEqual(errors, []) + self.assertFalse(second_started_before_release) + self.assertEqual( + chat_nodes._chat_sessions[session_id], + [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "first-reply"}, + {"role": "user", "content": "second"}, + {"role": "assistant", "content": "second-reply"}, + ], + ) + + def test_chat_node_has_provider_neutral_display_name(self): + self.assertEqual(chat_nodes.NODE_DISPLAY_NAME_MAPPINGS["ChatNode"], "Chat (LLM)") + + def test_readme_documents_supported_local_server_examples(self): + readme = (PROJECT_ROOT / "README.md").read_text(encoding="utf-8") + + expected_text = ( + "`provider=openai_compatible`", + "LM Studio", + "http://127.0.0.1:1234/v1", + "llama.cpp", + "http://127.0.0.1:8080/v1", + "LMDeploy", + "http://127.0.0.1:23333/v1", + ) + for text in expected_text: + with self.subTest(text=text): + self.assertIn(text, readme) + + +if __name__ == "__main__": + unittest.main()