25d54dd896
* auto fixes * auto fix pt2 and transitive deps and undefined var checking locals() * manual fixes (ignored or letta-code fixed) * fix circular import
565 lines
22 KiB
Python
565 lines
22 KiB
Python
"""
|
|
Shared helper functions for Anthropic-compatible proxy endpoints.
|
|
|
|
These helpers are used by both the Anthropic and Z.ai proxy routers to reduce code duplication.
|
|
"""
|
|
|
|
import json
|
|
|
|
from fastapi import Request
|
|
|
|
from letta.log import get_logger
|
|
from letta.server.rest_api.utils import capture_and_persist_messages
|
|
from letta.settings import model_settings
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def strip_policy_specs(text: str) -> str:
|
|
"""
|
|
Remove Claude policy injection blocks from message text.
|
|
|
|
Claude injects policy instructions in two forms:
|
|
1. Appended with prefix: 'user: <policy_spec>...'
|
|
2. As entire message: '<policy_spec>...'
|
|
|
|
We truncate everything from the policy start marker onwards since it's all injected policy content.
|
|
"""
|
|
# Check if entire message is a policy spec (starts with tag)
|
|
if text.startswith("<policy_spec>"):
|
|
logger.info("[Proxy Helpers] Stripped policy injection (entire message)")
|
|
return ""
|
|
|
|
# Check if policy spec is appended (with prefix)
|
|
policy_start = text.find("user: <policy_spec>")
|
|
if policy_start != -1:
|
|
logger.info(f"[Proxy Helpers] Stripped policy injection from position {policy_start}")
|
|
# Truncate everything from this point onwards
|
|
cleaned = text[:policy_start].strip()
|
|
return cleaned
|
|
|
|
# No policy injection found, return original text
|
|
return text
|
|
|
|
|
|
def extract_user_messages(body: bytes) -> list[str]:
|
|
"""Extract user messages from request body."""
|
|
messages = []
|
|
try:
|
|
request_data = json.loads(body)
|
|
messages = request_data.get("messages", [])
|
|
|
|
user_messages = []
|
|
for msg in messages:
|
|
if msg.get("role") == "user":
|
|
content = msg.get("content", "")
|
|
if isinstance(content, str):
|
|
# Strip policy specs before adding
|
|
cleaned = strip_policy_specs(content)
|
|
if cleaned: # Only add if not empty after stripping
|
|
user_messages.append(cleaned)
|
|
elif isinstance(content, list):
|
|
for block in content:
|
|
if isinstance(block, dict):
|
|
if block.get("type") == "text":
|
|
text = block.get("text", "")
|
|
# Strip policy specs from text blocks
|
|
cleaned = strip_policy_specs(text)
|
|
if cleaned: # Only add if not empty after stripping
|
|
user_messages.append(cleaned)
|
|
elif block.get("type") == "image":
|
|
user_messages.append("[IMAGE]")
|
|
|
|
return user_messages
|
|
except Exception as e:
|
|
logger.warning(f"[Proxy Helpers] Failed to extract user messages from request {messages}: {e}")
|
|
return []
|
|
|
|
|
|
def extract_assistant_message(response_data: dict) -> str:
|
|
"""Extract assistant message from response data."""
|
|
content_blocks = []
|
|
try:
|
|
content_blocks = response_data.get("content", [])
|
|
text_parts = []
|
|
|
|
for block in content_blocks:
|
|
if isinstance(block, dict) and block.get("type") == "text":
|
|
text_parts.append(block.get("text", ""))
|
|
|
|
return "\n".join(text_parts)
|
|
except Exception as e:
|
|
logger.warning(f"[Proxy Helpers] Failed to extract assistant message from response {content_blocks}: {e}")
|
|
return ""
|
|
|
|
|
|
def is_topic_detection_response(message: str) -> bool:
|
|
"""
|
|
Check if the assistant message is a topic detection response (contains isNewTopic key).
|
|
These are Claude Code metadata responses that should not be persisted as conversation.
|
|
"""
|
|
try:
|
|
stripped = message.strip()
|
|
if stripped.startswith("{") and stripped.endswith("}"):
|
|
parsed = json.loads(stripped)
|
|
# Check for isNewTopic key which indicates topic detection
|
|
if "isNewTopic" in parsed:
|
|
return True
|
|
except (json.JSONDecodeError, AttributeError):
|
|
pass
|
|
return False
|
|
|
|
|
|
def prepare_headers(request: Request, proxy_name: str, use_bearer_auth: bool = False) -> dict | None:
|
|
"""
|
|
Prepare headers for forwarding to Anthropic-compatible API.
|
|
|
|
Args:
|
|
request: The incoming FastAPI request
|
|
proxy_name: Name of the proxy for logging (e.g., "Anthropic Proxy", "Z.ai Proxy")
|
|
use_bearer_auth: If True, convert x-api-key to Bearer token in Authorization header (for Z.ai)
|
|
|
|
Returns:
|
|
Dictionary of headers to forward, or None if authentication fails
|
|
"""
|
|
skip_headers = {
|
|
"host",
|
|
"connection",
|
|
"content-length",
|
|
"transfer-encoding",
|
|
"content-encoding",
|
|
"te",
|
|
"upgrade",
|
|
"proxy-authenticate",
|
|
"proxy-authorization",
|
|
"authorization",
|
|
}
|
|
|
|
headers = {}
|
|
for key, value in request.headers.items():
|
|
if key.lower() not in skip_headers:
|
|
headers[key] = value
|
|
|
|
# Extract API key from headers or fallback to letta's key
|
|
api_key = None
|
|
if "x-api-key" in headers:
|
|
api_key = headers["x-api-key"]
|
|
elif "anthropic-api-key" in headers:
|
|
api_key = headers["anthropic-api-key"]
|
|
else:
|
|
# Fallback to letta's anthropic api key if not provided
|
|
api_key = model_settings.anthropic_api_key
|
|
if api_key:
|
|
logger.info(f"[{proxy_name}] Falling back to Letta's anthropic api key instead of user's key")
|
|
|
|
# Handle authentication based on proxy type
|
|
if use_bearer_auth:
|
|
# Z.ai: use Bearer token in Authorization header
|
|
if api_key:
|
|
headers["authorization"] = f"Bearer {api_key}"
|
|
# Keep x-api-key in headers too (doesn't hurt)
|
|
if "x-api-key" not in headers and api_key:
|
|
headers["x-api-key"] = api_key
|
|
else:
|
|
# Anthropic: use x-api-key header
|
|
if api_key and "x-api-key" not in headers:
|
|
headers["x-api-key"] = api_key
|
|
|
|
if "content-type" not in headers:
|
|
headers["content-type"] = "application/json"
|
|
|
|
return headers
|
|
|
|
|
|
def format_memory_blocks(blocks, agent_id: str) -> str:
|
|
"""Format memory blocks for injection into system prompt."""
|
|
blocks_with_content = [block for block in blocks if block.value]
|
|
|
|
if not blocks_with_content:
|
|
return ""
|
|
|
|
memory_context = (
|
|
"<letta>\n"
|
|
"You have persistent memory powered by Letta that is maintained across conversations. "
|
|
"A background agent updates these memory blocks based on conversation content.\n"
|
|
"<memory_blocks>\n"
|
|
"The following memory blocks are currently engaged in your core memory unit:\n\n"
|
|
)
|
|
|
|
for idx, block in enumerate(blocks_with_content):
|
|
label = block.label or "block"
|
|
value = block.value or ""
|
|
desc = block.description or ""
|
|
chars_current = len(value)
|
|
limit = block.limit if block.limit is not None else 0
|
|
|
|
memory_context += f"<{label}>\n"
|
|
if desc:
|
|
memory_context += "<description>\n"
|
|
memory_context += f"{desc}\n"
|
|
memory_context += "</description>\n"
|
|
memory_context += "<metadata>\n"
|
|
memory_context += f"- chars_current={chars_current}\n"
|
|
memory_context += f"- chars_limit={limit}\n"
|
|
memory_context += "</metadata>\n"
|
|
memory_context += "<value>\n"
|
|
memory_context += f"{value}\n"
|
|
memory_context += "</value>\n"
|
|
memory_context += f"</{label}>\n"
|
|
|
|
if idx != len(blocks_with_content) - 1:
|
|
memory_context += "\n"
|
|
|
|
memory_context += "\n</memory_blocks>\n\n"
|
|
memory_context += (
|
|
"<memory_management>\n"
|
|
f"Users can view and edit their memory blocks at:\n"
|
|
f"https://app.letta.com/agents/{agent_id}\n\n"
|
|
"Share this link when users ask how to manage their memory, what you remember about them, or how to view, edit, or delete stored information.\n"
|
|
"</memory_management>\n\n"
|
|
"<documentation>\n"
|
|
"- Memory blocks: https://docs.letta.com/guides/agents/memory-blocks/index.md\n"
|
|
"- Full Letta documentation: https://docs.letta.com/llms.txt\n\n"
|
|
"Reference these when users ask how Letta memory works or want to learn more about the platform.\n"
|
|
"</documentation>\n"
|
|
"</letta>"
|
|
)
|
|
return memory_context
|
|
|
|
|
|
def build_response_from_chunks(chunks: list[bytes]) -> str:
|
|
"""Build complete response text from streaming chunks."""
|
|
try:
|
|
text_parts = []
|
|
full_data = b"".join(chunks).decode("utf-8")
|
|
|
|
for line in full_data.split("\n"):
|
|
if line.startswith("data: "):
|
|
data_str = line[6:] # Remove "data: " prefix
|
|
|
|
if data_str.strip() in ["[DONE]", ""]:
|
|
continue
|
|
|
|
try:
|
|
event_data = json.loads(data_str)
|
|
event_type = event_data.get("type")
|
|
|
|
if event_type == "content_block_delta":
|
|
delta = event_data.get("delta", {})
|
|
if delta.get("type") == "text_delta":
|
|
text_parts.append(delta.get("text", ""))
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
return "".join(text_parts)
|
|
except Exception as e:
|
|
logger.warning(f"[Proxy Helpers] Failed to build response from chunks: {e}")
|
|
return ""
|
|
|
|
|
|
async def inject_memory_context(
|
|
server,
|
|
agent,
|
|
actor,
|
|
request_data: dict,
|
|
proxy_name: str,
|
|
) -> dict:
|
|
"""
|
|
Inject memory context into the request system prompt.
|
|
|
|
Args:
|
|
server: SyncServer instance
|
|
agent: Agent to get memory from
|
|
actor: Actor performing the operation
|
|
request_data: Request data dictionary to modify
|
|
proxy_name: Name of the proxy for logging (e.g., "Anthropic Proxy", "Z.ai Proxy")
|
|
|
|
Returns:
|
|
Modified request data with memory context injected
|
|
"""
|
|
try:
|
|
messages = request_data.get("messages", [])
|
|
if not messages:
|
|
return request_data
|
|
|
|
memory_context = format_memory_blocks(agent.blocks, agent.id)
|
|
|
|
if not memory_context:
|
|
logger.debug(f"[{proxy_name}] No memory blocks found, skipping memory injection")
|
|
return request_data
|
|
|
|
block_count = len([b for b in agent.blocks if b.value])
|
|
logger.info(f"[{proxy_name}] Injecting {block_count} memory block(s) into request")
|
|
|
|
# Inject into system prompt
|
|
modified_data = request_data.copy()
|
|
|
|
# Check if there's already a system prompt
|
|
# Anthropic API accepts system as either a string or list of content blocks
|
|
existing_system = modified_data.get("system", "")
|
|
|
|
# Handle both string and list system prompts
|
|
if isinstance(existing_system, list):
|
|
# If it's a list, prepend our context as a text block
|
|
modified_data["system"] = existing_system + [{"type": "text", "text": memory_context.rstrip()}]
|
|
elif existing_system:
|
|
# If it's a non-empty string, prepend our context
|
|
modified_data["system"] = memory_context + existing_system
|
|
else:
|
|
# No existing system prompt
|
|
modified_data["system"] = memory_context.rstrip()
|
|
|
|
# Fix max_tokens if using extended thinking
|
|
# Anthropic requires max_tokens > thinking.budget_tokens
|
|
if "thinking" in modified_data and isinstance(modified_data["thinking"], dict):
|
|
budget_tokens = modified_data["thinking"].get("budget_tokens", 0)
|
|
current_max_tokens = modified_data.get("max_tokens", 0)
|
|
|
|
if budget_tokens > 0 and current_max_tokens <= budget_tokens:
|
|
# Set max_tokens to budget_tokens + reasonable buffer for response
|
|
# Claude Code typically uses budget_tokens around 10000-20000
|
|
modified_data["max_tokens"] = budget_tokens + 4096
|
|
logger.info(
|
|
f"[{proxy_name}] Adjusted max_tokens from {current_max_tokens} to {modified_data['max_tokens']} (thinking.budget_tokens={budget_tokens})"
|
|
)
|
|
|
|
return modified_data
|
|
|
|
except Exception as e:
|
|
logger.exception(f"[{proxy_name}] Failed to inject memory context: {e}")
|
|
return request_data
|
|
|
|
|
|
async def persist_messages_background(
|
|
server,
|
|
agent,
|
|
actor,
|
|
user_messages: list[str],
|
|
assistant_message: str,
|
|
model_name: str,
|
|
proxy_name: str,
|
|
):
|
|
"""
|
|
Background task to persist messages without blocking the response.
|
|
|
|
This runs asynchronously after the response is returned to minimize latency.
|
|
|
|
Args:
|
|
server: SyncServer instance
|
|
agent: Agent to persist messages for
|
|
actor: Actor performing the operation
|
|
user_messages: List of user messages to persist
|
|
assistant_message: Assistant message to persist
|
|
model_name: Model name for the messages
|
|
proxy_name: Name of the proxy for logging (e.g., "Anthropic Proxy", "Z.ai Proxy")
|
|
"""
|
|
try:
|
|
result = await capture_and_persist_messages(
|
|
server=server,
|
|
agent=agent,
|
|
actor=actor,
|
|
user_messages=user_messages,
|
|
assistant_message=assistant_message,
|
|
model=model_name,
|
|
)
|
|
if result.get("success"):
|
|
logger.info(f"[{proxy_name}] Persisted messages: {result['messages_created']} messages saved")
|
|
else:
|
|
logger.debug(f"[{proxy_name}] Skipped persistence: {result.get('reason', 'unknown')}")
|
|
except Exception as e:
|
|
logger.error(f"[{proxy_name}] Failed to persist messages in background: {e}")
|
|
|
|
|
|
async def check_for_duplicate_message(server, agent, actor, user_messages: list[str], proxy_name: str) -> list[str]:
|
|
"""
|
|
Check if the last user message is a duplicate of the most recent persisted message.
|
|
|
|
Returns a filtered list with duplicates removed to prevent race conditions.
|
|
|
|
Args:
|
|
server: SyncServer instance
|
|
agent: Agent to check messages for
|
|
actor: Actor performing the operation
|
|
user_messages: List of user messages to check
|
|
proxy_name: Name of the proxy for logging
|
|
|
|
Returns:
|
|
Filtered list of user messages (empty if duplicate detected)
|
|
"""
|
|
user_messages_to_persist = user_messages.copy() if user_messages else []
|
|
if user_messages_to_persist:
|
|
try:
|
|
from letta.schemas.enums import MessageRole
|
|
|
|
recent_messages = await server.message_manager.list_messages(
|
|
agent_id=agent.id,
|
|
actor=actor,
|
|
limit=5,
|
|
roles=[MessageRole.user],
|
|
ascending=False,
|
|
)
|
|
if recent_messages:
|
|
last_user_msg = recent_messages[0]
|
|
last_message_text = ""
|
|
if last_user_msg.content:
|
|
for content_block in last_user_msg.content:
|
|
if hasattr(content_block, "text"):
|
|
last_message_text += content_block.text
|
|
|
|
incoming_msg = user_messages_to_persist[-1]
|
|
if last_message_text and last_message_text == incoming_msg:
|
|
logger.info(f"[{proxy_name}] Skipping duplicate user message: {incoming_msg[:100]}...")
|
|
user_messages_to_persist = []
|
|
except Exception as e:
|
|
logger.warning(f"[{proxy_name}] Failed to check for duplicate messages: {e}")
|
|
|
|
return user_messages_to_persist
|
|
|
|
|
|
async def backfill_agent_project_id(server, agent, actor, project_id: str):
|
|
"""
|
|
Temporary helper to backfill project_id for legacy agents.
|
|
|
|
TODO(@caren): Remove this function after all existing Claude Code agents have been backfilled.
|
|
|
|
Args:
|
|
server: SyncServer instance
|
|
agent: Agent to update
|
|
actor: Actor performing the operation
|
|
project_id: Project ID to set
|
|
|
|
Returns:
|
|
Updated agent or original agent if update fails
|
|
"""
|
|
from letta.schemas.agent import UpdateAgent
|
|
|
|
try:
|
|
updated_agent = await server.update_agent_async(
|
|
agent_id=agent.id,
|
|
request=UpdateAgent(project_id=project_id),
|
|
actor=actor,
|
|
)
|
|
logger.info(f"[Backfill] Successfully updated agent {agent.id} with project_id {project_id}")
|
|
return updated_agent
|
|
except Exception as e:
|
|
logger.warning(f"[Backfill] Failed to update agent project_id: {e}. Continuing with in-memory update.")
|
|
# Fallback: continue with in-memory update
|
|
agent.project_id = project_id
|
|
return agent
|
|
|
|
|
|
async def get_or_create_claude_code_agent(
|
|
server,
|
|
actor,
|
|
project_id: str = None,
|
|
agent_id: str = None,
|
|
):
|
|
"""
|
|
Get or create a special agent for Claude Code sessions.
|
|
|
|
Args:
|
|
server: SyncServer instance
|
|
actor: Actor performing the operation (user ID)
|
|
project_id: Optional project ID to associate the agent with
|
|
agent_id: Optional specific agent ID to use (from X-LETTA-AGENT-ID header)
|
|
|
|
Returns:
|
|
Agent instance
|
|
"""
|
|
from letta.schemas.agent import CreateAgent
|
|
|
|
# If a specific agent ID is provided, try to use it directly
|
|
if agent_id:
|
|
logger.debug(f"Attempting to fetch agent by ID: {agent_id}")
|
|
try:
|
|
agent = await server.agent_manager.get_agent_by_id_async(agent_id=agent_id, actor=actor)
|
|
logger.info(f"Found agent via X-LETTA-AGENT-ID header: {agent.id} (name: {agent.name})")
|
|
return agent
|
|
except Exception as e:
|
|
logger.warning(f"Could not find agent with ID {agent_id}: {e}. Falling back to default behavior.")
|
|
# Fall through to default behavior below
|
|
|
|
# Create short user identifier from UUID (first 8 chars)
|
|
if actor:
|
|
user_short_id = str(actor.id)[:8] if hasattr(actor, "id") else str(actor)[:8]
|
|
else:
|
|
user_short_id = "default"
|
|
|
|
agent_name = f"claude-code-{user_short_id}"
|
|
|
|
try:
|
|
# Try to find existing agent by name (most reliable)
|
|
# Note: Search by name only, not tags, since name is unique and more reliable
|
|
logger.debug(f"Searching for agent with name: {agent_name}")
|
|
agents = await server.agent_manager.list_agents_async(
|
|
actor=actor,
|
|
limit=10, # Get a few in case of duplicates
|
|
name=agent_name,
|
|
include=["agent.blocks", "agent.managed_group", "agent.tags"],
|
|
)
|
|
|
|
# list_agents_async returns a list directly, not an object with .agents
|
|
logger.debug(f"Agent search returned {len(agents) if agents else 0} results")
|
|
if agents and len(agents) > 0:
|
|
# Return the first matching agent
|
|
logger.info(f"Found existing Claude Code agent: {agents[0].id} (name: {agent_name})")
|
|
agent = agents[0]
|
|
|
|
# Temporary patch: Fix project_id if it's missing (legacy bug)
|
|
# TODO(@caren): Remove this after all existing Claude Code agents have been backfilled
|
|
if not agent.project_id and project_id:
|
|
logger.info(f"[Backfill] Agent {agent.id} missing project_id, backfilling with {project_id}")
|
|
agent = await backfill_agent_project_id(server, agent, actor, project_id)
|
|
|
|
return agent
|
|
else:
|
|
logger.debug(f"No existing agent found with name: {agent_name}")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Could not find existing agent: {e}", exc_info=True)
|
|
|
|
# Create new agent
|
|
try:
|
|
logger.info(f"Creating new Claude Code agent: {agent_name} with project_id: {project_id}")
|
|
|
|
# Create minimal agent config
|
|
agent_config = CreateAgent(
|
|
name=agent_name,
|
|
description="Agent for capturing Claude Code conversations",
|
|
memory_blocks=[
|
|
{
|
|
"label": "human",
|
|
"value": "This is my section of core memory devoted to information about the human.\nI don't yet know anything about them.\nWhat's their name? Where are they from? What do they do? Who are they?\nI should update this memory over time as I interact with the human and learn more about them.",
|
|
"description": "A memory block for keeping track of the human (user) the agent is interacting with.",
|
|
},
|
|
{
|
|
"label": "persona",
|
|
"value": "This is my section of core memory devoted to information myself.\nThere's nothing here yet.\nI should update this memory over time as I develop my personality.",
|
|
"description": "A memory block for storing the agent's core personality details and behavior profile.",
|
|
},
|
|
{
|
|
"label": "project",
|
|
"value": "This is my section of core memory devoted to information about what the agent is working on.\nI don't yet know anything about it.\nI should update this memory over time with high level understanding and learnings.",
|
|
"description": "A memory block for storing the information about the project the agent is working on.",
|
|
},
|
|
],
|
|
tags=["claude-code"],
|
|
enable_sleeptime=True,
|
|
agent_type="letta_v1_agent",
|
|
model="anthropic/claude-sonnet-4-5-20250929",
|
|
embedding="openai/text-embedding-ada-002",
|
|
project_id=project_id,
|
|
)
|
|
|
|
new_agent = await server.create_agent_async(
|
|
request=agent_config,
|
|
actor=actor,
|
|
)
|
|
|
|
logger.info(f"Created Claude Code agent {new_agent.name}: {new_agent.id}")
|
|
return new_agent
|
|
|
|
except Exception as e:
|
|
logger.exception(f"Failed to create Claude Code agent: {e}")
|
|
raise
|