PersonaAgent: When Large Language Model Agents Meet Personalization at Test Time

# PersonaAgent core loop implementation sketch (LangChain-based)

from langchain.chat_models import ChatAnthropic
from langchain.vectorstores import FAISS
from langchain.embeddings import BedrockEmbeddings

# 1. Initialize Persona (system prompt) per user
def init_persona(user_profile_summary: str) -> str:
    return f"""You are a helpful personalized assistant.
User summary: {user_profile_summary}
STRICT RULES:
1. Think step-by-step about what information you need.
2. MUST use at least TWO tools to answer the question.
3. Provide clear, concise responses."""

# 2. Episodic Memory storage and retrieval
class EpisodicMemory:
    def __init__(self, embeddings):
        self.store = {}  # user_id -> FAISS index
        self.embeddings = embeddings

    def add(self, user_id, query, response, metadata={}):
        # In practice, store in FAISS or another vector DB
        pass

    def retrieve(self, user_id, query, top_k=4):
        # Return Top-K interactions based on similarity
        pass

# 3. Core logic for Test-Time Persona optimization
def compute_textual_gradient(llm, question, agent_response, ground_truth):
    """Have the LLM analyze response differences to generate feedback"""
    prompt = f"""You are a meticulous evaluator of personalized AI agent responses.
Analyze the following and give feedback on how to improve the system prompt.

Question: {question}
Expected Answer: {ground_truth}
Agent Response: {agent_response}

Feedback should focus on:
1. How to improve search keywords for this user.
2. User's prior interactions and preferences.
3. Explicit user profile descriptions not specific to this task.

Feedback:"""
    return llm.predict(prompt)

def update_persona(llm, current_persona, feedbacks):
    """Aggregate feedback and update Persona"""
    aggregated = "\n".join(feedbacks)
    prompt = f"""You are a prompt engineering assistant.
Current system prompt: {current_persona}
Provided Feedback: {aggregated}

Generate an updated system prompt that highlights the user's unique preferences.
New system prompt:"""
    return llm.predict(prompt)

# 4. Test-time alignment loop
def test_time_alignment(llm, agent, persona, recent_interactions, n_iter=1):
    for _ in range(n_iter):
        feedbacks = []
        for q, gt in recent_interactions:
            agent_resp = agent.run(q, system_prompt=persona)
            feedback = compute_textual_gradient(llm, q, agent_resp, gt)
            feedbacks.append(feedback)
        persona = update_persona(llm, persona, feedbacks)
    return persona