← Back to DevBytes

Agent Memory Architectures with vLLM: Complete Guide

Introduction to Agent Memory Architectures with vLLM

Building autonomous agents that can reason, plan, and act over extended interactions requires more than just a powerful language model — it requires a robust memory architecture. When you pair an agent's memory system with vLLM, a high-throughput, memory-efficient inference engine, you unlock the ability to serve agents at scale without sacrificing latency or context quality.

This guide walks through the theory and practice of agent memory architectures, shows how vLLM fits into the stack, and provides working code you can adapt for production deployments.

What Is an Agent Memory Architecture?

An agent memory architecture defines how an autonomous agent stores, retrieves, and uses information across interactions. Unlike a stateless chatbot that only sees the current prompt, a memory-enabled agent can recall past conversations, reference long-term knowledge, and maintain working memory for multi-step reasoning.

The Three Core Memory Types

Why Memory Matters for Agents

Without memory, agents are amnesiac. They repeat mistakes, forget user preferences, and cannot perform tasks that span multiple sessions. A well-designed memory architecture enables:

Why vLLM for Agent Memory Systems?

vLLM is an open-source inference engine optimized for high-throughput LLM serving. Its features align particularly well with agent workloads:

For memory-heavy agents, the combination of prefix caching and efficient KV management means you can serve dozens of concurrent agents with large context windows on a single GPU.

Architecture Overview

A production agent memory stack with vLLM typically has these layers:

Setting Up vLLM

First, install vLLM and launch a server. We'll use a mid-size model suitable for agent reasoning.

# Install vLLM
pip install vllm

# Launch the server with prefix caching enabled
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Meta-Llama-3-8B-Instruct \
  --enable-prefix-caching \
  --max-model-len 8192 \
  --port 8000

The --enable-prefix-caching flag is critical for agents. When your memory manager prepends a shared system prompt or retrieved context to many requests, vLLM reuses the cached KV states, reducing time-to-first-token significantly.

Building the Memory Manager

The memory manager is the heart of the architecture. It handles storage, retrieval, and context assembly. Below is a practical implementation using ChromaDB for long-term memory and an in-memory store for working memory.

import chromadb
import json
from datetime import datetime
from openai import OpenAI

class AgentMemory:
    def __init__(self, agent_id: str, vllm_base_url: str = "http://localhost:8000/v1"):
        self.agent_id = agent_id
        self.client = OpenAI(base_url=vllm_base_url, api_key="vllm")

        # Long-term semantic memory
        self.chroma = chromadb.PersistentClient(path="./memory_store")
        self.collection = self.chroma.get_or_create_collection(
            name=f"agent_{agent_id}",
            metadata={"hnsw:space": "cosine"}
        )

        # Short-term working memory
        self.working_memory: list[dict] = []
        self.max_working_messages = 20

    def add_to_working_memory(self, role: str, content: str):
        """Add a message to the rolling working memory buffer."""
        self.working_memory.append({
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        })
        # Keep working memory bounded
        if len(self.working_memory) > self.max_working_messages:
            # Offload oldest messages to long-term memory
            old = self.working_memory.pop(0)
            self.consolidate_to_long_term(old)

    def consolidate_to_long_term(self, message: dict):
        """Store a message in the vector database for later retrieval."""
        doc_id = f"{self.agent_id}_{message['timestamp']}"
        self.collection.add(
            ids=[doc_id],
            documents=[message["content"]],
            metadatas=[{
                "role": message["role"],
                "timestamp": message["timestamp"],
                "agent_id": self.agent_id
            }]
        )

    def retrieve_long_term(self, query: str, top_k: int = 5) -> list[str]:
        """Semantic retrieval of relevant past memories."""
        results = self.collection.query(
            query_texts=[query],
            n_results=top_k
        )
        return results["documents"][0] if results["documents"] else []

    def build_context(self, user_input: str, system_prompt: str) -> list[dict]:
        """Assemble the full context window for the LLM call."""
        # Retrieve relevant long-term memories
        retrieved = self.retrieve_long_term(user_input, top_k=5)

        memory_block = ""
        if retrieved:
            memory_block = "\n\nRelevant past context:\n"
            for i, mem in enumerate(retrieved, 1):
                memory_block += f"{i}. {mem}\n"

        messages = [{"role": "system", "content": system_prompt + memory_block}]
        messages.extend(self.working_memory)
        messages.append({"role": "user", "content": user_input})
        return messages

The Agent Loop

With the memory manager in place, the agent loop ties everything together. It calls vLLM, processes the response, and updates memory.

SYSTEM_PROMPT = """You are a helpful autonomous agent.
You have access to long-term memory that provides context from past interactions.
Use the provided past context when relevant, and reason step by step.
When you need to take an action, clearly state your plan."""

class Agent:
    def __init__(self, agent_id: str):
        self.memory = AgentMemory(agent_id)

    def respond(self, user_input: str) -> str:
        # Build context with retrieved memories
        messages = self.memory.build_context(user_input, SYSTEM_PROMPT)

        # Call vLLM
        response = self.memory.client.chat.completions.create(
            model="meta-llama/Meta-Llama-3-8B-Instruct",
            messages=messages,
            temperature=0.7,
            max_tokens=512,
        )

        reply = response.choices[0].message.content

        # Update both working and long-term memory
        self.memory.add_to_working_memory("user", user_input)
        self.memory.add_to_working_memory("assistant", reply)

        return reply

    def reflect(self):
        """Periodic self-reflection: summarize working memory and store insights."""
        if len(self.memory.working_memory) < 4:
            return

        transcript = "\n".join(
            f"{m['role']}: {m['content']}" for m in self.memory.working_memory
        )

        reflection_prompt = [
            {"role": "system", "content": "Summarize key facts and preferences from this conversation that would be useful to remember long-term. Be concise."},
            {"role": "user", "content": transcript}
        ]

        response = self.memory.client.chat.completions.create(
            model="meta-llama/Meta-Llama-3-8B-Instruct",
            messages=reflection_prompt,
            temperature=0.3,
            max_tokens=256,
        )

        summary = response.choices[0].message.content
        self.memory.consolidate_to_long_term({
            "role": "system_reflection",
            "content": f"Reflection: {summary}",
            "timestamp": datetime.now().isoformat()
        })
        print(f"[Reflection stored]: {summary[:100]}...")

# Usage
if __name__ == "__main__":
    agent = Agent("user_001")

    print(agent.respond("My name is Alice and I prefer Python over JavaScript."))
    print(agent.respond("What programming language do I prefer?"))

    agent.reflect()

Advanced: Episodic Memory with Structured Traces

For agents that execute multi-step tasks, episodic memory stores full execution traces. This enables the agent to learn from past successes and failures.

import sqlite3
from dataclasses import dataclass, asdict
from typing import Any

@dataclass
class Episode:
    episode_id: str
    task: str
    steps: list[dict[str, Any]]
    outcome: str  # "success" or "failure"
    lesson: str
    timestamp: str

class EpisodicMemory:
    def __init__(self, db_path: str = "episodes.db"):
        self.conn = sqlite3.connect(db_path)
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS episodes (
                episode_id TEXT PRIMARY KEY,
                task TEXT,
                steps TEXT,
                outcome TEXT,
                lesson TEXT,
                timestamp TEXT
            )
        """)
        self.conn.commit()

    def store(self, episode: Episode):
        self.conn.execute(
            "INSERT OR REPLACE INTO episodes VALUES (?, ?, ?, ?, ?, ?)",
            (episode.episode_id, episode.task, json.dumps(episode.steps),
             episode.outcome, episode.lesson, episode.timestamp)
        )
        self.conn.commit()

    def retrieve_similar(self, task: str, outcome: str = None) -> list[dict]:
        """Retrieve past episodes by keyword match on task description."""
        query = "SELECT * FROM episodes WHERE task LIKE ?"
        params = [f"%{task}%"]
        if outcome:
            query += " AND outcome = ?"
            params.append(outcome)
        rows = self.conn.execute(query, params).fetchall()
        return [
            {"episode_id": r[0], "task": r[1], "steps": json.loads(r[2]),
             "outcome": r[3], "lesson": r[4], "timestamp": r[5]}
            for r in rows
        ]

# Integrating episodic memory into context building
def build_context_with_episodes(self, user_input, system_prompt):
    messages = self.build_context(user_input, system_prompt)

    # Add relevant past episodes
    past_episodes = self.episodic.retrieve_similar(user_input, outcome="success")
    if past_episodes:
        episode_text = "\n\nPast successful approaches:\n"
        for ep in past_episodes[:3]:
            episode_text += f"- Task: {ep['task']}\n  Lesson: {ep['lesson']}\n"
        # Insert after system prompt
        messages[0]["content"] += episode_text

    return messages

Optimizing for vLLM Prefix Caching

To get the most out of vLLM's prefix caching, structure your prompts so that shared content appears at the beginning. The memory block and system prompt should be stable across calls within a session.

# GOOD: Stable prefix, variable suffix — prefix cache hits
def build_context_optimized(self, user_input, system_prompt):
    # Stable prefix: system prompt + retrieved memories (changes infrequently)
    retrieved = self.retrieve_long_term(user_input, top_k=5)
    memory_block = "\n".join(f"- {m}" for m in retrieved)

    stable_prefix = f"{system_prompt}\n\nMemory:\n{memory_block}\n"

    messages = [
        {"role": "system", "content": stable_prefix},
    ]
    # Working memory grows but the prefix stays cacheable
    messages.extend(self.working_memory[-6:])
    messages.append({"role": "user", "content": user_input})
    return messages

# BAD: User input injected into system prompt — breaks prefix caching
def build_context_bad(self, user_input, system_prompt):
    # This makes every request have a unique prefix — no cache hits
    messages = [
        {"role": "system", "content": f"{system_prompt}\nUser said: {user_input}"},
    ]
    return messages

Best Practices

Memory Hygiene

Retrieval Quality

vLLM-Specific Tuning

Security and Privacy

Scaling Considerations

As your agent fleet grows, the memory architecture must scale horizontally. Key strategies include:

import asyncio

class AsyncAgentMemory(AgentMemory):
    async def add_to_working_memory_async(self, role: str, content: str):
        self.add_to_working_memory(role, content)
        # Fire-and-forget long-term storage
        asyncio.create_task(
            asyncio.to_thread(
                self.consolidate_to_long_term,
                {"role": role, "content": content,
                 "timestamp": datetime.now().isoformat()}
            )
        )

Conclusion

Agent memory architectures transform LLMs from stateless text generators into persistent, learning, context-aware agents. By combining short-term working memory, long-term semantic retrieval, and episodic traces, you give agents the ability to remember, reflect, and improve. vLLM's PagedAttention and prefix caching make this practical at scale by keeping inference fast and memory-efficient even with large context windows and high concurrency. The key to success lies in thoughtful memory hygiene — consolidating aggressively, retrieving precisely, and structuring prompts to maximize cache hits. Start with the memory manager and agent loop patterns shown here, then iterate on retrieval quality and consolidation strategies as your agent's tasks grow in complexity.

— Ad —

Google AdSense will appear here after approval

← Back to all articles