更新

qwen_agent/llm/qwen_oai.py

@@ -124,7 +124,7 @@ _TEXT_COMPLETION_CMD = object()
 #
 def parse_messages(messages, functions):
     if all(m.role != "user" for m in messages):
-        raise Exception( f"Invalid request: Expecting at least one user message.",)
+        raise Exception(f"Invalid request: Expecting at least one user message.", )
     messages = copy.deepcopy(messages)
     default_system = "You are a helpful assistant."
     system = ""
@@ -381,7 +381,7 @@ def predict(
     stop_words_ids = [tokenizer.encode(s) for s in stop_words] if stop_words else None
     if stop_words:
         # TODO: It's a little bit tricky to trim stop words in the stream mode.
-        raise Exception("Invalid request: custom stop words are not yet supported for stream mode.",)
+        raise Exception("Invalid request: custom stop words are not yet supported for stream mode.", )
     response_generator = qmodel.chat_stream(
         tokenizer, query, history=history, stop_words_ids=stop_words_ids, **gen_kwargs
     )
@@ -420,35 +420,34 @@ class QwenChatAsOAI(BaseChatModel):
         self.model = model
         super().__init__()
         tokenizer = AutoTokenizer.from_pretrained(
-            self.checkpoint_path,
+            self.model,
             trust_remote_code=True,
             resume_download=True,
         )
         device_map = "cpu"
         # device_map = "auto"
         qmodel = AutoModelForCausalLM.from_pretrained(
-            self.checkpoint_path,
+            self.model,
             device_map=device_map,
             trust_remote_code=True,
             resume_download=True,
         ).eval()
 
         qmodel.generation_config = GenerationConfig.from_pretrained(
-            self.checkpoint_path,
+            self.model,
             trust_remote_code=True,
             resume_download=True,
         )
 
-
     def _chat_stream(
-        self,
-        messages: List[Dict],
-        stop: Optional[List[str]] = None,
+            self,
+            messages: List[Dict],
+            stop: Optional[List[str]] = None,
     ) -> Iterator[str]:
         _request = ChatCompletionRequest(model=self.model,
-                                                messages=messages,
-                                                stop=stop,
-                                                stream=True)
+                                         messages=messages,
+                                         stop=stop,
+                                         stream=True)
         response = create_chat_completion(_request)
         # TODO: error handling
         for chunk in response:
@@ -456,14 +455,11 @@ class QwenChatAsOAI(BaseChatModel):
                 yield chunk.choices[0].delta.content
 
     def _chat_no_stream(
-        self,
-        messages: List[Dict],
-        stop: Optional[List[str]] = None,
+            self,
+            messages: List[Dict],
+            stop: Optional[List[str]] = None,
     ) -> str:
-        _request = ChatCompletionRequest(model=self.model,
-                                                messages=messages,
-                                                stop=stop,
-                                                stream=False)
+        _request = ChatCompletionRequest(model=self.model, messages=messages, stop=stop, stream=False)
         response = create_chat_completion(_request)
         # TODO: error handling
         return response.choices[0].message.content
@@ -472,16 +468,13 @@ class QwenChatAsOAI(BaseChatModel):
                             messages: List[Dict],
                             functions: Optional[List[Dict]] = None) -> Dict:
         if functions:
-            _request = ChatCompletionRequest(model=self.model,
-                                                    messages=messages,
-                                                    functions=functions)
+            _request = ChatCompletionRequest(model=self.model, messages=messages, functions=functions)
             response = create_chat_completion(_request)
         else:
-            _request = ChatCompletionRequest(model=self.model,
-                                             messages=messages)
+            _request = ChatCompletionRequest(model=self.model, messages=messages)
             response = create_chat_completion(_request)
         # TODO: error handling
-        return response.choices[0].message.dict()
+        return response.choices[0].message.model_dump()
 
 
 class QwenChatAsOAI1(BaseChatModel):
@@ -495,9 +488,9 @@ class QwenChatAsOAI1(BaseChatModel):
         self.model = model
 
     def _chat_stream(
-        self,
-        messages: List[Dict],
-        stop: Optional[List[str]] = None,
+            self,
+            messages: List[Dict],
+            stop: Optional[List[str]] = None,
     ) -> Iterator[str]:
         response = openai.ChatCompletion.create(model=self.model,
                                                 messages=messages,
@@ -509,9 +502,9 @@ class QwenChatAsOAI1(BaseChatModel):
                 yield chunk.choices[0].delta.content
 
     def _chat_no_stream(
-        self,
-        messages: List[Dict],
-        stop: Optional[List[str]] = None,
+            self,
+            messages: List[Dict],
+            stop: Optional[List[str]] = None,
     ) -> str:
         response = openai.ChatCompletion.create(model=self.model,
                                                 messages=messages,
@@ -531,4 +524,4 @@ class QwenChatAsOAI1(BaseChatModel):
             response = openai.ChatCompletion.create(model=self.model,
                                                     messages=messages)
         # TODO: error handling
-        return response.choices[0].message
+        return response.choices[0].message