File size: 1,633 Bytes
7626706 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
from dataclasses import dataclass
from llama_cpp import Llama
from llama_cpp import ChatCompletionRequestMessage as Message
from llama_cpp import ChatCompletionRequestSystemMessage as SystemMessage
from llama_cpp import ChatCompletionRequestAssistantMessage as AssistantMessage
from llama_cpp import ChatCompletionRequestUserMessage as UserMessage
@dataclass
class MessageRole:
ASSISTANT: str = "assistant"
SYSTEM: str = "system"
USER: str = "user"
EXIT: str = "exit"
class ConversationHandler:
def __init__(self, model: Llama, message_role: MessageRole) -> None:
self.model: Llama = model
self.message_role = message_role
self.messages: list[Message] = [
SystemMessage(
role=self.message_role.SYSTEM,
content='You are a helpful developer assistant, answer all the questions correctly and concisely.'
),
AssistantMessage(role=self.message_role.ASSISTANT, content='Hello, do you have any question?'),
]
def send_message(self, content: str):
new_message = UserMessage(role=self.message_role.USER, content=content)
self.messages.append(new_message)
def generate_reply(self) -> str:
response = self.model.create_chat_completion(
messages=self.messages,
temperature=0.7,
top_p=0.9,
top_k=20,
max_tokens=128
)
response_content = response['choices'][0]['message']
self.messages.append(AssistantMessage(role=self.message_role.ASSISTANT, content=response_content))
return response_content
|