Skip to content

Commit

Permalink
refactor: Remove usage of deprecated ChatMessage.to_openai_format
Browse files Browse the repository at this point in the history
  • Loading branch information
shadeMe committed Aug 19, 2024
1 parent 13cdaab commit 8932f0d
Showing 1 changed file with 4 additions and 16 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from typing import Any, Callable, ClassVar, Dict, List, Optional

from botocore.eventstream import EventStream
from haystack.components.generators.openai_utils import _convert_message_to_openai_format
from haystack.dataclasses import ChatMessage, ChatRole, StreamingChunk
from transformers import AutoTokenizer, PreTrainedTokenizer

Expand Down Expand Up @@ -389,27 +390,14 @@ def prepare_chat_messages(self, messages: List[ChatMessage]) -> str:
# default is https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1/blob/main/tokenizer_config.json
# but we'll use our custom chat template
prepared_prompt: str = self.prompt_handler.tokenizer.apply_chat_template(
conversation=[self.to_openai_format(m) for m in messages], tokenize=False, chat_template=self.chat_template
conversation=[_convert_message_to_openai_format(m) for m in messages],
tokenize=False,
chat_template=self.chat_template,
)
if self.truncate:
prepared_prompt = self._ensure_token_limit(prepared_prompt)
return prepared_prompt

def to_openai_format(self, m: ChatMessage) -> Dict[str, Any]:
"""
Convert the message to the format expected by OpenAI's Chat API.
See the [API reference](https://platform.openai.com/docs/api-reference/chat/create) for details.
:returns: A dictionary with the following key:
- `role`
- `content`
- `name` (optional)
"""
msg = {"role": m.role.value, "content": m.content}
if m.name:
msg["name"] = m.name
return msg

def check_prompt(self, prompt: str) -> Dict[str, Any]:
"""
Checks the prompt length and resizes it if necessary. If the prompt is too long, it will be truncated.
Expand Down

0 comments on commit 8932f0d

Please sign in to comment.