jalr/tabbyAPI-ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: list support in message.content (#122)

Browse source
This commit is contained in:
Orion 2024-06-04 01:57:15 +08:00 • committed by GitHub
parent 1951f7521c
commit 6cc3bd9752
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 15 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

3
endpoints/OAI/types/chat_completion.py
View file

@ -41,7 +41,8 @@ class ChatCompletionStreamChoice(BaseModel):
class ChatCompletionRequest(CommonCompletionRequest):
# Messages
# Take in a string as well even though it's not part of the OAI spec
messages: Union[str, List[Dict[str, str]]]
# support messages.content as a list of dict
messages: Union[str, List[Dict[str, Union[str, List[Dict[str, str]]]]]]
prompt_template: Optional[str] = None
add_generation_prompt: Optional[bool] = True
template_vars: Optional[dict] = {}

13
endpoints/OAI/utils/chat_completion.py
View file

@ -151,6 +151,19 @@ def format_prompt_with_template(data: ChatCompletionRequest):
unwrap(data.ban_eos_token, False),
)
# Deal with list in messages.content
# Just replace the content list with the very first text message
for message in data.messages:
if message["role"] == "user" and isinstance(message["content"], list):
message["content"] = next(
(
content["text"]
for content in message["content"]
if content["type"] == "text"
),
"",
)
# Overwrite any protected vars with their values
data.template_vars.update(
{