93 lines
3.2 KiB
Python
93 lines
3.2 KiB
Python
from stream import Stream
|
|
import threading
|
|
|
|
LLAMA3_MODEL = "meta/meta-llama-3-8b-instruct"
|
|
|
|
# Instances of this class lives very briefly, created when needed and immediately disposed.
|
|
# Subclasses must be careful in not setting up long term resources
|
|
class AgentControl:
|
|
def __init__(self, app, agent_id):
|
|
self.app = app
|
|
self.agent_id = agent_id
|
|
self.system_prompt = "You are a useful and precise AI assistant. Respond concisely."
|
|
|
|
# Receives a single string prompt, and must respond with a string
|
|
async def prompt(self, content):
|
|
response = await replicate_llama_completion(self.app, [
|
|
LlamaMessage("system", self.system_prompt),
|
|
LlamaMessage("user", content)
|
|
])
|
|
return response
|
|
|
|
async def received_message(self, message):
|
|
chat_context = self.app.get_chat_context(message.chat_id)
|
|
def sender_name(message):
|
|
sender = message['sender']
|
|
is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
|
|
return "assistant" if is_digit else sender
|
|
chat_summary = "\n".join(f"**{sender_name(message)}**: {message['content']}" for message in chat_context.messages)
|
|
|
|
prologue = "\nYou are given the entire context of the conversation, respond with your next message in the conversation."
|
|
print(chat_summary)
|
|
|
|
response = await replicate_llama_completion(self.app, [
|
|
LlamaMessage("system", self.system_prompt + prologue),
|
|
LlamaMessage("user", chat_summary)
|
|
])
|
|
|
|
self.app.send_chat_message(
|
|
chat_id=message.chat_id,
|
|
sender=self.agent_id,
|
|
content=response,
|
|
)
|
|
|
|
def create_from_data(app, agent_id, kind, data):
|
|
return AgentControl(app, agent_id)
|
|
|
|
class LlamaMessage:
|
|
def __init__(self, role, content):
|
|
self.role = role
|
|
self.content = content
|
|
|
|
async def replicate_llama_completion(app, messages):
|
|
# Llama 3 instruct
|
|
full_prompt = "<|begin_of_text|>"
|
|
for message in messages:
|
|
header_name = message.role if message.role != "agent" else "assistant"
|
|
full_prompt += f"<|start_header_id|>{header_name}<|end_header_id|>\n\n{message.content}<|eot_id|>"
|
|
|
|
full_prompt += f"<|start_header_id|>assistant<|end_header_id|>"
|
|
|
|
content_stream = Stream(app)
|
|
thread = threading.Thread(
|
|
target=run_replicate_inference_loop,
|
|
kwargs={
|
|
'replicate_client': app.replicate_client,
|
|
'model': LLAMA3_MODEL,
|
|
'prompt': full_prompt,
|
|
'content_stream': content_stream,
|
|
}
|
|
)
|
|
|
|
thread.start()
|
|
response = await content_stream.content_wait()
|
|
|
|
return response.lstrip()
|
|
|
|
|
|
def run_replicate_inference_loop(replicate_client, model, prompt, content_stream):
|
|
print("START INFERENCE")
|
|
completion = replicate_client.run(
|
|
model,
|
|
input={
|
|
"prompt": prompt,
|
|
"temperature": 0.75,
|
|
"max_new_tokens": 500,
|
|
"system_prompt": "" # We handle system prompt in the formatted prompt
|
|
}
|
|
)
|
|
for event in completion:
|
|
if event:
|
|
content_stream.feed(str(event))
|
|
print("END INFERENCE")
|
|
content_stream.close() |