mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 07:09:11 +02:00
merge conflict
fix merge conflict multi
This commit is contained in:
zzstoatzz
parent
d6efd774ff
commit
7947d69185
3 changed files with 308 additions and 86 deletions
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -16,4 +16,4 @@ src/fastmcp/_version.py
|
|||
|
||||
# editors
|
||||
.cursorrules
|
||||
.vscode/
|
||||
.vscode/
|
||||
|
|
|
|||
|
|
@ -1,122 +1,344 @@
|
|||
# /// script
|
||||
# dependencies = ["pydantic-ai-slim[openai]", "asyncpg", "numpy", "pgvector", "fastmcp"]
|
||||
# ///
|
||||
|
||||
"""
|
||||
filesystem mcp server with basic memory capabilities.
|
||||
keeps a user profile that can be updated and summarized by an llm.
|
||||
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 datetime import UTC, datetime
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
from typing import Annotated, Self
|
||||
|
||||
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 fastmcp import FastMCP
|
||||
|
||||
MAX_MEMORIES = 3
|
||||
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"
|
||||
|
||||
class Memory(BaseModel):
|
||||
"""a single memory/observation about the user"""
|
||||
|
||||
content: str
|
||||
timestamp: float
|
||||
importance: Annotated[int, Field(ge=1, le=5)] = Field(default=3)
|
||||
|
||||
|
||||
class Profile(BaseModel):
|
||||
"""user profile built from memories"""
|
||||
|
||||
memories: list[Memory] = Field(default_factory=list, max_length=MAX_MEMORIES)
|
||||
summary: str = Field(default="")
|
||||
|
||||
|
||||
class MemoryUpdate(BaseModel):
|
||||
"""llm analysis of how to update the profile"""
|
||||
|
||||
keep_indices: list[int] = Field(description="indices of memories to keep")
|
||||
new_memory: Memory = Field(description="processed version of the new memory")
|
||||
updated_summary: str = Field(description="brief summary of all memories")
|
||||
|
||||
|
||||
memory_agent = Agent(
|
||||
"openai:gpt-4o",
|
||||
result_type=MemoryUpdate,
|
||||
system_prompt="""
|
||||
you help maintain a concise user memory profile. when given a new memory:
|
||||
1. analyze its importance relative to existing memories
|
||||
2. if we're at max capacity of memories, decide which to keep
|
||||
3. provide a brief summary of all memories
|
||||
focus on keeping the most important and relevant information.
|
||||
""",
|
||||
mcp = FastMCP(
|
||||
"memory",
|
||||
dependencies=[
|
||||
"pydantic-ai-slim[openai]",
|
||||
"asyncpg",
|
||||
"numpy",
|
||||
"pgvector",
|
||||
],
|
||||
)
|
||||
|
||||
mcp = FastMCP("memory", dependencies=["pydantic-ai-slim[openai]"])
|
||||
|
||||
DB_DSN = "postgresql://postgres:postgres@localhost:54320/memory_db"
|
||||
# reset memory with rm ~/.fastmcp/{USER}/memory/*
|
||||
PROFILE_DIR = (
|
||||
Path.home() / ".fastmcp" / os.environ.get("USER", "anon") / "memory"
|
||||
).resolve()
|
||||
PROFILE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def remember(
|
||||
content: Annotated[str, Field(description="new observation/memory to store")],
|
||||
importance: Annotated[int, Field(ge=1, le=5, description="importance (1-5)")],
|
||||
) -> str:
|
||||
"""store a new memory/observation about the user"""
|
||||
profile_path = PROFILE_DIR / "profile.json"
|
||||
|
||||
if profile_path.exists():
|
||||
profile = Profile.model_validate_json(profile_path.read_text())
|
||||
else:
|
||||
profile = Profile()
|
||||
|
||||
new_memory = Memory(
|
||||
content=content,
|
||||
timestamp=datetime.now(UTC).timestamp(),
|
||||
importance=importance,
|
||||
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)
|
||||
)
|
||||
|
||||
if len(profile.memories) >= MAX_MEMORIES:
|
||||
result = await memory_agent.run(
|
||||
f"""
|
||||
new memory: {content} (importance: {importance})
|
||||
|
||||
current memories:
|
||||
{[f"{i}: {m.content} (importance: {m.importance})"
|
||||
for i, m in enumerate(profile.memories)]}
|
||||
"""
|
||||
|
||||
async def do_ai[T](
|
||||
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)
|
||||
|
||||
profile.memories = [profile.memories[i] for i in result.data.keep_indices]
|
||||
profile.memories.append(result.data.new_memory)
|
||||
profile.summary = result.data.updated_summary
|
||||
else:
|
||||
profile.memories.append(new_memory)
|
||||
def get_effective_importance(self):
|
||||
return self.importance * (1 + math.log(self.access_count + 1))
|
||||
|
||||
profile_path.write_text(profile.model_dump_json(indent=2))
|
||||
return f"remembered: {content}"
|
||||
|
||||
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:
|
||||
"""read and display the current memory profile"""
|
||||
profile_path = PROFILE_DIR / "profile.json"
|
||||
if not profile_path.exists():
|
||||
return "no profile found"
|
||||
deps = Deps(openai=AsyncOpenAI(), pool=await get_db_pool())
|
||||
profile = await display_memory_tree(deps)
|
||||
await deps.pool.close()
|
||||
return profile
|
||||
|
||||
profile = Profile.model_validate_json(profile_path.read_text())
|
||||
|
||||
output = ["current memories:"]
|
||||
for i, memory in enumerate(profile.memories):
|
||||
output.append(
|
||||
f"{i}. {memory.content} "
|
||||
f"(importance: {memory.importance}, "
|
||||
f"timestamp: {datetime.fromtimestamp(memory.timestamp, UTC)})"
|
||||
)
|
||||
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()
|
||||
|
||||
if profile.summary:
|
||||
output.append(f"\nsummary: {profile.summary}")
|
||||
pool = await asyncpg.create_pool(DB_DSN)
|
||||
try:
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute("CREATE EXTENSION IF NOT EXISTS vector;")
|
||||
|
||||
return "\n".join(output)
|
||||
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())
|
||||
|
|
|
|||
2
uv.lock
generated
2
uv.lock
generated
|
|
@ -228,7 +228,7 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "fastmcp"
|
||||
version = "0.3.6.dev5+g6a13ab9.d20241203"
|
||||
version = "0.3.6.dev8+g3b5ae20"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "httpx" },
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue