letta-nightly 0.8.3.dev20250612104349__py3-none-any.whl → 0.8.4.dev20250614104137__py3-none-any.whl

letta/otel/metrics.py

@@ -1,3 +1,7 @@
+import re
+import time
+from typing import List
+
 from fastapi import FastAPI, Request
 from opentelemetry import metrics
 from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
@@ -5,8 +9,9 @@ from opentelemetry.metrics import NoOpMeter
 from opentelemetry.sdk.metrics import MeterProvider
 from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
 
+from letta.helpers.datetime_helpers import ns_to_ms
 from letta.log import get_logger
-from letta.otel.context import add_ctx_attribute
+from letta.otel.context import add_ctx_attribute, get_ctx_attributes
 from letta.otel.resource import get_resource, is_pytest_environment
 
 logger = get_logger(__name__)
@@ -14,26 +19,85 @@ logger = get_logger(__name__)
 _meter: metrics.Meter = NoOpMeter("noop")
 _is_metrics_initialized: bool = False
 
+# Endpoints to include in endpoint metrics tracking (opt-in) vs tracing.py opt-out
+_included_v1_endpoints_regex: List[str] = [
+    "^POST /v1/agents/(?P<agent_id>[^/]+)/messages$",
+    "^POST /v1/agents/(?P<agent_id>[^/]+)/messages/stream$",
+    "^POST /v1/agents/(?P<agent_id>[^/]+)/messages/async$",
+]
+
+# Header attributes to set context with
+header_attributes = {
+    "x-organization-id": "organization.id",
+    "x-project-id": "project.id",
+    "x-base-template-id": "base_template.id",
+    "x-template-id": "template.id",
+    "x-agent-id": "agent.id",
+}
+
 
 async def _otel_metric_middleware(request: Request, call_next):
     if not _is_metrics_initialized:
         return await call_next(request)
 
-    header_attributes = {
-        "x-organization-id": "organization.id",
-        "x-project-id": "project.id",
-        "x-base-template-id": "base_template.id",
-        "x-template-id": "template.id",
-        "x-agent-id": "agent.id",
-    }
-    try:
-        for header_key, otel_key in header_attributes.items():
-            header_value = request.headers.get(header_key)
-            if header_value:
-                add_ctx_attribute(otel_key, header_value)
+    for header_key, otel_key in header_attributes.items():
+        header_value = request.headers.get(header_key)
+        if header_value:
+            add_ctx_attribute(otel_key, header_value)
+
+    # Opt-in check for latency / error tracking
+    endpoint_path = f"{request.method} {request.url.path}"
+    should_track_endpoint_metrics = any(re.match(regex, endpoint_path) for regex in _included_v1_endpoints_regex)
+
+    if not should_track_endpoint_metrics:
         return await call_next(request)
-    except Exception:
+
+    # --- Opt-in endpoint metrics ---
+    start_perf_counter_ns = time.perf_counter_ns()
+    response = None
+    status_code = 500  # reasonable default
+
+    try:
+        response = await call_next(request)
+        status_code = response.status_code
+        return response
+    except Exception as e:
+        # Determine status code from exception
+        status_code = getattr(e, "status_code", 500)
         raise
+    finally:
+        end_to_end_ms = ns_to_ms(time.perf_counter_ns() - start_perf_counter_ns)
+        _record_endpoint_metrics(
+            request=request,
+            latency_ms=end_to_end_ms,
+            status_code=status_code,
+        )
+
+
+def _record_endpoint_metrics(
+    request: Request,
+    latency_ms: float,
+    status_code: int,
+):
+    """Record endpoint latency and request count metrics."""
+    try:
+        # Get the route pattern for better endpoint naming
+        route = request.scope.get("route")
+        endpoint_name = route.path if route and hasattr(route, "path") else "unknown"
+
+        attrs = {
+            "endpoint_path": endpoint_name,
+            "method": request.method,
+            "status_code": status_code,
+            **get_ctx_attributes(),
+        }
+        from letta.otel.metric_registry import MetricRegistry
+
+        MetricRegistry().endpoint_e2e_ms_histogram.record(latency_ms, attributes=attrs)
+        MetricRegistry().endpoint_request_counter.add(1, attributes=attrs)
+
+    except Exception as e:
+        logger.warning(f"Failed to record endpoint metrics: {e}")
 
 
 def setup_metrics(

letta/schemas/letta_message_content.py

@@ -6,6 +6,7 @@ from pydantic import BaseModel, Field
 
 class MessageContentType(str, Enum):
     text = "text"
+    image = "image"
     tool_call = "tool_call"
     tool_return = "tool_return"
     reasoning = "reasoning"
@@ -18,7 +19,7 @@ class MessageContent(BaseModel):
 
 
 # -------------------------------
-# User Content Types
+# Text Content
 # -------------------------------
 
 
@@ -27,8 +28,62 @@ class TextContent(MessageContent):
     text: str = Field(..., description="The text content of the message.")
 
 
+# -------------------------------
+# Image Content
+# -------------------------------
+
+
+class ImageSourceType(str, Enum):
+    url = "url"
+    base64 = "base64"
+    letta = "letta"
+
+
+class ImageSource(BaseModel):
+    type: ImageSourceType = Field(..., description="The source type for the image.")
+
+
+class UrlImage(ImageSource):
+    type: Literal[ImageSourceType.url] = Field(ImageSourceType.url, description="The source type for the image.")
+    url: str = Field(..., description="The URL of the image.")
+
+
+class Base64Image(ImageSource):
+    type: Literal[ImageSourceType.base64] = Field(ImageSourceType.base64, description="The source type for the image.")
+    media_type: str = Field(..., description="The media type for the image.")
+    data: str = Field(..., description="The base64 encoded image data.")
+    detail: Optional[str] = Field(
+        None,
+        description="What level of detail to use when processing and understanding the image (low, high, or auto to let the model decide)",
+    )
+
+
+class LettaImage(ImageSource):
+    type: Literal[ImageSourceType.letta] = Field(ImageSourceType.letta, description="The source type for the image.")
+    file_id: str = Field(..., description="The unique identifier of the image file persisted in storage.")
+    media_type: Optional[str] = Field(None, description="The media type for the image.")
+    data: Optional[str] = Field(None, description="The base64 encoded image data.")
+    detail: Optional[str] = Field(
+        None,
+        description="What level of detail to use when processing and understanding the image (low, high, or auto to let the model decide)",
+    )
+
+
+ImageSourceUnion = Annotated[Union[UrlImage, Base64Image, LettaImage], Field(discriminator="type")]
+
+
+class ImageContent(MessageContent):
+    type: Literal[MessageContentType.image] = Field(MessageContentType.image, description="The type of the message.")
+    source: ImageSourceUnion = Field(..., description="The source of the image.")
+
+
+# -------------------------------
+# User Content Types
+# -------------------------------
+
+
 LettaUserMessageContentUnion = Annotated[
-    Union[TextContent],
+    Union[TextContent, ImageContent],
     Field(discriminator="type"),
 ]
 
@@ -37,11 +92,13 @@ def create_letta_user_message_content_union_schema():
     return {
         "oneOf": [
             {"$ref": "#/components/schemas/TextContent"},
+            {"$ref": "#/components/schemas/ImageContent"},
         ],
         "discriminator": {
             "propertyName": "type",
             "mapping": {
                 "text": "#/components/schemas/TextContent",
+                "image": "#/components/schemas/ImageContent",
             },
         },
     }
@@ -150,7 +207,9 @@ class OmittedReasoningContent(MessageContent):
 
 
 LettaMessageContentUnion = Annotated[
-    Union[TextContent, ToolCallContent, ToolReturnContent, ReasoningContent, RedactedReasoningContent, OmittedReasoningContent],
+    Union[
+        TextContent, ImageContent, ToolCallContent, ToolReturnContent, ReasoningContent, RedactedReasoningContent, OmittedReasoningContent
+    ],
     Field(discriminator="type"),
 ]
 
@@ -159,6 +218,7 @@ def create_letta_message_content_union_schema():
     return {
         "oneOf": [
             {"$ref": "#/components/schemas/TextContent"},
+            {"$ref": "#/components/schemas/ImageContent"},
             {"$ref": "#/components/schemas/ToolCallContent"},
             {"$ref": "#/components/schemas/ToolReturnContent"},
             {"$ref": "#/components/schemas/ReasoningContent"},
@@ -169,6 +229,7 @@ def create_letta_message_content_union_schema():
             "propertyName": "type",
             "mapping": {
                 "text": "#/components/schemas/TextContent",
+                "image": "#/components/schemas/ImageContent",
                 "tool_call": "#/components/schemas/ToolCallContent",
                 "tool_return": "#/components/schemas/ToolCallContent",
                 "reasoning": "#/components/schemas/ReasoningContent",

letta/schemas/letta_request.py

@@ -2,13 +2,17 @@ from typing import List, Optional
 
 from pydantic import BaseModel, Field, HttpUrl
 
-from letta.constants import DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
+from letta.constants import DEFAULT_MAX_STEPS, DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
 from letta.schemas.letta_message import MessageType
 from letta.schemas.message import MessageCreate
 
 
 class LettaRequest(BaseModel):
     messages: List[MessageCreate] = Field(..., description="The messages to be sent to the agent.")
+    max_steps: int = Field(
+        default=DEFAULT_MAX_STEPS,
+        description="Maximum number of steps the agent should take to process the request.",
+    )
     use_assistant_message: bool = Field(
         default=True,
         description="Whether the server should parse specific tool call arguments (default `send_message`) as `AssistantMessage` objects.",

letta/schemas/message.py

@@ -31,6 +31,7 @@ from letta.schemas.letta_message import (
     UserMessage,
 )
 from letta.schemas.letta_message_content import (
+    ImageContent,
     LettaMessageContentUnion,
     OmittedReasoningContent,
     ReasoningContent,
@@ -415,15 +416,17 @@ class Message(BaseMessage):
             # This is type UserMessage
             if self.content and len(self.content) == 1 and isinstance(self.content[0], TextContent):
                 text_content = self.content[0].text
+            elif self.content:
+                text_content = self.content
             else:
                 raise ValueError(f"Invalid user message (no text object on message): {self.content}")
 
-            message_str = unpack_message(text_content)
+            message = unpack_message(text_content)
             messages.append(
                 UserMessage(
                     id=self.id,
                     date=self.created_at,
-                    content=message_str or text_content,
+                    content=message,
                     name=self.name,
                     otid=self.otid,
                     sender_id=self.sender_id,
@@ -658,13 +661,14 @@ class Message(BaseMessage):
             text_content = self.content[0].text
         elif self.content and len(self.content) == 1 and isinstance(self.content[0], ToolReturnContent):
             text_content = self.content[0].content
+        elif self.content and len(self.content) == 1 and isinstance(self.content[0], ImageContent):
+            text_content = "[Image Here]"
         # Otherwise, check if we have TextContent and multiple other parts
         elif self.content and len(self.content) > 1:
             text = [content for content in self.content if isinstance(content, TextContent)]
-            if len(text) > 1:
-                assert len(text) == 1, f"multiple text content parts found in a single message: {self.content}"
-                text_content = text[0].text
-                parse_content_parts = True
+            assert len(text) == 1, f"multiple text content parts found in a single message: {self.content}"
+            text_content = text[0].text
+            parse_content_parts = True
         else:
             text_content = None
 
@@ -778,11 +782,35 @@ class Message(BaseMessage):
             }
 
         elif self.role == "user":
-            assert all([v is not None for v in [text_content, self.role]]), vars(self)
-            anthropic_message = {
-                "content": text_content,
-                "role": self.role,
-            }
+            # special case for text-only message
+            if text_content is not None:
+                anthropic_message = {
+                    "content": text_content,
+                    "role": self.role,
+                }
+            else:
+                content_parts = []
+                for content in self.content:
+                    if isinstance(content, TextContent):
+                        content_parts.append({"type": "text", "text": content.text})
+                    elif isinstance(content, ImageContent):
+                        content_parts.append(
+                            {
+                                "type": "image",
+                                "source": {
+                                    "type": "base64",
+                                    "data": content.source.data,
+                                    "media_type": content.source.media_type,
+                                },
+                            }
+                        )
+                    else:
+                        raise ValueError(f"Unsupported content type: {content.type}")
+
+                anthropic_message = {
+                    "content": content_parts,
+                    "role": self.role,
+                }
 
         elif self.role == "assistant":
             assert self.tool_calls is not None or text_content is not None
@@ -887,10 +915,27 @@ class Message(BaseMessage):
             }
 
         elif self.role == "user":
-            assert all([v is not None for v in [text_content, self.role]]), vars(self)
+            assert self.content, vars(self)
+
+            content_parts = []
+            for content in self.content:
+                if isinstance(content, TextContent):
+                    content_parts.append({"text": content.text})
+                elif isinstance(content, ImageContent):
+                    content_parts.append(
+                        {
+                            "inline_data": {
+                                "data": content.source.data,
+                                "mime_type": content.source.media_type,
+                            }
+                        }
+                    )
+                else:
+                    raise ValueError(f"Unsupported content type: {content.type}")
+
             google_ai_message = {
                 "role": "user",
-                "parts": [{"text": text_content}],
+                "parts": content_parts,
             }
 
         elif self.role == "assistant":
@@ -1006,8 +1051,10 @@ class Message(BaseMessage):
         # embedded function calls in multi-turn conversation become more clear
         if self.content and len(self.content) == 1 and isinstance(self.content[0], TextContent):
             text_content = self.content[0].text
-        if self.content and len(self.content) == 1 and isinstance(self.content[0], ToolReturnContent):
+        elif self.content and len(self.content) == 1 and isinstance(self.content[0], ToolReturnContent):
             text_content = self.content[0].content
+        elif self.content and len(self.content) == 1 and isinstance(self.content[0], ImageContent):
+            text_content = "[Image Here]"
         else:
             text_content = None
         if self.role == "system":

letta/schemas/openai/chat_completion_request.py

@@ -10,7 +10,7 @@ class SystemMessage(BaseModel):
 
 
 class UserMessage(BaseModel):
-    content: Union[str, List[str]]
+    content: Union[str, List[str], List[dict]]
     role: str = "user"
     name: Optional[str] = None
 

letta/schemas/providers.py

@@ -2,6 +2,8 @@ import warnings
 from datetime import datetime
 from typing import List, Literal, Optional
 
+import aiohttp
+import requests
 from pydantic import BaseModel, Field, model_validator
 
 from letta.constants import DEFAULT_EMBEDDING_CHUNK_SIZE, LETTA_MODEL_ENDPOINT, LLM_MAX_TOKENS, MIN_CONTEXT_WINDOW
@@ -872,9 +874,6 @@ class OllamaProvider(OpenAIProvider):
     async def list_llm_models_async(self) -> List[LLMConfig]:
         """Async version of list_llm_models below"""
         endpoint = f"{self.base_url}/api/tags"
-
-        import aiohttp
-
         async with aiohttp.ClientSession() as session:
             async with session.get(endpoint) as response:
                 if response.status != 200:
@@ -903,8 +902,6 @@ class OllamaProvider(OpenAIProvider):
 
     def list_llm_models(self) -> List[LLMConfig]:
         # https://github.com/ollama/ollama/blob/main/docs/api.md#list-local-models
-        import requests
-
         response = requests.get(f"{self.base_url}/api/tags")
         if response.status_code != 200:
             raise Exception(f"Failed to list Ollama models: {response.text}")
@@ -931,9 +928,6 @@ class OllamaProvider(OpenAIProvider):
         return configs
 
     def get_model_context_window(self, model_name: str) -> Optional[int]:
-
-        import requests
-
         response = requests.post(f"{self.base_url}/api/show", json={"name": model_name, "verbose": True})
         response_json = response.json()
 
@@ -965,11 +959,19 @@ class OllamaProvider(OpenAIProvider):
                 return value
         return None
 
-    def get_model_embedding_dim(self, model_name: str):
-        import requests
-
+    def _get_model_embedding_dim(self, model_name: str):
         response = requests.post(f"{self.base_url}/api/show", json={"name": model_name, "verbose": True})
         response_json = response.json()
+        return self._get_model_embedding_dim_impl(response_json, model_name)
+
+    async def _get_model_embedding_dim_async(self, model_name: str):
+        async with aiohttp.ClientSession() as session:
+            async with session.post(f"{self.base_url}/api/show", json={"name": model_name, "verbose": True}) as response:
+                response_json = await response.json()
+            return self._get_model_embedding_dim_impl(response_json, model_name)
+
+    @staticmethod
+    def _get_model_embedding_dim_impl(response_json: dict, model_name: str):
         if "model_info" not in response_json:
             if "error" in response_json:
                 print(f"Ollama fetch model info error for {model_name}: {response_json['error']}")
@@ -979,10 +981,35 @@ class OllamaProvider(OpenAIProvider):
                 return value
         return None
 
+    async def list_embedding_models_async(self) -> List[EmbeddingConfig]:
+        """Async version of list_embedding_models below"""
+        endpoint = f"{self.base_url}/api/tags"
+        async with aiohttp.ClientSession() as session:
+            async with session.get(endpoint) as response:
+                if response.status != 200:
+                    raise Exception(f"Failed to list Ollama models: {response.text}")
+                response_json = await response.json()
+
+        configs = []
+        for model in response_json["models"]:
+            embedding_dim = await self._get_model_embedding_dim_async(model["name"])
+            if not embedding_dim:
+                print(f"Ollama model {model['name']} has no embedding dimension")
+                continue
+            configs.append(
+                EmbeddingConfig(
+                    embedding_model=model["name"],
+                    embedding_endpoint_type="ollama",
+                    embedding_endpoint=self.base_url,
+                    embedding_dim=embedding_dim,
+                    embedding_chunk_size=300,
+                    handle=self.get_handle(model["name"], is_embedding=True),
+                )
+            )
+        return configs
+
     def list_embedding_models(self) -> List[EmbeddingConfig]:
         # https://github.com/ollama/ollama/blob/main/docs/api.md#list-local-models
-        import requests
-
         response = requests.get(f"{self.base_url}/api/tags")
         if response.status_code != 200:
             raise Exception(f"Failed to list Ollama models: {response.text}")
@@ -990,7 +1017,7 @@ class OllamaProvider(OpenAIProvider):
 
         configs = []
         for model in response_json["models"]:
-            embedding_dim = self.get_model_embedding_dim(model["name"])
+            embedding_dim = self._get_model_embedding_dim(model["name"])
             if not embedding_dim:
                 print(f"Ollama model {model['name']} has no embedding dimension")
                 continue

letta/schemas/tool_rule.py

@@ -1,20 +1,48 @@
 import json
+import logging
 from typing import Annotated, Any, Dict, List, Literal, Optional, Set, Union
 
+from jinja2 import Template
 from pydantic import Field
 
 from letta.schemas.enums import ToolRuleType
 from letta.schemas.letta_base import LettaBase
 
+logger = logging.getLogger(__name__)
+
 
 class BaseToolRule(LettaBase):
     __id_prefix__ = "tool_rule"
     tool_name: str = Field(..., description="The name of the tool. Must exist in the database for the user's organization.")
     type: ToolRuleType = Field(..., description="The type of the message.")
+    prompt_template: Optional[str] = Field(
+        None,
+        description="Optional Jinja2 template for generating agent prompt about this tool rule. Template can use variables like 'tool_name' and rule-specific attributes.",
+    )
 
     def get_valid_tools(self, tool_call_history: List[str], available_tools: Set[str], last_function_response: Optional[str]) -> set[str]:
         raise NotImplementedError
 
+    def render_prompt(self) -> Optional[str]:
+        """Render the prompt template with this rule's attributes."""
+        template_to_use = self.prompt_template or self._get_default_template()
+        if not template_to_use:
+            return None
+
+        try:
+            template = Template(template_to_use)
+            return template.render(**self.model_dump())
+        except Exception as e:
+            logger.warning(
+                f"Failed to render prompt template for tool rule '{self.tool_name}' (type: {self.type}). "
+                f"Template: '{template_to_use}'. Error: {e}"
+            )
+            return None
+
+    def _get_default_template(self) -> Optional[str]:
+        """Get the default template for this rule type. Override in subclasses."""
+        return None
+
 
 class ChildToolRule(BaseToolRule):
     """
@@ -23,11 +51,18 @@ class ChildToolRule(BaseToolRule):
 
     type: Literal[ToolRuleType.constrain_child_tools] = ToolRuleType.constrain_child_tools
     children: List[str] = Field(..., description="The children tools that can be invoked.")
+    prompt_template: Optional[str] = Field(
+        default="<tool_constraint>After using {{ tool_name }}, you can only use these tools: {{ children | join(', ') }}</tool_constraint>",
+        description="Optional Jinja2 template for generating agent prompt about this tool rule.",
+    )
 
     def get_valid_tools(self, tool_call_history: List[str], available_tools: Set[str], last_function_response: Optional[str]) -> Set[str]:
         last_tool = tool_call_history[-1] if tool_call_history else None
         return set(self.children) if last_tool == self.tool_name else available_tools
 
+    def _get_default_template(self) -> Optional[str]:
+        return "<tool_constraint>After using {{ tool_name }}, you can only use these tools: {{ children | join(', ') }}</tool_constraint>"
+
 
 class ParentToolRule(BaseToolRule):
     """
@@ -36,11 +71,18 @@ class ParentToolRule(BaseToolRule):
 
     type: Literal[ToolRuleType.parent_last_tool] = ToolRuleType.parent_last_tool
     children: List[str] = Field(..., description="The children tools that can be invoked.")
+    prompt_template: Optional[str] = Field(
+        default="<tool_constraint>{{ children | join(', ') }} can only be used after {{ tool_name }}</tool_constraint>",
+        description="Optional Jinja2 template for generating agent prompt about this tool rule.",
+    )
 
     def get_valid_tools(self, tool_call_history: List[str], available_tools: Set[str], last_function_response: Optional[str]) -> Set[str]:
         last_tool = tool_call_history[-1] if tool_call_history else None
         return set(self.children) if last_tool == self.tool_name else available_tools - set(self.children)
 
+    def _get_default_template(self) -> Optional[str]:
+        return "<tool_constraint>{{ children | join(', ') }} can only be used after {{ tool_name }}</tool_constraint>"
+
 
 class ConditionalToolRule(BaseToolRule):
     """
@@ -51,6 +93,10 @@ class ConditionalToolRule(BaseToolRule):
     default_child: Optional[str] = Field(None, description="The default child tool to be called. If None, any tool can be called.")
     child_output_mapping: Dict[Any, str] = Field(..., description="The output case to check for mapping")
     require_output_mapping: bool = Field(default=False, description="Whether to throw an error when output doesn't match any case")
+    prompt_template: Optional[str] = Field(
+        default="<tool_constraint>{{ tool_name }} will determine which tool to use next based on its output</tool_constraint>",
+        description="Optional Jinja2 template for generating agent prompt about this tool rule.",
+    )
 
     def get_valid_tools(self, tool_call_history: List[str], available_tools: Set[str], last_function_response: Optional[str]) -> Set[str]:
         """Determine valid tools based on function output mapping."""
@@ -96,6 +142,9 @@ class ConditionalToolRule(BaseToolRule):
         else:  # Assume string
             return str(function_output) == str(key)
 
+    def _get_default_template(self) -> Optional[str]:
+        return "<tool_constraint>{{ tool_name }} will determine which tool to use next based on its output</tool_constraint>"
+
 
 class InitToolRule(BaseToolRule):
     """
@@ -111,6 +160,13 @@ class TerminalToolRule(BaseToolRule):
     """
 
     type: Literal[ToolRuleType.exit_loop] = ToolRuleType.exit_loop
+    prompt_template: Optional[str] = Field(
+        default="<tool_constraint>{{ tool_name }} ends the conversation when called</tool_constraint>",
+        description="Optional Jinja2 template for generating agent prompt about this tool rule.",
+    )
+
+    def _get_default_template(self) -> Optional[str]:
+        return "<tool_constraint>{{ tool_name }} ends the conversation when called</tool_constraint>"
 
 
 class ContinueToolRule(BaseToolRule):
@@ -119,6 +175,10 @@ class ContinueToolRule(BaseToolRule):
     """
 
     type: Literal[ToolRuleType.continue_loop] = ToolRuleType.continue_loop
+    prompt_template: Optional[str] = Field(
+        default="<tool_constraint>{{ tool_name }} requires continuing the conversation when called</tool_constraint>",
+        description="Optional Jinja2 template for generating agent prompt about this tool rule.",
+    )
 
 
 class MaxCountPerStepToolRule(BaseToolRule):
@@ -128,6 +188,10 @@ class MaxCountPerStepToolRule(BaseToolRule):
 
     type: Literal[ToolRuleType.max_count_per_step] = ToolRuleType.max_count_per_step
     max_count_limit: int = Field(..., description="The max limit for the total number of times this tool can be invoked in a single step.")
+    prompt_template: Optional[str] = Field(
+        default="<tool_constraint>{{ tool_name }}: max {{ max_count_limit }} use(s) per turn</tool_constraint>",
+        description="Optional Jinja2 template for generating agent prompt about this tool rule.",
+    )
 
     def get_valid_tools(self, tool_call_history: List[str], available_tools: Set[str], last_function_response: Optional[str]) -> Set[str]:
         """Restricts the tool if it has been called max_count_limit times in the current step."""
@@ -139,6 +203,9 @@ class MaxCountPerStepToolRule(BaseToolRule):
 
         return available_tools
 
+    def _get_default_template(self) -> Optional[str]:
+        return "<tool_constraint>{{ tool_name }}: max {{ max_count_limit }} use(s) per turn</tool_constraint>"
+
 
 ToolRule = Annotated[
     Union[ChildToolRule, InitToolRule, TerminalToolRule, ConditionalToolRule, ContinueToolRule, MaxCountPerStepToolRule, ParentToolRule],

letta/schemas/user.py

@@ -3,8 +3,8 @@ from typing import Optional
 
 from pydantic import Field
 
+from letta.constants import DEFAULT_ORG_ID
 from letta.schemas.letta_base import LettaBase
-from letta.services.organization_manager import OrganizationManager
 
 
 class UserBase(LettaBase):
@@ -22,7 +22,7 @@ class User(UserBase):
     """
 
     id: str = UserBase.generate_id_field()
-    organization_id: Optional[str] = Field(OrganizationManager.DEFAULT_ORG_ID, description="The organization id of the user")
+    organization_id: Optional[str] = Field(DEFAULT_ORG_ID, description="The organization id of the user")
     name: str = Field(..., description="The name of the user.")
     created_at: Optional[datetime] = Field(default_factory=datetime.utcnow, description="The creation date of the user.")
     updated_at: Optional[datetime] = Field(default_factory=datetime.utcnow, description="The update date of the user.")