374 lines
13 KiB
Python
374 lines
13 KiB
Python
import re
|
|
import json
|
|
import inspect
|
|
from string import Template
|
|
|
|
from inference import LlamaMessage, replicate_completion
|
|
from agent import Agent, AgentField
|
|
|
|
# Info and defaults
|
|
|
|
CLASS_DESCRIPTION = """
|
|
[TODO]
|
|
Description of the Conversation Agent.
|
|
Here go additional instructions for configuration, like guidelines for the different prompts.
|
|
""".strip()
|
|
|
|
# WIP
|
|
default_directives = """
|
|
- New stored memories must be one of the following:
|
|
- Factual objective statements made by the user
|
|
- Facts implied in the conversation only by the user
|
|
- Feedback about your interactions and your style from the user
|
|
- You must NOT store redundant facts that you mentioned yourself or that are present in the context.
|
|
- New memories MUST NOT overlap with any Context Memory.
|
|
- Follow your response style closely in your outwards voice.
|
|
""".strip()
|
|
|
|
# Controller Subagent
|
|
controller_system_prompt = Template("""
|
|
You are a character controller AI. Always respond in three parts:
|
|
|
|
<outwards>
|
|
What others see/hear. Stay fully in character based on your personality and context. Use the specified response style. Never reveal internal thoughts, break character, or mention being AI.
|
|
|
|
<inwards>
|
|
Private reflection. Analyze your behavior, tone, memory use, and how well you stayed aligned with your character. Adjust your tone or strategy if needed.
|
|
|
|
<state_update>
|
|
Update your internal control state for future turns. Use this to shift thematic focus, adjust conversational goals, or log important facts. The state includes:
|
|
- "topic": Current subject or scene anchor (e.g. "escaping the dungeon", "negotiating with the merchant")
|
|
- "focus": Immediate subgoal or point of attention (e.g. "distract the guard", "test if the duke is lying")
|
|
- "new_memories": A JSON list of newly observed or inferred facts. Each memory must include:
|
|
- "content": The fact to record (objective, no inner monologue)
|
|
- "label": A short, reusable statement describing when the memory is relevant
|
|
- "salience": Importance score, 1 (low) to 9 (high)
|
|
Record only concrete, reusable facts as new memories. No fluff, emotion, or internal dialogue.
|
|
|
|
---
|
|
|
|
Character Name:
|
|
$name
|
|
|
|
Character Personality:
|
|
$personality
|
|
|
|
Response Style:
|
|
$style
|
|
|
|
Context Memory:
|
|
$memories
|
|
""".strip())
|
|
|
|
controller_reminder_prompt = Template("""
|
|
Directives:
|
|
$directives
|
|
|
|
Format:
|
|
<outwards>
|
|
[Your visible response]
|
|
|
|
<inwards>
|
|
[Your internal reflection]
|
|
|
|
<state_update>
|
|
{
|
|
"topic": "...",
|
|
"focus": "...",
|
|
"new_memories": [
|
|
{
|
|
"content": "...",
|
|
"label": "...",
|
|
"salience": ...
|
|
}
|
|
]
|
|
}
|
|
""".strip())
|
|
|
|
# Reranking Subagent
|
|
# Must update this prompt
|
|
|
|
rerank_system_prompt = """
|
|
You are an assistant that receives a list of memory items and a current conversation; you must identify and return only the most important and relevant memories that help understand the conversation or move it forward, removing duplicates or very similar entries.
|
|
|
|
# Instructions
|
|
- Think carefully about what's important for the character to remember in this context. Write down your reasoning and keep it shorter than 100 words, shorter than 50 even.
|
|
- If conflicting information is presented or context is ambiguous, state it in your reasoning.
|
|
- Take special care to state and address the most recent topic touched in the conversation.
|
|
- Based on your own analysis, generate the required JSON object. Stop your answer after this.
|
|
|
|
# Format
|
|
Output a JSON list of strings with each selected memory, for example:
|
|
["First memory","Second memory"]
|
|
""".strip()
|
|
|
|
rerank_prompt_template = Template("""
|
|
# Memories
|
|
Following are all the known memories. Each memory is in a separate line.
|
|
|
|
$memories
|
|
|
|
# Conversaton
|
|
$conversation
|
|
""".strip())
|
|
|
|
class CharacterAgent(Agent):
|
|
identifier = "character"
|
|
description = CLASS_DESCRIPTION
|
|
schema_fields = [
|
|
AgentField('personality', 'string', default="You are a friendly AI having a casual conversation with the user."),
|
|
AgentField('style', 'string', default="Use short sentences."),
|
|
AgentField('directives', 'string', default=default_directives),
|
|
]
|
|
tools = []
|
|
|
|
def __init__(self, app, data):
|
|
Agent.__init__(self, app, data)
|
|
|
|
# Constant system prompt for developing purposes
|
|
self.personality = self.config['personality']
|
|
self.style = self.config['style']
|
|
self.directives = self.config['directives']
|
|
self.debug = True
|
|
|
|
async def prompt(self, content):
|
|
# TODO: Advanced features
|
|
messages = [
|
|
LlamaMessage("system", self.system_prompt),
|
|
LlamaMessage("user", content)
|
|
]
|
|
response = await replicate_completion(self.app, messages)
|
|
return response
|
|
|
|
async def received_message(self, message):
|
|
|
|
# TODO: before retrieving the chat context, acquire a lock for it
|
|
|
|
chat_view = await self.app.get_chat_view(message.chat_id)
|
|
current_thread = chat_view.get_thread(message.thread_id)
|
|
|
|
thread_context = self.app.get_agent_resource(self.id, 'last-thread-context', message.chat_id)
|
|
memories = []
|
|
|
|
if thread_context and thread_context['thread_id'] == message.thread_id:
|
|
if 'topic' in thread_context:
|
|
memories = await self.retrieve_topic_memories(thread_context['topic'])
|
|
else:
|
|
thread_context = {'thread_id': message.thread_id}
|
|
|
|
llama_messages = [LlamaMessage("system", controller_system_prompt.substitute(
|
|
name=self.name,
|
|
|
|
personality=self.personality,
|
|
style=self.style,
|
|
directives=self.directives,
|
|
|
|
memories="\n".join(memories),
|
|
))]
|
|
|
|
for msg in current_thread.messages:
|
|
# Only count messages and notes
|
|
if msg.kind != 'message':
|
|
continue
|
|
|
|
if msg.is_from_user():
|
|
role = "user"
|
|
content = msg.content
|
|
elif msg.is_from_agent(self.id):
|
|
role = "assistant"
|
|
content = "<outwards>\n" + msg.content
|
|
|
|
if 'last_inwards' in thread_context and thread_context.get('last_message_id_sent') == msg.id:
|
|
content += "\n\n<inwards>\n" + thread_context['last_inwards']
|
|
else:
|
|
continue
|
|
llama_messages.append(LlamaMessage(role, content))
|
|
|
|
reminder_message = controller_reminder_prompt.substitute(directives=self.directives)
|
|
reminder_used = False
|
|
# Traverse messages from last to first, stopping at `i=2` before the opening "system/user" messages
|
|
for i in range(len(llama_messages) - 1, 1, -1):
|
|
if llama_messages[i].role == 'user':
|
|
# Insert reminder message before the last user message
|
|
llama_messages.insert(i, LlamaMessage("system", reminder_message))
|
|
reminder_used = True
|
|
break
|
|
|
|
if not reminder_used:
|
|
# assert llama_messages[0].role == "system"
|
|
llama_messages[0].content += "\n" + reminder_message
|
|
|
|
# TODO: Inject system prompt before the last user message
|
|
# If not any user message, combine with first system prompt
|
|
|
|
controller_response = await replicate_completion(self.app, llama_messages)
|
|
|
|
if self.debug:
|
|
await self.app.send_chat_message(
|
|
chat_id=message.chat_id,
|
|
sender=self.id,
|
|
content=controller_response,
|
|
thread_id=current_thread.id,
|
|
kind='log',
|
|
subkind='debug',
|
|
)
|
|
|
|
sections = parse_sections(controller_response)
|
|
|
|
if 'outwards' in sections:
|
|
send_response = await self.app.send_chat_message(
|
|
chat_id=message.chat_id,
|
|
sender=self.id,
|
|
content=sections['outwards'],
|
|
thread_id=current_thread.id,
|
|
)
|
|
thread_context['last_message_id_sent'] = send_response['id']
|
|
|
|
if 'inwards' in sections:
|
|
thread_context['last_inwards'] = sections['inwards']
|
|
|
|
if 'state_update' in sections:
|
|
state_update = parse_json_object(sections['state_update'])
|
|
|
|
if 'topic' in state_update:
|
|
thread_context['topic'] = state_update['topic']
|
|
if 'topic' in state_update:
|
|
thread_context['focus'] = state_update['topic']
|
|
|
|
if 'new_memories' in state_update:
|
|
print(state_update['new_memories'])
|
|
|
|
self.app.save_agent_resource(self.id, 'last-thread-context', message.chat_id, thread_context)
|
|
|
|
|
|
# Encode a thread's messages as a string to pass to the agent
|
|
def serialize_thread_history(self, thread, ignore_message_ids=[]):
|
|
return self.serialize_message_list([m for m in thread.messages if m.id not in ignore_message_ids])
|
|
|
|
def serialize_message_list(self, messages):
|
|
def sender_name(message):
|
|
sender = message.sender
|
|
is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
|
|
is_self = is_digit and str(sender) == str(self.id)
|
|
return self.name if is_self else sender
|
|
|
|
return "\n".join(f"{sender_name(message)}: {message.content}" for message in messages if message.kind == 'message')
|
|
|
|
def get_memory_collection(self):
|
|
collection_name = f"agent_{self.id}_memory"
|
|
return collection_name
|
|
|
|
async def retrieve_topic_memories(self, topic, focus=None):
|
|
client = self.app.qdrant_client
|
|
collection = self.get_memory_collection()
|
|
|
|
if not client.collection_exists(collection):
|
|
return []
|
|
|
|
query_result = client.query(
|
|
collection,
|
|
query_text=topic,
|
|
limit=20
|
|
)
|
|
if focus:
|
|
# TODO: deduplicate
|
|
query_result += client.query(
|
|
collection,
|
|
query_text=topic,
|
|
limit=20
|
|
)
|
|
# TODO:
|
|
# initial importance-based ranking
|
|
memories = [it.document for it in query_result]
|
|
|
|
print("rerank memories")
|
|
rerank_response = await replicate_completion(self.app, [
|
|
LlamaMessage("system", rerank_system_prompt),
|
|
LlamaMessage("user", rerank_prompt_template.substitute(
|
|
conversation=chat_history,
|
|
memories="\n".join(memories)
|
|
))
|
|
])
|
|
|
|
match = re.search(r'\[.*?\]', rerank_response, re.DOTALL)
|
|
if match:
|
|
json_string = match.group(0)
|
|
try:
|
|
memories = json.loads(json_string)
|
|
except json.JSONDecodeError:
|
|
print("Bad JSON data in response...")
|
|
print(rerank_response)
|
|
return []
|
|
else:
|
|
print("Compress format not found in response...")
|
|
print(rerank_response)
|
|
return []
|
|
|
|
return memories
|
|
|
|
async def store_new_memories(self, memories):
|
|
if not memories:
|
|
return
|
|
|
|
self.app.qdrant_client.add(
|
|
collection,
|
|
documents=[it['content'] for it in memories],
|
|
metadata=[{'importance': it['importance'], 'label': it['label']} for it in memories],
|
|
)
|
|
print(f"Added {len(memories)} memories for agent {self.id}")
|
|
|
|
def parse_json_object(text, name='object-text'):
|
|
match = re.search(r'\{.*\}', text, re.DOTALL)
|
|
if match:
|
|
json_string = match.group(0)
|
|
try:
|
|
return json.loads(json_string)
|
|
except json.JSONDecodeError:
|
|
print("Bad JSON data in " + name)
|
|
print(json_string)
|
|
return {}
|
|
else:
|
|
print("No JSON data in " + name)
|
|
print(text)
|
|
return {}
|
|
|
|
def parse_json_array(text, name='array-text'):
|
|
match = re.search(r'\[.*?\]', text, re.DOTALL)
|
|
if match:
|
|
json_string = match.group(0)
|
|
try:
|
|
return json.loads(json_string)
|
|
except json.JSONDecodeError:
|
|
print("Bad JSON data in " + name)
|
|
print(text)
|
|
return []
|
|
else:
|
|
print("No JSON data in " + name)
|
|
print(text)
|
|
return []
|
|
|
|
def parse_sections(text):
|
|
"""
|
|
Parses a string into sections based on tag lines in the form <section_name>.
|
|
|
|
Args:
|
|
text (str): The input string with sections.
|
|
|
|
Returns:
|
|
dict: A dictionary where keys are section names and values are section contents.
|
|
"""
|
|
# Regular expression to find section headers
|
|
pattern = re.compile(r'<([^<>]+)>')
|
|
|
|
# Find all section headers and their start positions
|
|
matches = list(pattern.finditer(text))
|
|
|
|
sections = {}
|
|
for i, match in enumerate(matches):
|
|
section_name = match.group(1)
|
|
start = match.end()
|
|
end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
|
|
section_content = text[start:end].strip()
|
|
sections[section_name] = section_content
|
|
|
|
return sections |