refactor: openai

apps/common/constants/permission_constants.py

@@ -154,6 +154,7 @@ class PermissionConstants(Enum):
                                                                                                RoleConstants.USER])
     KNOWLEDGE_MODULE_DELETE = Permission(group=Group.KNOWLEDGE, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
                                                                                                    RoleConstants.USER])
+
     def get_workspace_application_permission(self):
         return lambda r, kwargs: Permission(group=self.value.group, operate=self.value.operate,
                                             resource_path=

apps/models_provider/base_model_provider.py

@@ -100,6 +100,9 @@ class MaxKBBaseModel(ABC):
         optional_params = {}
         for key, value in model_kwargs.items():
             if key not in ['model_id', 'use_local', 'streaming', 'show_ref_label']:
+                if key == 'extra_body' and isinstance(value, dict):
+                    optional_params = {**optional_params, **value}
+                else:
                     optional_params[key] = value
         return optional_params
 

apps/models_provider/impl/base_chat_open_ai.py

@@ -1,15 +1,16 @@
 # coding=utf-8
-import warnings
+from typing import Dict, Optional, Any, Iterator, cast, Union, Sequence, Callable, Mapping
-from typing import List, Dict, Optional, Any, Iterator, cast, Type, Union
 
-import openai
-from langchain_core.callbacks import CallbackManagerForLLMRun
 from langchain_core.language_models import LanguageModelInput
-from langchain_core.messages import BaseMessage, get_buffer_string, BaseMessageChunk, AIMessageChunk
+from langchain_core.messages import BaseMessage, get_buffer_string, BaseMessageChunk, HumanMessageChunk, AIMessageChunk, \
-from langchain_core.outputs import ChatGenerationChunk, ChatGeneration
+    SystemMessageChunk, FunctionMessageChunk, ChatMessageChunk
+from langchain_core.messages.ai import UsageMetadata
+from langchain_core.messages.tool import tool_call_chunk, ToolMessageChunk
+from langchain_core.outputs import ChatGenerationChunk
 from langchain_core.runnables import RunnableConfig, ensure_config
-from langchain_core.utils.pydantic import is_basemodel_subclass
+from langchain_core.tools import BaseTool
 from langchain_openai import ChatOpenAI
+from langchain_openai.chat_models.base import _create_usage_metadata
 
 from common.config.tokenizer_manage_config import TokenizerManage
 
@@ -19,6 +20,65 @@ def custom_get_token_ids(text: str):
     return tokenizer.encode(text)
 
 
+def _convert_delta_to_message_chunk(
+        _dict: Mapping[str, Any], default_class: type[BaseMessageChunk]
+) -> BaseMessageChunk:
+    id_ = _dict.get("id")
+    role = cast(str, _dict.get("role"))
+    content = cast(str, _dict.get("content") or "")
+    additional_kwargs: dict = {}
+    if 'reasoning_content' in _dict:
+        additional_kwargs['reasoning_content'] = _dict.get('reasoning_content')
+    if _dict.get("function_call"):
+        function_call = dict(_dict["function_call"])
+        if "name" in function_call and function_call["name"] is None:
+            function_call["name"] = ""
+        additional_kwargs["function_call"] = function_call
+    tool_call_chunks = []
+    if raw_tool_calls := _dict.get("tool_calls"):
+        additional_kwargs["tool_calls"] = raw_tool_calls
+        try:
+            tool_call_chunks = [
+                tool_call_chunk(
+                    name=rtc["function"].get("name"),
+                    args=rtc["function"].get("arguments"),
+                    id=rtc.get("id"),
+                    index=rtc["index"],
+                )
+                for rtc in raw_tool_calls
+            ]
+        except KeyError:
+            pass
+
+    if role == "user" or default_class == HumanMessageChunk:
+        return HumanMessageChunk(content=content, id=id_)
+    elif role == "assistant" or default_class == AIMessageChunk:
+        return AIMessageChunk(
+            content=content,
+            additional_kwargs=additional_kwargs,
+            id=id_,
+            tool_call_chunks=tool_call_chunks,  # type: ignore[arg-type]
+        )
+    elif role in ("system", "developer") or default_class == SystemMessageChunk:
+        if role == "developer":
+            additional_kwargs = {"__openai_role__": "developer"}
+        else:
+            additional_kwargs = {}
+        return SystemMessageChunk(
+            content=content, id=id_, additional_kwargs=additional_kwargs
+        )
+    elif role == "function" or default_class == FunctionMessageChunk:
+        return FunctionMessageChunk(content=content, name=_dict["name"], id=id_)
+    elif role == "tool" or default_class == ToolMessageChunk:
+        return ToolMessageChunk(
+            content=content, tool_call_id=_dict["tool_call_id"], id=id_
+        )
+    elif role or default_class == ChatMessageChunk:
+        return ChatMessageChunk(content=content, role=role, id=id_)
+    else:
+        return default_class(content=content, id=id_)  # type: ignore
+
+
 class BaseChatOpenAI(ChatOpenAI):
     usage_metadata: dict = {}
     custom_get_token_ids = custom_get_token_ids
@@ -26,7 +86,13 @@ class BaseChatOpenAI(ChatOpenAI):
     def get_last_generation_info(self) -> Optional[Dict[str, Any]]:
         return self.usage_metadata
 
-    def get_num_tokens_from_messages(self, messages: List[BaseMessage]) -> int:
+    def get_num_tokens_from_messages(
+            self,
+            messages: list[BaseMessage],
+            tools: Optional[
+                Sequence[Union[dict[str, Any], type, Callable, BaseTool]]
+            ] = None,
+    ) -> int:
         if self.usage_metadata is None or self.usage_metadata == {}:
             try:
                 return super().get_num_tokens_from_messages(messages)
@@ -44,114 +110,77 @@ class BaseChatOpenAI(ChatOpenAI):
                 return len(tokenizer.encode(text))
         return self.get_last_generation_info().get('output_tokens', 0)
 
-    def _stream(
+    def _stream(self, *args: Any, **kwargs: Any) -> Iterator[ChatGenerationChunk]:
+        kwargs['stream_usage'] = True
+        for chunk in super()._stream(*args, **kwargs):
+            if chunk.message.usage_metadata is not None:
+                self.usage_metadata = chunk.message.usage_metadata
+            yield chunk
+
+    def _convert_chunk_to_generation_chunk(
             self,
-            messages: List[BaseMessage],
+            chunk: dict,
-            stop: Optional[List[str]] = None,
+            default_chunk_class: type,
-            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            base_generation_info: Optional[dict],
-            **kwargs: Any,
+    ) -> Optional[ChatGenerationChunk]:
-    ) -> Iterator[ChatGenerationChunk]:
+        if chunk.get("type") == "content.delta":  # from beta.chat.completions.stream
-        kwargs["stream"] = True
+            return None
-        kwargs["stream_options"] = {"include_usage": True}
+        token_usage = chunk.get("usage")
-        """Set default stream_options."""
+        choices = (
-        stream_usage = self._should_stream_usage(kwargs.get('stream_usage'), **kwargs)
+                chunk.get("choices", [])
-        # Note: stream_options is not a valid parameter for Azure OpenAI.
+                # from beta.chat.completions.stream
-        # To support users proxying Azure through ChatOpenAI, here we only specify
+                or chunk.get("chunk", {}).get("choices", [])
-        # stream_options if include_usage is set to True.
-        # See https://learn.microsoft.com/en-us/azure/ai-services/openai/whats-new
-        # for release notes.
-        if stream_usage:
-            kwargs["stream_options"] = {"include_usage": stream_usage}
-
-        payload = self._get_request_payload(messages, stop=stop, **kwargs)
-        default_chunk_class: Type[BaseMessageChunk] = AIMessageChunk
-        base_generation_info = {}
-
-        if "response_format" in payload and is_basemodel_subclass(
-                payload["response_format"]
-        ):
-            # TODO: Add support for streaming with Pydantic response_format.
-            warnings.warn("Streaming with Pydantic response_format not yet supported.")
-            chat_result = self._generate(
-                messages, stop, run_manager=run_manager, **kwargs
         )
-            msg = chat_result.generations[0].message
+
-            yield ChatGenerationChunk(
+        usage_metadata: Optional[UsageMetadata] = (
-                message=AIMessageChunk(
+            _create_usage_metadata(token_usage) if token_usage else None
-                    **msg.dict(exclude={"type", "additional_kwargs"}),
-                    # preserve the "parsed" Pydantic object without converting to dict
-                    additional_kwargs=msg.additional_kwargs,
-                ),
-                generation_info=chat_result.generations[0].generation_info,
         )
-            return
+        if len(choices) == 0:
-        if self.include_response_headers:
+            # logprobs is implicitly None
-            raw_response = self.client.with_raw_response.create(**payload)
+            generation_chunk = ChatGenerationChunk(
-            response = raw_response.parse()
+                message=default_chunk_class(content="", usage_metadata=usage_metadata)
-            base_generation_info = {"headers": dict(raw_response.headers)}
-        else:
-            response = self.client.create(**payload)
-        with response:
-            is_first_chunk = True
-            for chunk in response:
-                if not isinstance(chunk, dict):
-                    chunk = chunk.model_dump()
-
-                generation_chunk = super()._convert_chunk_to_generation_chunk(
-                    chunk,
-                    default_chunk_class,
-                    base_generation_info if is_first_chunk else {},
             )
-                if generation_chunk is None:
+            return generation_chunk
-                    continue
 
-                # custom code
+        choice = choices[0]
-                if len(chunk['choices']) > 0 and 'reasoning_content' in chunk['choices'][0]['delta']:
+        if choice["delta"] is None:
-                    generation_chunk.message.additional_kwargs["reasoning_content"] = chunk['choices'][0]['delta'][
+            return None
-                        'reasoning_content']
 
-                default_chunk_class = generation_chunk.message.__class__
+        message_chunk = _convert_delta_to_message_chunk(
-                logprobs = (generation_chunk.generation_info or {}).get("logprobs")
+            choice["delta"], default_chunk_class
-                if run_manager:
-                    run_manager.on_llm_new_token(
-                        generation_chunk.text, chunk=generation_chunk, logprobs=logprobs
         )
-                is_first_chunk = False
+        generation_info = {**base_generation_info} if base_generation_info else {}
-                # custom code
-                if generation_chunk.message.usage_metadata is not None:
-                    self.usage_metadata = generation_chunk.message.usage_metadata
-                yield generation_chunk
 
-    def _create_chat_result(self,
+        if finish_reason := choice.get("finish_reason"):
-                            response: Union[dict, openai.BaseModel],
+            generation_info["finish_reason"] = finish_reason
-                            generation_info: Optional[Dict] = None):
+            if model_name := chunk.get("model"):
-        result = super()._create_chat_result(response, generation_info)
+                generation_info["model_name"] = model_name
-        try:
+            if system_fingerprint := chunk.get("system_fingerprint"):
-            reasoning_content = ''
+                generation_info["system_fingerprint"] = system_fingerprint
-            reasoning_content_enable = False
+
-            for res in response.choices:
+        logprobs = choice.get("logprobs")
-                if 'reasoning_content' in res.message.model_extra:
+        if logprobs:
-                    reasoning_content_enable = True
+            generation_info["logprobs"] = logprobs
-                    _reasoning_content = res.message.model_extra.get('reasoning_content')
+
-                    if _reasoning_content is not None:
+        if usage_metadata and isinstance(message_chunk, AIMessageChunk):
-                        reasoning_content += _reasoning_content
+            message_chunk.usage_metadata = usage_metadata
-            if reasoning_content_enable:
+
-                result.llm_output['reasoning_content'] = reasoning_content
+        generation_chunk = ChatGenerationChunk(
-        except Exception as e:
+            message=message_chunk, generation_info=generation_info or None
-            pass
+        )
-        return result
+        return generation_chunk
 
     def invoke(
             self,
             input: LanguageModelInput,
             config: Optional[RunnableConfig] = None,
             *,
-            stop: Optional[List[str]] = None,
+            stop: Optional[list[str]] = None,
             **kwargs: Any,
     ) -> BaseMessage:
         config = ensure_config(config)
         chat_result = cast(
-            ChatGeneration,
+            "ChatGeneration",
             self.generate_prompt(
                 [self._convert_input(input)],
                 stop=stop,
@@ -162,7 +191,9 @@ class BaseChatOpenAI(ChatOpenAI):
                 run_id=config.pop("run_id", None),
                 **kwargs,
             ).generations[0][0],
+
         ).message
+
         self.usage_metadata = chat_result.response_metadata[
             'token_usage'] if 'token_usage' in chat_result.response_metadata else chat_result.usage_metadata
         return chat_result