From 4f7a0f56a59461a8a261ec42de21ca2e5be170dc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=A7=91=E6=9E=97=5FCyril=20Pilgrim?=
 <onepigge1999@gmail.com>
Date: Fri, 31 Jan 2025 11:04:03 +0800
Subject: [PATCH] push

---
 .../openai_compatible_loader.cpython-312.pyc  | Bin 4888 -> 6249 bytes
 openai_compatible_loader.py                   |  83 +++++++++++++++---
 2 files changed, 69 insertions(+), 14 deletions(-)

diff --git a/__pycache__/openai_compatible_loader.cpython-312.pyc b/__pycache__/openai_compatible_loader.cpython-312.pyc
index 337741d1eaea251decf84ef04a964a9b4d2748c3..ca2b49769be2508bfd734e6165bcde149e05d9e9 100644
GIT binary patch
delta 2118
zcmZuxZA?>F7(VBI_R`zZ7Di!wc*}=SlrJfuFvYftb55h#bP8FZ_TCl<P<L<9i8ze9
z0gK6GCsPvCxMWLM*cJ>XvMkGJ*$<QbsUsOpN3#8NaX*TfK&DxC&MnS`c#`v+=f3ZI
z-t+eKea`siKda3@n@ni{5BkH2;63Y@c?n(}d*nGP1~}PxnQ1l22E{0gidhyEQBG6D
zfI&84d%6JRSDxGCbnLU>X;mz;O|i-uicJx7aMplX&b-XFW}G!Uvl3aTipiR}3>zrJ
zNuQ>)_p3cK^YkI;P+!HNYAyypgW>t$F}Z~2$}QBy79OQ0fDCU^G9}Z4^aQ|+xj}u+
zPEo0N7q|pvaE=Dx8hxGqh-JZfW{~N|-y6r&x&eYg2Hj;V_OgR)H{2KnfXoHK^Xz-T
zC!C)O0zdUGF6%ru$jLnUfn(Gafx*`E9Lg7S9P1}T@D^wQS32-)^ixpxpCJ}A+RE*N
z1{CHR(u6|;d{u8qR(jBnT+<;ZOz@qGPgYbd!N|U_FF|$51KBOUt{z2}!XYW3g#De8
zzCml0oLqtqD&d4h^>yL9420BnE$j>T<IUY#s1J4Ug**SdR<#6|s#XyoMJJ%58SpBt
z4^<>aQ9qzPb-+Lo{<G|D|2L5w+V2=1g0%}ENr{3R#_a{y&yAgnXXIWjAKw?V6#r3N
zA1|$(JTuu9E7`ux(mA#y3$lw=0GDlE1%hcGOr}soO@@M@n!0=0cd}<U&lN7z?ETf=
zGRcm!<MM>@7kkTU28mAXd93eQOIkp7#oq~gOH%i(rJN~T)<d2?$pn_%#{lBmmTbAP
zg8IAKY|YH+>CRYQ^MY-6Ox*o6NfYlg?Ir+2;3D*z;AP2E0=|GAmiv;%Y_iy6p(7&n
z@La@{O?Zju5x7VVMR=`<Ln6)4W9FiO!92>498ne$&O;m%6<C61-D2T9g?uy%dUC|5
zohQ7291?ZlU1cy@sDG3Q-ZVN|Y4@`9Jj-G}pGM=P3I3GtY){B1&lIyaV6Iw@Z(rR)
zU@L~ND7BLCS`5g@@A9hKo`0%DN2-fJ9R{aCH4<J%U^{`AwMIR5pj-TY=$d)Re-6bs
z*2FLvhQrh_gOIQ*Z<HA{46?ZO*&EEq76FV<BXEQnv5bHKJM%~wfp8}Z8{J$X-uM}d
zzBg73aW50#AQ+`bsi}2oUIH6r3-LMdG*r{WEUri&GPYs)Q#a|Ub)jB_qg*$6AvO%0
zo7J=HF9|Y=J523ZcrAG5;Aun$#mXTkjkR}>NKYyqkV>V#P*@6t`uk+5t55Qs#**`$
z>`~UWR_SC^QmWd3A(3{>*VC`;Rn?HHHe-X6*Kak!rF5geVyW9DN=scfxqFS&DoGc}
zKzGDFL#@P!dl0PYdRodYn`mAIjI|Xjgy3W_UR1hR<a$u#nto@t^FF^&ba*lU@YR;M
zJ@0nQ1AF;WLD^)(blr4mY{!Y1vvr}MEtcIDca$tTTn`*%PjK%2*pcH4juVTyCsqZL
zGnM}_$v*_w;`xWyk{rlug=<OOjRo~0YFX#cv;Z#gQo1J{ehS^&I7DN1&F*fbzkzNS
z_iZI+kv!<G=OXzu;gX00B3q$*I~S?o2zNnZuh(%qB$h_bBkt#;0D5RK3VDynM-4o7
z3}jZFuBq~O`odNIQ17W8CEOM2(>8Ry+K3BMn+Uu{U?=*^bYK%{^@KrH-s@Miu%_-J
zo)>DUL&waOJN6Lc>jd5)pudTv=hZg}kk>K6X=nTV?TY@lX;y6M=;)9-&=2N@G_qTB
z;qH5e{xKi06$nAN%CIK#cM)7Mqoe7+?VKucTS3Gsc%5`4*PZUErlRE#e?+SVuY-;5
zDDvj9^LZ@h^A6fuOwSjy*sBUSxib^>TF%;Z!KQkRk-Pk1wHd8g8sJZ8i?wN!iFL!t
ziX<TGTy1KN9eI02HzW&Cz*@e08ykeTGm?O;lUiz>^`_SQbZU*~7ba;E5U;LXC%ana
vwC{SK(A!YVIt**kaht@E9#z}W2evxsM|W*y>AesJ;WE&dWCx8N*(&}8^@j;r

delta 948
zcmYL{O=uHA6vy{vKax$e*`!;Ww5g`4MpLz^)oNO+t=otn6cj<xRwP~BZLDZpXH&2k
zv1p-D6x2Zx^`L^F9#qW1da+G=DR@v(Xox*}P`tzsXh09nB%%xRd;fXw?SF=S`(@3|
z<<8fVWM{~;<lwFNu=A>O9+FqTc6La5nIZMrbjKLi;}jXaLbn^mobDu1Cd*~Wy4!H+
z<%Zjks)$3MQ&-0Lp7If?eo=Z9c3W)PP=D;OW*X@^b{8l)>yS2bd{EO-HR(n^g7g*f
z4&WccM)xC2&U)0)bOV_=?wC0^fVY>mdvY$M^_xblZxF@Igq9dKQ4@I$pSIT3)Kimx
zPAEvWIZoFSS`G|6TecH|>2k+CuzvuC8E3{}+=3^i_Ua^;v?ckZp!0XQi!PB#LCTVH
zrI<c0UFA#xG~@fSDohi%>m2?fHvo6dULhoz3+x5JmtCsZ6k#rT$P)2KS9mQFNoTzU
zvGg`=RP!iA_5Yg{FIq-&CQnd{$J{&4cG2x&g2HE}Fv`sPLbp2>7;yJMzQ~91Y#=xn
zXnq%Hezt!iK6UEto&)a#-E+R~1(tS|dcPNY3t3-xAzw@vO8FAM6|29^J>;14vg(Ix
zpmG9ciBfNPRXh6xR6%&k67}3clCMCu36la*839N#V*{$~f?8wExIk4L8Ml>cuT}Nq
z`Esv+6>&zZDH<pmDb`Y~qiDv8(&dj3nX@6|<gj5TOthX>Hc-$(q85Bc3Ab*cvXf#f
z1sx^YNU@Egl^`dWBLhcb2BII+bQARU_SV+oPfEm2EzK~Q@-G~$=yLePI4p1$sRw@8
z@VSbQn;$PzWhNtw(G_f_N|igqmgx|X&2(6d`uXX|(r7h5Q_T}S<7cCF>`Wa`s-?jy
z_|~EkiTqdnM-r%o9E93%urdM<@S(~!{GxK@mPS4fca=Ot<3qC)C!(b|%a&qxab2FJ
z4cRs8{)+Z(Q|8MPKiNhc@Qgzfe&neY_EP_S_>E^h9KyUeBp(Eb!#AcF@<|Kc<Zbu^
D#d`0<

diff --git a/openai_compatible_loader.py b/openai_compatible_loader.py
index 365ab27..7965e9c 100644
--- a/openai_compatible_loader.py
+++ b/openai_compatible_loader.py
@@ -17,8 +17,9 @@ class OpenAICompatibleLoader:
     def INPUT_TYPES(cls):
         return {
             "required": {
-                "model": ("STRING", {"default": "gpt-3.5-turbo"}),
-                "api_base": ("STRING", {"default": "https://api.openai.com"}),
+                "auto_detect": ("BOOLEAN", {"default": True, "label": "自动检测模型"}),
+                "model": ("STRING", {"default": "", "label": "手动指定模型(当自动检测失败时)"}),
+                "base_url": ("STRING", {"default": "https://api.openai.com"}),
                 "api_key": ("STRING", {"default": ""}),
                 "prompt": ("STRING", {"multiline": True}),
                 "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0}),
@@ -34,10 +35,25 @@ class OpenAICompatibleLoader:
     FUNCTION = "generate"
     CATEGORY = "DeepSeek_Toolkit"
 
-    async def async_generate(self, session: ClientSession, payload: dict, api_base: str, api_key: str):
+    async def fetch_models(self, session: ClientSession, base_url: str, api_key: str):
+        try:
+            async with session.get(
+                f"{base_url}/models",
+                headers={
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {api_key}"
+                }
+            ) as response:
+                response.raise_for_status()
+                data = await response.json()
+                return [model['id'] for model in data.get('data', [])]
+        except ClientError as e:
+            raise Exception(f"Failed to fetch models: {str(e)}")
+
+    async def async_generate(self, session: ClientSession, payload: dict, base_url: str, api_key: str):
         try:
             async with session.post(
-                f"{api_base}/chat/completions",
+                f"{base_url}/chat/completions",
                 headers={
                     "Content-Type": "application/json",
                     "Authorization": f"Bearer {api_key}"
@@ -50,21 +66,57 @@ class OpenAICompatibleLoader:
         except ClientError as e:
             raise Exception(f"API request failed: {str(e)}")
 
-    def generate(self, api_base: str, api_key: str, prompt: str, model: str,
+    def generate(self, base_url: str, api_key: str, prompt: str, model: str,
                 temperature: float, max_tokens: int, context: Optional[str] = None):
-        messages = [{
-            "role": "user",
-            "content": prompt
-        }]
+        # 模型检测逻辑
+        async def detect_model():
+            models = await self.fetch_models(session, base_url, api_key)
+            if not models:
+                raise ValueError("无法从API获取模型列表")
+            return models
         
+        # 构建消息体
+        messages = []
         if context:
-            messages.insert(0, {
+            messages.append({
                 "role": "system",
                 "content": context
             })
-
+        messages.append({
+            "role": "user",
+            "content": prompt
+        })
+        
+        # 确定最终使用的模型
+        available_models = []
+        selected_model = model
+        async def run_model_detection():
+            if auto_detect:
+                try:
+                    models = await self.fetch_models(session, base_url, api_key)
+                    if not models:
+                        raise ValueError("无法从API获取模型列表")
+                    return models
+                except Exception as e:
+                    if not model:
+                        raise ValueError(f"自动检测模型失败且未提供手动模型: {str(e)}")
+            return []
+        
+        # 模型检测逻辑
+        available_models = await run_model_detection()
+                if model and model not in available_models:
+                    print(f"警告: 手动指定的模型 '{model}' 不在可用模型列表中")
+                elif not model:
+                    selected_model = available_models[0] if available_models else ""
+            except Exception as e:
+                if not model:
+                    raise ValueError(f"自动检测模型失败且未提供手动模型: {str(e)}")
+        
+        if not selected_model:
+            raise ValueError("无法确定使用的模型，请检查base url或手动指定有效模型")
+        
         payload = {
-            "model": model,
+            "model": selected_model,
             "messages": messages,
             "temperature": temperature,
             "max_tokens": max_tokens
@@ -74,7 +126,10 @@ class OpenAICompatibleLoader:
             # 使用异步请求（需要ComfyUI运行在支持async的环境中）
             async def run_async():
                 async with ClientSession() as session:
-                    return await self.async_generate(session, payload, api_base, api_key)
+                    models = await self.fetch_models(session, base_url, api_key)
+                    if model not in models:
+                        raise ValueError(f"Model '{model}' not found in available models.")
+                    return await self.async_generate(session, payload, base_url, api_key)
 
             import asyncio
             response = asyncio.run(run_async())
@@ -86,7 +141,7 @@ class OpenAICompatibleLoader:
             try:
                 import requests
                 response = requests.post(
-                    f"{api_base}/v1/chat/completions",
+                    f"{base_url}/v1/chat/completions",
                     headers={
                         "Content-Type": "application/json",
                         "Authorization": f"Bearer {api_key}"