feat: add zhipuai (#1188)

api/controllers/console/workspace/model_providers.py

@@ -246,7 +246,8 @@ class ModelProviderModelParameterRuleApi(Resource):
                 'enabled': v.enabled,
                 'min': v.min,
                 'max': v.max,
-                'default': v.default
+                'default': v.default,
+                'precision': v.precision
             }
             for k, v in vars(parameter_rules).items()
         }
@@ -290,10 +291,15 @@ class ModelProviderFreeQuotaQualificationVerifyApi(Resource):
     @login_required
     @account_initialization_required
     def get(self, provider_name: str):
+        parser = reqparse.RequestParser()
+        parser.add_argument('token', type=str, required=False, nullable=True, location='args')
+        args = parser.parse_args()
+
         provider_service = ProviderService()
         result = provider_service.free_quota_qualification_verify(
             tenant_id=current_user.current_tenant_id,
-            provider_name=provider_name
+            provider_name=provider_name,
+            token=args['token']
         )
 
         return result

api/core/callback_handler/llm_callback_handler.py

@@ -63,7 +63,18 @@ class LLMCallbackHandler(BaseCallbackHandler):
             self.conversation_message_task.append_message_text(response.generations[0][0].text)
             self.llm_message.completion = response.generations[0][0].text
 
-        self.llm_message.completion_tokens = self.model_instance.get_num_tokens([PromptMessage(content=self.llm_message.completion)])
+        if response.llm_output and 'token_usage' in response.llm_output:
+            if 'prompt_tokens' in response.llm_output['token_usage']:
+                self.llm_message.prompt_tokens = response.llm_output['token_usage']['prompt_tokens']
+
+            if 'completion_tokens' in response.llm_output['token_usage']:
+                self.llm_message.completion_tokens = response.llm_output['token_usage']['completion_tokens']
+            else:
+                self.llm_message.completion_tokens = self.model_instance.get_num_tokens(
+                    [PromptMessage(content=self.llm_message.completion)])
+        else:
+            self.llm_message.completion_tokens = self.model_instance.get_num_tokens(
+                [PromptMessage(content=self.llm_message.completion)])
 
         self.conversation_message_task.save_message(self.llm_message)
 

api/core/chain/sensitive_word_avoidance_chain.py

@@ -2,13 +2,8 @@ import enum
 import logging
 from typing import List, Dict, Optional, Any
 
-import openai
-from flask import current_app
 from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains.base import Chain
-from openai import InvalidRequestError
-from openai.error import APIConnectionError, APIError, ServiceUnavailableError, Timeout, RateLimitError, \
-    AuthenticationError, OpenAIError
 from pydantic import BaseModel
 
 from core.model_providers.error import LLMBadRequestError
@@ -86,6 +81,12 @@ class SensitiveWordAvoidanceChain(Chain):
             result = self._check_moderation(text)
 
         if not result:
-            raise LLMBadRequestError(self.sensitive_word_avoidance_rule.canned_response)
+            raise SensitiveWordAvoidanceError(self.sensitive_word_avoidance_rule.canned_response)
 
         return {self.output_key: text}
+
+
+class SensitiveWordAvoidanceError(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+        self.message = message

api/core/completion.py

@@ -7,6 +7,7 @@ from requests.exceptions import ChunkedEncodingError
 from core.agent.agent_executor import AgentExecuteResult, PlanningStrategy
 from core.callback_handler.main_chain_gather_callback_handler import MainChainGatherCallbackHandler
 from core.callback_handler.llm_callback_handler import LLMCallbackHandler
+from core.chain.sensitive_word_avoidance_chain import SensitiveWordAvoidanceError
 from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException
 from core.model_providers.error import LLMBadRequestError
 from core.memory.read_only_conversation_token_db_buffer_shared_memory import \
@@ -76,28 +77,53 @@ class Completion:
             app_model_config=app_model_config
         )
 
-        # parse sensitive_word_avoidance_chain
-        chain_callback = MainChainGatherCallbackHandler(conversation_message_task)
-        sensitive_word_avoidance_chain = orchestrator_rule_parser.to_sensitive_word_avoidance_chain(final_model_instance, [chain_callback])
-        if sensitive_word_avoidance_chain:
-            query = sensitive_word_avoidance_chain.run(query)
-
-        # get agent executor
-        agent_executor = orchestrator_rule_parser.to_agent_executor(
-            conversation_message_task=conversation_message_task,
-            memory=memory,
-            rest_tokens=rest_tokens_for_context_and_memory,
-            chain_callback=chain_callback
-        )
-
-        # run agent executor
-        agent_execute_result = None
-        if agent_executor:
-            should_use_agent = agent_executor.should_use_agent(query)
-            if should_use_agent:
-                agent_execute_result = agent_executor.run(query)
-        # run the final llm
         try:
+            # parse sensitive_word_avoidance_chain
+            chain_callback = MainChainGatherCallbackHandler(conversation_message_task)
+            sensitive_word_avoidance_chain = orchestrator_rule_parser.to_sensitive_word_avoidance_chain(
+                final_model_instance, [chain_callback])
+            if sensitive_word_avoidance_chain:
+                try:
+                    query = sensitive_word_avoidance_chain.run(query)
+                except SensitiveWordAvoidanceError as ex:
+                    cls.run_final_llm(
+                        model_instance=final_model_instance,
+                        mode=app.mode,
+                        app_model_config=app_model_config,
+                        query=query,
+                        inputs=inputs,
+                        agent_execute_result=None,
+                        conversation_message_task=conversation_message_task,
+                        memory=memory,
+                        fake_response=ex.message
+                    )
+                    return
+
+            # get agent executor
+            agent_executor = orchestrator_rule_parser.to_agent_executor(
+                conversation_message_task=conversation_message_task,
+                memory=memory,
+                rest_tokens=rest_tokens_for_context_and_memory,
+                chain_callback=chain_callback,
+                retriever_from=retriever_from
+            )
+
+            # run agent executor
+            agent_execute_result = None
+            if agent_executor:
+                should_use_agent = agent_executor.should_use_agent(query)
+                if should_use_agent:
+                    agent_execute_result = agent_executor.run(query)
+
+            # When no extra pre prompt is specified,
+            # the output of the agent can be used directly as the main output content without calling LLM again
+            fake_response = None
+            if not app_model_config.pre_prompt and agent_execute_result and agent_execute_result.output \
+                    and agent_execute_result.strategy not in [PlanningStrategy.ROUTER,
+                                                              PlanningStrategy.REACT_ROUTER]:
+                fake_response = agent_execute_result.output
+
+            # run the final llm
             cls.run_final_llm(
                 model_instance=final_model_instance,
                 mode=app.mode,
@@ -106,7 +132,8 @@ class Completion:
                 inputs=inputs,
                 agent_execute_result=agent_execute_result,
                 conversation_message_task=conversation_message_task,
-                memory=memory
+                memory=memory,
+                fake_response=fake_response
             )
         except ConversationTaskStoppedException:
             return
@@ -121,14 +148,8 @@ class Completion:
                       inputs: dict,
                       agent_execute_result: Optional[AgentExecuteResult],
                       conversation_message_task: ConversationMessageTask,
-                      memory: Optional[ReadOnlyConversationTokenDBBufferSharedMemory]):
-        # When no extra pre prompt is specified,
-        # the output of the agent can be used directly as the main output content without calling LLM again
-        fake_response = None
-        if not app_model_config.pre_prompt and agent_execute_result and agent_execute_result.output \
-                and agent_execute_result.strategy not in [PlanningStrategy.ROUTER, PlanningStrategy.REACT_ROUTER]:
-            fake_response = agent_execute_result.output
-
+                      memory: Optional[ReadOnlyConversationTokenDBBufferSharedMemory],
+                      fake_response: Optional[str]):
         # get llm prompt
         prompt_messages, stop_words = model_instance.get_prompt(
             mode=mode,

api/core/helper/moderation.py

@@ -1,32 +1,34 @@
 import logging
 
 import openai
-from flask import current_app
 
 from core.model_providers.error import LLMBadRequestError
 from core.model_providers.providers.base import BaseModelProvider
+from core.model_providers.providers.hosted import hosted_config, hosted_model_providers
 from models.provider import ProviderType
 
 
 def check_moderation(model_provider: BaseModelProvider, text: str) -> bool:
-    if current_app.config['HOSTED_MODERATION_ENABLED'] and current_app.config['HOSTED_MODERATION_PROVIDERS']:
-        moderation_providers = current_app.config['HOSTED_MODERATION_PROVIDERS'].split(',')
-
+    if hosted_config.moderation.enabled is True and hosted_model_providers.openai:
         if model_provider.provider.provider_type == ProviderType.SYSTEM.value \
-                and model_provider.provider_name in moderation_providers:
+                and model_provider.provider_name in hosted_config.moderation.providers:
             # 2000 text per chunk
             length = 2000
-            chunks = [text[i:i + length] for i in range(0, len(text), length)]
+            text_chunks = [text[i:i + length] for i in range(0, len(text), length)]
 
-            try:
-                moderation_result = openai.Moderation.create(input=chunks,
-                                                             api_key=current_app.config['HOSTED_OPENAI_API_KEY'])
-            except Exception as ex:
-                logging.exception(ex)
-                raise LLMBadRequestError('Rate limit exceeded, please try again later.')
+            max_text_chunks = 32
+            chunks = [text_chunks[i:i + max_text_chunks] for i in range(0, len(text_chunks), max_text_chunks)]
 
-            for result in moderation_result.results:
-                if result['flagged'] is True:
-                    return False
+            for text_chunk in chunks:
+                try:
+                    moderation_result = openai.Moderation.create(input=text_chunk,
+                                                                 api_key=hosted_model_providers.openai.api_key)
+                except Exception as ex:
+                    logging.exception(ex)
+                    raise LLMBadRequestError('Rate limit exceeded, please try again later.')
+
+                for result in moderation_result.results:
+                    if result['flagged'] is True:
+                        return False
 
     return True

api/core/model_providers/model_provider_factory.py

@@ -45,6 +45,9 @@ class ModelProviderFactory:
         elif provider_name == 'wenxin':
             from core.model_providers.providers.wenxin_provider import WenxinProvider
             return WenxinProvider
+        elif provider_name == 'zhipuai':
+            from core.model_providers.providers.zhipuai_provider import ZhipuAIProvider
+            return ZhipuAIProvider
         elif provider_name == 'chatglm':
             from core.model_providers.providers.chatglm_provider import ChatGLMProvider
             return ChatGLMProvider

api/core/model_providers/models/embedding/zhipuai_embedding.py

@@ -0,0 +1,22 @@
+from core.model_providers.error import LLMBadRequestError
+from core.model_providers.providers.base import BaseModelProvider
+from core.model_providers.models.embedding.base import BaseEmbedding
+from core.third_party.langchain.embeddings.zhipuai_embedding import ZhipuAIEmbeddings
+
+
+class ZhipuAIEmbedding(BaseEmbedding):
+    def __init__(self, model_provider: BaseModelProvider, name: str):
+        credentials = model_provider.get_model_credentials(
+            model_name=name,
+            model_type=self.type
+        )
+
+        client = ZhipuAIEmbeddings(
+            model=name,
+            **credentials,
+        )
+
+        super().__init__(model_provider, client, name)
+
+    def handle_exceptions(self, ex: Exception) -> Exception:
+        return LLMBadRequestError(f"ZhipuAI embedding: {str(ex)}")

api/core/model_providers/models/entity/model_params.py

@@ -49,6 +49,7 @@ class KwargRule(Generic[T], BaseModel):
     max: Optional[T] = None
     default: Optional[T] = None
     alias: Optional[str] = None
+    precision: Optional[int] = None
 
 
 class ModelKwargsRules(BaseModel):

api/core/model_providers/models/llm/zhipuai_model.py

@@ -0,0 +1,61 @@
+from typing import List, Optional, Any
+
+from langchain.callbacks.manager import Callbacks
+from langchain.schema import LLMResult
+
+from core.model_providers.error import LLMBadRequestError
+from core.model_providers.models.llm.base import BaseLLM
+from core.model_providers.models.entity.message import PromptMessage
+from core.model_providers.models.entity.model_params import ModelMode, ModelKwargs
+from core.third_party.langchain.llms.zhipuai_llm import ZhipuAIChatLLM
+
+
+class ZhipuAIModel(BaseLLM):
+    model_mode: ModelMode = ModelMode.CHAT
+
+    def _init_client(self) -> Any:
+        provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, self.model_kwargs)
+        return ZhipuAIChatLLM(
+            streaming=self.streaming,
+            callbacks=self.callbacks,
+            **self.credentials,
+            **provider_model_kwargs
+        )
+
+    def _run(self, messages: List[PromptMessage],
+             stop: Optional[List[str]] = None,
+             callbacks: Callbacks = None,
+             **kwargs) -> LLMResult:
+        """
+        run predict by prompt messages and stop words.
+
+        :param messages:
+        :param stop:
+        :param callbacks:
+        :return:
+        """
+        prompts = self._get_prompt_from_messages(messages)
+        return self._client.generate([prompts], stop, callbacks)
+
+    def get_num_tokens(self, messages: List[PromptMessage]) -> int:
+        """
+        get num tokens of prompt messages.
+
+        :param messages:
+        :return:
+        """
+        prompts = self._get_prompt_from_messages(messages)
+        return max(self._client.get_num_tokens_from_messages(prompts), 0)
+
+    def _set_model_kwargs(self, model_kwargs: ModelKwargs):
+        provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, model_kwargs)
+        for k, v in provider_model_kwargs.items():
+            if hasattr(self.client, k):
+                setattr(self.client, k, v)
+
+    def handle_exceptions(self, ex: Exception) -> Exception:
+        return LLMBadRequestError(f"ZhipuAI: {str(ex)}")
+
+    @property
+    def support_streaming(self):
+        return True

api/core/model_providers/models/moderation/openai_moderation.py

@@ -23,14 +23,18 @@ class OpenAIModeration(BaseModeration):
 
         # 2000 text per chunk
         length = 2000
-        chunks = [text[i:i + length] for i in range(0, len(text), length)]
+        text_chunks = [text[i:i + length] for i in range(0, len(text), length)]
 
-        moderation_result = self._client.create(input=chunks,
-                                                api_key=credentials['openai_api_key'])
+        max_text_chunks = 32
+        chunks = [text_chunks[i:i + max_text_chunks] for i in range(0, len(text_chunks), max_text_chunks)]
 
-        for result in moderation_result.results:
-            if result['flagged'] is True:
-                return False
+        for text_chunk in chunks:
+            moderation_result = self._client.create(input=text_chunk,
+                                                    api_key=credentials['openai_api_key'])
+
+            for result in moderation_result.results:
+                if result['flagged'] is True:
+                    return False
 
         return True
 

api/core/model_providers/providers/anthropic_provider.py

@@ -69,11 +69,11 @@ class AnthropicProvider(BaseModelProvider):
         :return:
         """
         return ModelKwargsRules(
-            temperature=KwargRule[float](min=0, max=1, default=1),
-            top_p=KwargRule[float](min=0, max=1, default=0.7),
+            temperature=KwargRule[float](min=0, max=1, default=1, precision=2),
+            top_p=KwargRule[float](min=0, max=1, default=0.7, precision=2),
             presence_penalty=KwargRule[float](enabled=False),
             frequency_penalty=KwargRule[float](enabled=False),
-            max_tokens=KwargRule[int](alias="max_tokens_to_sample", min=10, max=100000, default=256),
+            max_tokens=KwargRule[int](alias="max_tokens_to_sample", min=10, max=100000, default=256, precision=0),
         )
 
     @classmethod

api/core/model_providers/providers/azure_openai_provider.py

@@ -164,14 +164,14 @@ class AzureOpenAIProvider(BaseModelProvider):
         model_credentials = self.get_model_credentials(model_name, model_type)
 
         return ModelKwargsRules(
-            temperature=KwargRule[float](min=0, max=2, default=1),
-            top_p=KwargRule[float](min=0, max=1, default=1),
-            presence_penalty=KwargRule[float](min=-2, max=2, default=0),
-            frequency_penalty=KwargRule[float](min=-2, max=2, default=0),
+            temperature=KwargRule[float](min=0, max=2, default=1, precision=2),
+            top_p=KwargRule[float](min=0, max=1, default=1, precision=2),
+            presence_penalty=KwargRule[float](min=-2, max=2, default=0, precision=2),
+            frequency_penalty=KwargRule[float](min=-2, max=2, default=0, precision=2),
             max_tokens=KwargRule[int](min=10, max=base_model_max_tokens.get(
                 model_credentials['base_model_name'],
                 4097
-            ), default=16),
+            ), default=16, precision=0),
         )
 
     @classmethod

api/core/model_providers/providers/chatglm_provider.py

@@ -64,11 +64,11 @@ class ChatGLMProvider(BaseModelProvider):
         }
 
         return ModelKwargsRules(
-            temperature=KwargRule[float](min=0, max=2, default=1),
-            top_p=KwargRule[float](min=0, max=1, default=0.7),
+            temperature=KwargRule[float](min=0, max=2, default=1, precision=2),
+            top_p=KwargRule[float](min=0, max=1, default=0.7, precision=2),
             presence_penalty=KwargRule[float](enabled=False),
             frequency_penalty=KwargRule[float](enabled=False),
-            max_tokens=KwargRule[int](alias='max_token', min=10, max=model_max_tokens.get(model_name), default=2048),
+            max_tokens=KwargRule[int](alias='max_token', min=10, max=model_max_tokens.get(model_name), default=2048, precision=0),
         )
 
     @classmethod

api/core/model_providers/providers/hosted.py

@@ -45,6 +45,18 @@ class HostedModelProviders(BaseModel):
 hosted_model_providers = HostedModelProviders()
 
 
+class HostedModerationConfig(BaseModel):
+    enabled: bool = False
+    providers: list[str] = []
+
+
+class HostedConfig(BaseModel):
+    moderation = HostedModerationConfig()
+
+
+hosted_config = HostedConfig()
+
+
 def init_app(app: Flask):
     if os.environ.get("DEBUG") and os.environ.get("DEBUG").lower() == 'true':
         langchain.verbose = True
@@ -78,3 +90,9 @@ def init_app(app: Flask):
             paid_min_quantity=app.config.get("HOSTED_ANTHROPIC_PAID_MIN_QUANTITY"),
             paid_max_quantity=app.config.get("HOSTED_ANTHROPIC_PAID_MAX_QUANTITY"),
         )
+
+    if app.config.get("HOSTED_MODERATION_ENABLED") and app.config.get("HOSTED_MODERATION_PROVIDERS"):
+        hosted_config.moderation = HostedModerationConfig(
+            enabled=app.config.get("HOSTED_MODERATION_ENABLED"),
+            providers=app.config.get("HOSTED_MODERATION_PROVIDERS").split(',')
+        )

api/core/model_providers/providers/huggingface_hub_provider.py

@@ -47,11 +47,11 @@ class HuggingfaceHubProvider(BaseModelProvider):
         :return:
         """
         return ModelKwargsRules(
-            temperature=KwargRule[float](min=0, max=2, default=1),
-            top_p=KwargRule[float](min=0.01, max=0.99, default=0.7),
+            temperature=KwargRule[float](min=0, max=2, default=1, precision=2),
+            top_p=KwargRule[float](min=0.01, max=0.99, default=0.7, precision=2),
             presence_penalty=KwargRule[float](enabled=False),
             frequency_penalty=KwargRule[float](enabled=False),
-            max_tokens=KwargRule[int](alias='max_new_tokens', min=10, max=4000, default=200),
+            max_tokens=KwargRule[int](alias='max_new_tokens', min=10, max=4000, default=200, precision=0),
         )
 
     @classmethod

api/core/model_providers/providers/localai_provider.py

@@ -52,9 +52,9 @@ class LocalAIProvider(BaseModelProvider):
         :return:
         """
         return ModelKwargsRules(
-            temperature=KwargRule[float](min=0, max=2, default=0.7),
-            top_p=KwargRule[float](min=0, max=1, default=1),
-            max_tokens=KwargRule[int](min=10, max=4097, default=16),
+            temperature=KwargRule[float](min=0, max=2, default=0.7, precision=2),
+            top_p=KwargRule[float](min=0, max=1, default=1, precision=2),
+            max_tokens=KwargRule[int](min=10, max=4097, default=16, precision=0),
         )
 
     @classmethod

api/core/model_providers/providers/minimax_provider.py

@@ -74,11 +74,11 @@ class MinimaxProvider(BaseModelProvider):
         }
 
         return ModelKwargsRules(
-            temperature=KwargRule[float](min=0.01, max=1, default=0.9),
-            top_p=KwargRule[float](min=0, max=1, default=0.95),
+            temperature=KwargRule[float](min=0.01, max=1, default=0.9, precision=2),
+            top_p=KwargRule[float](min=0, max=1, default=0.95, precision=2),
             presence_penalty=KwargRule[float](enabled=False),
             frequency_penalty=KwargRule[float](enabled=False),
-            max_tokens=KwargRule[int](min=10, max=model_max_tokens.get(model_name, 6144), default=1024),
+            max_tokens=KwargRule[int](min=10, max=model_max_tokens.get(model_name, 6144), default=1024, precision=0),
         )
 
     @classmethod

api/core/model_providers/providers/openai_provider.py

@@ -133,11 +133,11 @@ class OpenAIProvider(BaseModelProvider):
         }
 
         return ModelKwargsRules(
-            temperature=KwargRule[float](min=0, max=2, default=1),
-            top_p=KwargRule[float](min=0, max=1, default=1),
-            presence_penalty=KwargRule[float](min=-2, max=2, default=0),
-            frequency_penalty=KwargRule[float](min=-2, max=2, default=0),
-            max_tokens=KwargRule[int](min=10, max=model_max_tokens.get(model_name, 4097), default=16),
+            temperature=KwargRule[float](min=0, max=2, default=1, precision=2),
+            top_p=KwargRule[float](min=0, max=1, default=1, precision=2),
+            presence_penalty=KwargRule[float](min=-2, max=2, default=0, precision=2),
+            frequency_penalty=KwargRule[float](min=-2, max=2, default=0, precision=2),
+            max_tokens=KwargRule[int](min=10, max=model_max_tokens.get(model_name, 4097), default=16, precision=0),
         )
 
     @classmethod

api/core/model_providers/providers/openllm_provider.py

@@ -45,11 +45,11 @@ class OpenLLMProvider(BaseModelProvider):
         :return:
         """
         return ModelKwargsRules(
-            temperature=KwargRule[float](min=0.01, max=2, default=1),
-            top_p=KwargRule[float](min=0, max=1, default=0.7),
-            presence_penalty=KwargRule[float](min=-2, max=2, default=0),
-            frequency_penalty=KwargRule[float](min=-2, max=2, default=0),
-            max_tokens=KwargRule[int](alias='max_new_tokens', min=10, max=4000, default=128),
+            temperature=KwargRule[float](min=0.01, max=2, default=1, precision=2),
+            top_p=KwargRule[float](min=0, max=1, default=0.7, precision=2),
+            presence_penalty=KwargRule[float](min=-2, max=2, default=0, precision=2),
+            frequency_penalty=KwargRule[float](min=-2, max=2, default=0, precision=2),
+            max_tokens=KwargRule[int](alias='max_new_tokens', min=10, max=4000, default=128, precision=0),
         )
 
     @classmethod

api/core/model_providers/providers/replicate_provider.py

@@ -72,6 +72,7 @@ class ReplicateProvider(BaseModelProvider):
                         min=float(value.get('minimum')) if value.get('minimum') is not None else None,
                         max=float(value.get('maximum')) if value.get('maximum') is not None else None,
                         default=float(value.get('default')) if value.get('default') is not None else None,
+                        precision = 2
                     )
                     if key == 'temperature':
                         model_kwargs_rules.temperature = kwarg_rule
@@ -84,6 +85,7 @@ class ReplicateProvider(BaseModelProvider):
                         min=int(value.get('minimum')) if value.get('minimum') is not None else 1,
                         max=int(value.get('maximum')) if value.get('maximum') is not None else 8000,
                         default=int(value.get('default')) if value.get('default') is not None else 500,
+                        precision = 0
                     )
 
         return model_kwargs_rules