+109
-159
tools/blocks.py
+109
-159
tools/blocks.py
···
1
1
"""Block management tools for user-specific memory blocks."""
2
-
import logging
3
-
from typing import List, Type
4
2
from pydantic import BaseModel, Field
5
-
from letta_client.client import BaseTool
6
-
from letta_client import Letta
7
-
8
-
9
-
logger = logging.getLogger(__name__)
3
+
from typing import List, Dict, Any
10
4
11
5
12
6
class AttachUserBlocksArgs(BaseModel):
13
7
handles: List[str] = Field(..., description="List of user Bluesky handles (e.g., ['user1.bsky.social', 'user2.bsky.social'])")
14
8
15
9
16
-
class AttachUserBlocksTool(BaseTool):
17
-
name: str = "attach_user_blocks"
18
-
args_schema: Type[BaseModel] = AttachUserBlocksArgs
19
-
description: str = "Attach user-specific memory blocks to the agent. Creates blocks if they don't exist."
20
-
tags: List[str] = ["memory", "blocks", "user"]
21
-
22
-
def run(self, handles: List[str], agent_state: "AgentState") -> str:
23
-
"""Attach user-specific memory blocks."""
24
-
import os
25
-
from letta_client import Letta
26
-
27
-
try:
28
-
client = Letta(token=os.environ["LETTA_API_KEY"])
29
-
results = []
30
-
31
-
# Get current blocks
32
-
current_blocks = agent_state.block_ids
33
-
current_block_labels = set()
34
-
for block_id in current_blocks:
35
-
block = client.blocks.get(block_id)
36
-
current_block_labels.add(block.label)
37
-
38
-
for handle in handles:
39
-
# Sanitize handle for block label - completely self-contained
40
-
clean_handle = handle.lstrip('@').replace('.', '_').replace('-', '_').replace(' ', '_')
41
-
block_label = f"user_{clean_handle}"
42
-
43
-
# Skip if already attached
44
-
if block_label in current_block_labels:
45
-
results.append(f"✓ {handle}: Already attached")
46
-
continue
47
-
48
-
# Check if block exists or create new one
49
-
try:
50
-
blocks = client.blocks.list(label=block_label)
51
-
if blocks and len(blocks) > 0:
52
-
block = blocks[0]
53
-
logger.info(f"Found existing block: {block_label}")
54
-
else:
55
-
block = client.blocks.create(
56
-
label=block_label,
57
-
value=f"# User: {handle}\n\nNo information about this user yet.",
58
-
limit=5000
59
-
)
60
-
logger.info(f"Created new block: {block_label}")
61
-
62
-
# Attach block individually to avoid race conditions
63
-
client.agents.blocks.attach(
64
-
agent_id=str(agent_state.id),
65
-
block_id=str(block.id),
66
-
)
67
-
results.append(f"✓ {handle}: Block attached")
68
-
69
-
except Exception as e:
70
-
results.append(f"✗ {handle}: Error - {str(e)}")
71
-
logger.error(f"Error processing block for {handle}: {e}")
72
-
73
-
return f"Attachment results:\n" + "\n".join(results)
74
-
75
-
except Exception as e:
76
-
logger.error(f"Error attaching user blocks: {e}")
77
-
raise e
78
-
79
-
80
10
class DetachUserBlocksArgs(BaseModel):
81
11
handles: List[str] = Field(..., description="List of user Bluesky handles (e.g., ['user1.bsky.social', 'user2.bsky.social'])")
82
12
83
13
84
-
class DetachUserBlocksTool(BaseTool):
85
-
name: str = "detach_user_blocks"
86
-
args_schema: Type[BaseModel] = DetachUserBlocksArgs
87
-
description: str = "Detach user-specific memory blocks from the agent. Blocks are preserved for later use."
88
-
tags: List[str] = ["memory", "blocks", "user"]
89
14
90
-
def run(self, handles: List[str], agent_state: "AgentState") -> str:
91
-
"""Detach user-specific memory blocks."""
92
-
import os
93
-
from letta_client import Letta
15
+
def attach_user_blocks(handles: list, agent_state: "AgentState") -> str:
16
+
"""
17
+
Attach user-specific memory blocks to the agent. Creates blocks if they don't exist.
18
+
19
+
Args:
20
+
handles: List of user Bluesky handles (e.g., ['user1.bsky.social', 'user2.bsky.social'])
21
+
agent_state: The agent state object containing agent information
22
+
23
+
Returns:
24
+
String with attachment results for each handle
25
+
"""
26
+
import os
27
+
import logging
28
+
from letta_client import Letta
29
+
30
+
logger = logging.getLogger(__name__)
31
+
32
+
try:
33
+
client = Letta(token=os.environ["LETTA_API_KEY"])
34
+
results = []
94
35
95
-
try:
96
-
client = Letta(token=os.environ["LETTA_API_KEY"])
97
-
results = []
98
-
blocks_to_remove = set()
36
+
# Get current blocks using the API
37
+
current_blocks = client.agents.blocks.list(agent_id=str(agent_state.id))
38
+
current_block_labels = set()
39
+
40
+
for block in current_blocks:
41
+
current_block_labels.add(block.label)
99
42
100
-
# Build mapping of block labels to IDs
101
-
current_blocks = agent_state.block_ids
102
-
block_label_to_id = {}
43
+
for handle in handles:
44
+
# Sanitize handle for block label - completely self-contained
45
+
clean_handle = handle.lstrip('@').replace('.', '_').replace('-', '_').replace(' ', '_')
46
+
block_label = f"user_{clean_handle}"
103
47
104
-
for block_id in current_blocks:
105
-
block = client.blocks.get(block_id)
106
-
block_label_to_id[block.label] = block_id
107
-
108
-
# Process each handle
109
-
for handle in handles:
110
-
# Sanitize handle for block label - completely self-contained
111
-
clean_handle = handle.lstrip('@').replace('.', '_').replace('-', '_').replace(' ', '_')
112
-
block_label = f"user_{clean_handle}"
48
+
# Skip if already attached
49
+
if block_label in current_block_labels:
50
+
results.append(f"✓ {handle}: Already attached")
51
+
continue
113
52
114
-
if block_label in block_label_to_id:
115
-
blocks_to_remove.add(block_label_to_id[block_label])
116
-
results.append(f"✓ {handle}: Detached")
53
+
# Check if block exists or create new one
54
+
try:
55
+
blocks = client.blocks.list(label=block_label)
56
+
if blocks and len(blocks) > 0:
57
+
block = blocks[0]
58
+
logger.info(f"Found existing block: {block_label}")
117
59
else:
118
-
results.append(f"✗ {handle}: Not attached")
60
+
block = client.blocks.create(
61
+
label=block_label,
62
+
value=f"# User: {handle}\n\nNo information about this user yet.",
63
+
limit=5000
64
+
)
65
+
logger.info(f"Created new block: {block_label}")
119
66
120
-
# Remove blocks from agent one by one
121
-
for block_id in blocks_to_remove:
122
-
client.agents.blocks.detach(
67
+
# Attach block atomically
68
+
client.agents.blocks.attach(
123
69
agent_id=str(agent_state.id),
124
-
block_id=block_id
70
+
block_id=str(block.id)
125
71
)
72
+
results.append(f"✓ {handle}: Block attached")
73
+
logger.info(f"Successfully attached block {block_label} to agent")
126
74
127
-
return f"Detachment results:\n" + "\n".join(results)
75
+
except Exception as e:
76
+
results.append(f"✗ {handle}: Error - {str(e)}")
77
+
logger.error(f"Error processing block for {handle}: {e}")
128
78
129
-
except Exception as e:
130
-
logger.error(f"Error detaching user blocks: {e}")
131
-
return f"Error detaching user blocks: {str(e)}"
79
+
return f"Attachment results:\n" + "\n".join(results)
80
+
81
+
except Exception as e:
82
+
logger.error(f"Error attaching user blocks: {e}")
83
+
raise Exception(f"Error attaching user blocks: {str(e)}")
132
84
133
85
134
-
# class UserBlockUpdate(BaseModel):
135
-
# handle: str = Field(..., description="User's Bluesky handle (e.g., 'user.bsky.social')")
136
-
# content: str = Field(..., description="New content for the user's memory block")
137
-
138
-
139
-
# class UpdateUserBlockArgs(BaseModel):
140
-
# updates: List[UserBlockUpdate] = Field(..., description="List of user block updates")
141
-
142
-
143
-
# class UpdateUserBlockTool(BaseTool):
144
-
# name: str = "update_user_blocks"
145
-
# args_schema: Type[BaseModel] = UpdateUserBlockArgs
146
-
# description: str = "Update the content of user-specific memory blocks"
147
-
# tags: List[str] = ["memory", "blocks", "user"]
148
-
149
-
# def run(self, updates: List[UserBlockUpdate]) -> str:
150
-
# """Update user-specific memory blocks."""
151
-
# import os
152
-
# from letta_client import Letta
153
-
154
-
# try:
155
-
# client = Letta(token=os.environ["LETTA_API_KEY"])
156
-
# results = []
157
-
158
-
# for update in updates:
159
-
# handle = update.handle
160
-
# new_content = update.content
161
-
# # Sanitize handle for block label - completely self-contained
162
-
# clean_handle = handle.lstrip('@').replace('.', '_').replace('-', '_').replace(' ', '_')
163
-
# block_label = f"user_{clean_handle}"
86
+
def detach_user_blocks(handles: list, agent_state: "AgentState") -> str:
87
+
"""
88
+
Detach user-specific memory blocks from the agent. Blocks are preserved for later use.
89
+
90
+
Args:
91
+
handles: List of user Bluesky handles (e.g., ['user1.bsky.social', 'user2.bsky.social'])
92
+
agent_state: The agent state object containing agent information
93
+
94
+
Returns:
95
+
String with detachment results for each handle
96
+
"""
97
+
import os
98
+
import logging
99
+
from letta_client import Letta
100
+
101
+
logger = logging.getLogger(__name__)
102
+
103
+
try:
104
+
client = Letta(token=os.environ["LETTA_API_KEY"])
105
+
results = []
164
106
165
-
# try:
166
-
# # Find the block
167
-
# blocks = client.blocks.list(label=block_label)
168
-
# if not blocks or len(blocks) == 0:
169
-
# results.append(f"✗ {handle}: Block not found - use attach_user_blocks first")
170
-
# continue
107
+
# Build mapping of block labels to IDs using the API
108
+
current_blocks = client.agents.blocks.list(agent_id=str(agent_state.id))
109
+
block_label_to_id = {}
171
110
172
-
# block = blocks[0]
111
+
for block in current_blocks:
112
+
block_label_to_id[block.label] = str(block.id)
173
113
174
-
# # Update block content
175
-
# updated_block = client.blocks.modify(
176
-
# block_id=str(block.id),
177
-
# value=new_content
178
-
# )
114
+
# Process each handle and detach atomically
115
+
for handle in handles:
116
+
# Sanitize handle for block label - completely self-contained
117
+
clean_handle = handle.lstrip('@').replace('.', '_').replace('-', '_').replace(' ', '_')
118
+
block_label = f"user_{clean_handle}"
179
119
180
-
# preview = new_content[:100] + "..." if len(new_content) > 100 else new_content
181
-
# results.append(f"✓ {handle}: Updated - {preview}")
120
+
if block_label in block_label_to_id:
121
+
try:
122
+
# Detach block atomically
123
+
client.agents.blocks.detach(
124
+
agent_id=str(agent_state.id),
125
+
block_id=block_label_to_id[block_label]
126
+
)
127
+
results.append(f"✓ {handle}: Detached")
128
+
logger.info(f"Successfully detached block {block_label} from agent")
129
+
except Exception as e:
130
+
results.append(f"✗ {handle}: Error during detachment - {str(e)}")
131
+
logger.error(f"Error detaching block {block_label}: {e}")
132
+
else:
133
+
results.append(f"✗ {handle}: Not attached")
182
134
183
-
# except Exception as e:
184
-
# results.append(f"✗ {handle}: Error - {str(e)}")
185
-
# logger.error(f"Error updating block for {handle}: {e}")
135
+
return f"Detachment results:\n" + "\n".join(results)
136
+
137
+
except Exception as e:
138
+
logger.error(f"Error detaching user blocks: {e}")
139
+
raise Exception(f"Error detaching user blocks: {str(e)}")
186
140
187
-
# return f"Update results:\n" + "\n".join(results)
188
141
189
-
# except Exception as e:
190
-
# logger.error(f"Error updating user blocks: {e}")
191
-
# return f"Error updating user blocks: {str(e)}"