Update zhipuai_LLM.py

zhipuai_LLM.py

@@ -1,161 +1,164 @@
-from typing import Any, Dict, Iterator, List, Optional, Union
-import os
-import time
-from zhipuai import ZhipuAI
-from langchain_core.callbacks import CallbackManagerForLLMRun
-from langchain_core.language_models import BaseChatModel
-from langchain_core.messages import (
-    AIMessage,
-    AIMessageChunk,
-    BaseMessage,
-    SystemMessage,
-    ChatMessage,
-    HumanMessage
-)
-from langchain_core.messages.ai import UsageMetadata
-from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
-
-def _convert_message_to_dict(message: Union[BaseMessage, dict, tuple]) -> dict:
-    role = "user"
-    content = ""
-    
-    if isinstance(message, tuple) and len(message) == 2:
-        msg_type, content = message
-        if msg_type == "system":
-            role = "system"
-        elif msg_type in ["ai", "assistant"]:
-            role = "assistant"
-        else:
-            role = "user"
-            
-    elif isinstance(message, dict):
-        msg_type = message.get("role", "user")
-        content = message.get("content", "")
-        if msg_type == "system":
-            role = "system"
-        elif msg_type in ["ai", "assistant"]:
-            role = "assistant"
-        else:
-            role = "user"
-            
-    elif isinstance(message, BaseMessage):
-        content = message.content
-        if isinstance(message, ChatMessage):
-            role = message.role
-        elif isinstance(message, HumanMessage):
-            role = "user"
-        elif isinstance(message, AIMessage):
-            role = "assistant"
-        elif isinstance(message, SystemMessage):
-            role = "system"
-        else:
-            role = "user"
-    else:
-        content = str(message)
-        
-    return {"role": role, "content": content}
-
-class ZhipuaiLLM(BaseChatModel):
-    model_name: str = "glm-4-flash" 
-    temperature: Optional[float] = 0.1
-    max_tokens: Optional[int] = None
-    timeout: Optional[int] = None
-    stop: Optional[List[str]] = None
-    max_retries: int = 3
-    api_key: str | None = None
-
-    def _get_client(self) -> ZhipuAI:
-        current_api_key = self.api_key or os.environ.get("ZHIPUAI_API_KEY")
-        return ZhipuAI(api_key=current_api_key)
-
-    def _generate(
-        self,
-        messages: List[Any],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> ChatResult:
-        zhipu_messages = [_convert_message_to_dict(message) for message in messages]
-        start_time = time.time()
-        
-        client = self._get_client()
-        response = client.chat.completions.create(
-            model=self.model_name,
-            temperature=self.temperature,
-            max_tokens=self.max_tokens,
-            timeout=self.timeout,
-            stop=stop,
-            messages=zhipu_messages,
-            **kwargs
-        )
-        
-        time_in_seconds = time.time() - start_time
-        message = AIMessage(
-            content=response.choices[0].message.content,
-            additional_kwargs={},
-            response_metadata={"time_in_seconds": round(time_in_seconds, 3)},
-            usage_metadata={
-                "input_tokens": response.usage.prompt_tokens,
-                "output_tokens": response.usage.completion_tokens,
-                "total_tokens": response.usage.total_tokens,
-            },
-        )
-        return ChatResult(generations=[ChatGeneration(message=message)])
-
-    def _stream(
-        self,
-        messages: List[Any],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> Iterator[ChatGenerationChunk]:
-        zhipu_messages = [_convert_message_to_dict(message) for message in messages]
-        start_time = time.time()
-        
-        client = self._get_client()
-        response = client.chat.completions.create(
-            model=self.model_name,
-            stream=True,
-            temperature=self.temperature,
-            max_tokens=self.max_tokens,
-            timeout=self.timeout,
-            stop=stop,
-            messages=zhipu_messages,
-            **kwargs
-        )
-        
-        usage_metadata = None
-        for res in response:
-            if hasattr(res, 'usage') and res.usage:
-                usage_metadata = UsageMetadata({
-                    "input_tokens": getattr(res.usage, 'prompt_tokens', 0),
-                    "output_tokens": getattr(res.usage, 'completion_tokens', 0),
-                    "total_tokens": getattr(res.usage, 'total_tokens', 0),
-                })
-                
-            chunk_content = res.choices[0].delta.content if res.choices and res.choices[0].delta.content else ""
-            chunk = ChatGenerationChunk(message=AIMessageChunk(content=chunk_content))
-
-            if run_manager and chunk_content:
-                run_manager.on_llm_new_token(chunk_content, chunk=chunk)
-                
-            yield chunk
-            
-        time_in_sec = time.time() - start_time
-        final_chunk = ChatGenerationChunk(
-            message=AIMessageChunk(
-                content="", 
-                response_metadata={"time_in_sec": round(time_in_sec, 3)}, 
-                usage_metadata=usage_metadata
-            )
-        )
-        if run_manager:
-            run_manager.on_llm_new_token("", chunk=final_chunk)
-        yield final_chunk
-
-    @property
-    def _llm_type(self) -> str:
-        return self.model_name
-
-    @property
-    def _identifying_params(self) -> Dict[str,
+from typing import Any, Dict, Iterator, List, Optional, Union
+import os
+import time
+from zhipuai import ZhipuAI
+from langchain_core.callbacks import CallbackManagerForLLMRun
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import (
+    AIMessage,
+    AIMessageChunk,
+    BaseMessage,
+    SystemMessage,
+    ChatMessage,
+    HumanMessage
+)
+from langchain_core.messages.ai import UsageMetadata
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
+
+def _convert_message_to_dict(message: Union[BaseMessage, dict, tuple]) -> dict:
+    role = "user"
+    content = ""
+    
+    if isinstance(message, tuple) and len(message) == 2:
+        msg_type, content = message
+        if msg_type == "system":
+            role = "system"
+        elif msg_type in ["ai", "assistant"]:
+            role = "assistant"
+        else:
+            role = "user"
+            
+    elif isinstance(message, dict):
+        msg_type = message.get("role", "user")
+        content = message.get("content", "")
+        if msg_type == "system":
+            role = "system"
+        elif msg_type in ["ai", "assistant"]:
+            role = "assistant"
+        else:
+            role = "user"
+            
+    elif isinstance(message, BaseMessage):
+        content = message.content
+        if isinstance(message, ChatMessage):
+            role = message.role
+        elif isinstance(message, HumanMessage):
+            role = "user"
+        elif isinstance(message, AIMessage):
+            role = "assistant"
+        elif isinstance(message, SystemMessage):
+            role = "system"
+        else:
+            role = "user"
+    else:
+        content = str(message)
+        
+    return {"role": role, "content": content}
+
+class ZhipuaiLLM(BaseChatModel):
+    model_name: str = "glm-4-flash" 
+    temperature: Optional[float] = 0.1
+    max_tokens: Optional[int] = None
+    timeout: Optional[int] = None
+    stop: Optional[List[str]] = None
+    max_retries: int = 3
+    api_key: str | None = None
+
+    def _get_client(self) -> ZhipuAI:
+        current_api_key = self.api_key or os.environ.get("ZHIPUAI_API_KEY")
+        return ZhipuAI(api_key=current_api_key)
+
+    def _generate(
+        self,
+        messages: List[Any],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        zhipu_messages = [_convert_message_to_dict(message) for message in messages]
+        start_time = time.time()
+        
+        client = self._get_client()
+        response = client.chat.completions.create(
+            model=self.model_name,
+            temperature=self.temperature,
+            max_tokens=self.max_tokens,
+            timeout=self.timeout,
+            stop=stop,
+            messages=zhipu_messages,
+            **kwargs
+        )
+        
+        time_in_seconds = time.time() - start_time
+        message = AIMessage(
+            content=response.choices[0].message.content,
+            additional_kwargs={},
+            response_metadata={"time_in_seconds": round(time_in_seconds, 3)},
+            usage_metadata={
+                "input_tokens": response.usage.prompt_tokens,
+                "output_tokens": response.usage.completion_tokens,
+                "total_tokens": response.usage.total_tokens,
+            },
+        )
+        return ChatResult(generations=[ChatGeneration(message=message)])
+
+    def _stream(
+        self,
+        messages: List[Any],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        zhipu_messages = [_convert_message_to_dict(message) for message in messages]
+        start_time = time.time()
+        
+        client = self._get_client()
+        response = client.chat.completions.create(
+            model=self.model_name,
+            stream=True,
+            temperature=self.temperature,
+            max_tokens=self.max_tokens,
+            timeout=self.timeout,
+            stop=stop,
+            messages=zhipu_messages,
+            **kwargs
+        )
+        
+        usage_metadata = None
+        for res in response:
+            if hasattr(res, 'usage') and res.usage:
+                usage_metadata = UsageMetadata({
+                    "input_tokens": getattr(res.usage, 'prompt_tokens', 0),
+                    "output_tokens": getattr(res.usage, 'completion_tokens', 0),
+                    "total_tokens": getattr(res.usage, 'total_tokens', 0),
+                })
+                
+            chunk_content = res.choices[0].delta.content if res.choices and res.choices[0].delta.content else ""
+            chunk = ChatGenerationChunk(message=AIMessageChunk(content=chunk_content))
+
+            if run_manager and chunk_content:
+                run_manager.on_llm_new_token(chunk_content, chunk=chunk)
+                
+            yield chunk
+            
+        time_in_sec = time.time() - start_time
+        final_chunk = ChatGenerationChunk(
+            message=AIMessageChunk(
+                content="", 
+                response_metadata={"time_in_sec": round(time_in_sec, 3)}, 
+                usage_metadata=usage_metadata
+            )
+        )
+        if run_manager:
+            run_manager.on_llm_new_token("", chunk=final_chunk)
+        yield final_chunk
+
+    @property
+    def _llm_type(self) -> str:
+        return self.model_name
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        return {
+            "model_name": self.model_name,
+        }