faiss_rag_enterprise/llama_index/multi_modal_llms/dashscope_utils.py

75 lines
2.3 KiB
Python

"""DashScope api utils."""
from http import HTTPStatus
from typing import Any, Dict, List, Sequence
from llama_index.core.llms.types import (
ChatMessage,
ChatResponse,
CompletionResponse,
)
from llama_index.schema import ImageDocument
def dashscope_response_to_completion_response(response: Any) -> CompletionResponse:
if response["status_code"] == HTTPStatus.OK:
content = response["output"]["choices"][0]["message"]["content"]
if content:
content = content[0]["text"]
else:
content = ""
return CompletionResponse(text=content, raw=response)
else:
return CompletionResponse(text="", raw=response)
def dashscope_response_to_chat_response(
response: Any,
) -> ChatResponse:
if response["status_code"] == HTTPStatus.OK:
content = response["output"]["choices"][0]["message"]["content"]
role = response["output"]["choices"][0]["message"]["role"]
return ChatResponse(
message=ChatMessage(role=role, content=content), raw=response
)
else:
return ChatResponse(message=ChatMessage(), raw=response)
def chat_message_to_dashscope_multi_modal_messages(
chat_messages: Sequence[ChatMessage],
) -> List[Dict]:
messages = []
for msg in chat_messages:
messages.append({"role": msg.role.value, "content": msg.content})
return messages
def create_dashscope_multi_modal_chat_message(
prompt: str, role: str, image_documents: Sequence[ImageDocument]
) -> ChatMessage:
if image_documents is None:
message = ChatMessage(role=role, content=[{"text": prompt}])
else:
content = []
for image_document in image_documents:
content.append(
{
"image": image_document.image_url
if image_document.image_url is not None
else image_document.image_path
}
)
content.append({"text": prompt})
message = ChatMessage(role=role, content=content)
return message
def load_local_images(local_images: List[str]) -> List[ImageDocument]:
# load images into image documents
image_documents = []
for _, img in enumerate(local_images):
new_image_document = ImageDocument(image_path=img)
image_documents.append(new_image_document)
return image_documents