python-sdk/examples/mcpserver/memory.py at main · modelcontextprotocol/python-sdk · GitHub

History

324 lines (270 loc) · 10.2 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

# /// script

# dependencies = ["pydantic-ai-slim[openai]", "asyncpg", "numpy", "pgvector"]

# ///

# uv pip install 'pydantic-ai-slim[openai]' asyncpg numpy pgvector

"""Recursive memory system inspired by the human brain's clustering of memories.

Uses OpenAI's 'text-embedding-3-small' model and pgvector for efficient

similarity search.

"""

import asyncio

import math

import os

from dataclasses import dataclass

from datetime import datetime, timezone

from pathlib import Path

from typing import Annotated, Self, TypeVar

import asyncpg

import numpy as np

from openai import AsyncOpenAI

from pgvector.asyncpg import register_vector # Import register_vector

from pydantic import BaseModel, Field

from pydantic_ai import Agent

from mcp.server.mcpserver import MCPServer

MAX_DEPTH = 5

SIMILARITY_THRESHOLD = 0.7

DECAY_FACTOR = 0.99

REINFORCEMENT_FACTOR = 1.1

DEFAULT_LLM_MODEL = "openai:gpt-4o"

DEFAULT_EMBEDDING_MODEL = "text-embedding-3-small"

T = TypeVar("T")

mcp = MCPServer("memory")

DB_DSN = "postgresql://postgres:postgres@localhost:54320/memory_db"

# reset memory with rm ~/.mcp/{USER}/memory/*

PROFILE_DIR = (Path.home() / ".mcp" / os.environ.get("USER", "anon") / "memory").resolve()

PROFILE_DIR.mkdir(parents=True, exist_ok=True)

def cosine_similarity(a: list[float], b: list[float]) -> float:

a_array = np.array(a, dtype=np.float64)

b_array = np.array(b, dtype=np.float64)

return np.dot(a_array, b_array) / (np.linalg.norm(a_array) * np.linalg.norm(b_array))

async def do_ai(

user_prompt: str,

system_prompt: str,

result_type: type[T] | Annotated,

deps=None,

) -> T:

agent = Agent(

DEFAULT_LLM_MODEL,

system_prompt=system_prompt,

result_type=result_type,

)

result = await agent.run(user_prompt, deps=deps)

return result.data

@dataclass

class Deps:

openai: AsyncOpenAI

pool: asyncpg.Pool

async def get_db_pool() -> asyncpg.Pool:

async def init(conn):

await conn.execute("CREATE EXTENSION IF NOT EXISTS vector;")

await register_vector(conn)

pool = await asyncpg.create_pool(DB_DSN, init=init)

return pool

class MemoryNode(BaseModel):

id: int | None = None

content: str

summary: str = ""

importance: float = 1.0

access_count: int = 0

timestamp: float = Field(default_factory=lambda: datetime.now(timezone.utc).timestamp())

embedding: list[float]

@classmethod

async def from_content(cls, content: str, deps: Deps):

embedding = await get_embedding(content, deps)

return cls(content=content, embedding=embedding)

async def save(self, deps: Deps):

async with deps.pool.acquire() as conn:

if self.id is None:

result = await conn.fetchrow(

"""

INSERT INTO memories (content, summary, importance, access_count,

timestamp, embedding)

VALUES ($1, $2, $3, $4, $5, $6)

RETURNING id

""",

self.content,

self.summary,

self.importance,

self.access_count,

self.timestamp,

self.embedding,

)

self.id = result["id"]

else:

await conn.execute(

"""

UPDATE memories

SET content = $1, summary = $2, importance = $3,

access_count = $4, timestamp = $5, embedding = $6

WHERE id = $7

""",

self.content,

self.summary,

self.importance,

self.access_count,

self.timestamp,

self.embedding,

self.id,

)

async def merge_with(self, other: Self, deps: Deps):

self.content = await do_ai(

f"{self.content}\n\n{other.content}",

"Combine the following two texts into a single, coherent text.",

str,

deps,

)

self.importance += other.importance

self.access_count += other.access_count

self.embedding = [(a + b) / 2 for a, b in zip(self.embedding, other.embedding)]

self.summary = await do_ai(self.content, "Summarize the following text concisely.", str, deps)

await self.save(deps)

# Delete the merged node from the database

if other.id is not None:

await delete_memory(other.id, deps)

def get_effective_importance(self):

return self.importance * (1 + math.log(self.access_count + 1))

async def get_embedding(text: str, deps: Deps) -> list[float]:

embedding_response = await deps.openai.embeddings.create(

input=text,

model=DEFAULT_EMBEDDING_MODEL,

)

return embedding_response.data[0].embedding

async def delete_memory(memory_id: int, deps: Deps):

async with deps.pool.acquire() as conn:

await conn.execute("DELETE FROM memories WHERE id = $1", memory_id)

async def add_memory(content: str, deps: Deps):

new_memory = await MemoryNode.from_content(content, deps)

await new_memory.save(deps)

similar_memories = await find_similar_memories(new_memory.embedding, deps)

for memory in similar_memories:

if memory.id != new_memory.id:

await new_memory.merge_with(memory, deps)

await update_importance(new_memory.embedding, deps)

await prune_memories(deps)

return f"Remembered: {content}"

async def find_similar_memories(embedding: list[float], deps: Deps) -> list[MemoryNode]:

async with deps.pool.acquire() as conn:

rows = await conn.fetch(

"""

SELECT id, content, summary, importance, access_count, timestamp, embedding

FROM memories

ORDER BY embedding <-> $1

LIMIT 5

""",

embedding,

)

memories = [

MemoryNode(

id=row["id"],

content=row["content"],

summary=row["summary"],

importance=row["importance"],

access_count=row["access_count"],

timestamp=row["timestamp"],

embedding=row["embedding"],

)

for row in rows

]

return memories

async def update_importance(user_embedding: list[float], deps: Deps):

async with deps.pool.acquire() as conn:

rows = await conn.fetch("SELECT id, importance, access_count, embedding FROM memories")

for row in rows:

memory_embedding = row["embedding"]

similarity = cosine_similarity(user_embedding, memory_embedding)

if similarity > SIMILARITY_THRESHOLD:

new_importance = row["importance"] * REINFORCEMENT_FACTOR

new_access_count = row["access_count"] + 1

else:

new_importance = row["importance"] * DECAY_FACTOR

new_access_count = row["access_count"]

await conn.execute(

"""

UPDATE memories

SET importance = $1, access_count = $2

WHERE id = $3

""",

new_importance,

new_access_count,

row["id"],

)

async def prune_memories(deps: Deps):

async with deps.pool.acquire() as conn:

rows = await conn.fetch(

"""

SELECT id, importance, access_count

FROM memories

ORDER BY importance DESC

OFFSET $1

""",

MAX_DEPTH,

)

for row in rows:

await conn.execute("DELETE FROM memories WHERE id = $1", row["id"])

async def display_memory_tree(deps: Deps) -> str:

async with deps.pool.acquire() as conn:

rows = await conn.fetch(

"""

SELECT content, summary, importance, access_count

FROM memories

ORDER BY importance DESC

LIMIT $1

""",

MAX_DEPTH,

)

result = ""

for row in rows:

effective_importance = row["importance"] * (1 + math.log(row["access_count"] + 1))

summary = row["summary"] or row["content"]

result += f"- {summary} (Importance: {effective_importance:.2f})\n"

return result

@mcp.tool()

async def remember(

contents: list[str] = Field(description="List of observations or memories to store"),

deps = Deps(openai=AsyncOpenAI(), pool=await get_db_pool())

try:

return "\n".join(await asyncio.gather(*[add_memory(content, deps) for content in contents]))

finally:

await deps.pool.close()

@mcp.tool()

async def read_profile() -> str:

deps = Deps(openai=AsyncOpenAI(), pool=await get_db_pool())

profile = await display_memory_tree(deps)

await deps.pool.close()

return profile

async def initialize_database():

pool = await asyncpg.create_pool("postgresql://postgres:postgres@localhost:54320/postgres")

try:

async with pool.acquire() as conn:

await conn.execute("""

SELECT pg_terminate_backend(pg_stat_activity.pid)

FROM pg_stat_activity

WHERE pg_stat_activity.datname = 'memory_db'

AND pid <> pg_backend_pid();

""")

await conn.execute("DROP DATABASE IF EXISTS memory_db;")

await conn.execute("CREATE DATABASE memory_db;")

finally:

await pool.close()

pool = await asyncpg.create_pool(DB_DSN)

try:

async with pool.acquire() as conn:

await conn.execute("CREATE EXTENSION IF NOT EXISTS vector;")

await register_vector(conn)

await conn.execute("""

CREATE TABLE IF NOT EXISTS memories (

id SERIAL PRIMARY KEY,

content TEXT NOT NULL,

summary TEXT,

importance REAL NOT NULL,

access_count INT NOT NULL,

timestamp DOUBLE PRECISION NOT NULL,

embedding vector(1536) NOT NULL

);

CREATE INDEX IF NOT EXISTS idx_memories_embedding ON memories

USING hnsw (embedding vector_l2_ops);

""")

finally:

await pool.close()

if __name__ == "__main__":

asyncio.run(initialize_database())

python-sdk/examples/mcpserver/memory.py at main · modelcontextprotocol/python-sdk · GitHub

Uh oh!

Expand file tree

Breadcrumbs

memory.py

Latest commit

History

Breadcrumbs

memory.py

File metadata and controls