Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: Tokens cannot be obtained from the model dialogue #2326

Merged
merged 1 commit into from
Feb 19, 2025
Merged

fix: Tokens cannot be obtained from the model dialogue #2326

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions apps/setting/models_provider/impl/base_chat_open_ai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,8 @@ def _stream(
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> Iterator[ChatGenerationChunk]:

kwargs["stream"] = True
kwargs["stream_options"] = {"include_usage": True}
"""Set default stream_options."""
stream_usage = self._should_stream_usage(kwargs.get('stream_usage'), **kwargs)
# Note: stream_options is not a valid parameter for Azure OpenAI.
Expand All @@ -63,7 +64,6 @@ def _stream(
if stream_usage:
kwargs["stream_options"] = {"include_usage": stream_usage}

kwargs["stream"] = True
payload = self._get_request_payload(messages, stop=stop, **kwargs)
default_chunk_class: Type[BaseMessageChunk] = AIMessageChunk
base_generation_info = {}
Expand Down Expand Up @@ -107,9 +107,6 @@ def _stream(
continue

# custom code
if generation_chunk.message.usage_metadata is not None:
self.usage_metadata = generation_chunk.message.usage_metadata
# custom code
if len(chunk['choices']) > 0 and 'reasoning_content' in chunk['choices'][0]['delta']:
generation_chunk.message.additional_kwargs["reasoning_content"] = chunk['choices'][0]['delta'][
'reasoning_content']
Expand All @@ -121,6 +118,9 @@ def _stream(
generation_chunk.text, chunk=generation_chunk, logprobs=logprobs
)
is_first_chunk = False
# 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,
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The provided code snippet appears to be a part of an implementation for handling streaming responses from a chat model, specifically within a class called Stream. Here's a breakdown with notes on potential improvements:

Key Issues Identified:

  1. Repeated Stream Option Setting: The stream option (stream attribute) is set multiple times without considering whether it was already defined.
  2. Redundant Usage Metadata Handling: There seems to be redundant logic around usage metadata retrieval.

Potential Improvements:

  1. Single Stream Option Check:
    Ensure that the stream option is only set once, ideally during initialization or just before initiating the stream process. This reduces ambiguity and potential bugs.
# Single check for stream option
if not kwargs.get('stream', False):
    kwargs["stream"] = True
  1. Avoid Redundant Usage Metadata Retrieval:
    Remove unnecessary checks and assignments related to retrieving usage metadata because it might lead to overwriting values if they are fetched more than once.
# Custom code removed for clarity
  1. Enhance Error Handling (Optional):
    While not directly addressed in this snippet, consider adding error handling mechanisms to manage edge cases such as invalid inputs or timeouts when fetching response chunks.

Suggested Changes:

Here’s how you could refactor the function based on these guidelines:

def _stream(
    messages: list[base_message.BaseMessage],
    stop: Optional[list[str]] = None,
    callbacks: Optional[List[CallbackHandler]] = None,
    verbose: bool = False,
    use_cache: bool = True,  # Assuming there's a need for caching
    llm_backend=None,
    run_manager: Optional[ CallbackManagerForLLMRun] = None,
    **kwargs: Any,
) -> Iterator[ChatGenerationChunk]:
    """
    Set default stream_options and initiate streaming response.
    """
    if llm_backend == "azure":
        del kwargs["stream"]
    
    # Ensure stream option is set correctly
    kwargs["stream"] = kwargs.get("stream", False)
    
    if kwargs.get("stream"):
        # Additional setup for streaming can go here
    
    payload = self._get_request_payload(messages, stop=stop, use_cache=use_cache, **kwargs)
    default_chunk_class: Type[BaseMessageChunk] = AIMessageChunk
    base_generation_info = {}
    
    # Rest of the code remains mostly unchanged

# Example usage
async def main():
    async for chunk in client.stream(["Hello"], use_stream=True):
        print(chunk.text)

Conclusion:

By ensuring the stream option is consistently managed and avoiding redundant operations concerning usage metadata, we improve both readability and robustness of the _stream method. These changes also enhance efficiency and reliability while maintaining consistency throughout the implementation.

Expand Down