diff --git a/backend/alembic/versions/v1_12_0_f066_configurable_memory_context.py b/backend/alembic/versions/v1_12_0_f066_configurable_memory_context.py new file mode 100644 index 000000000..ef1af704a --- /dev/null +++ b/backend/alembic/versions/v1_12_0_f066_configurable_memory_context.py @@ -0,0 +1,65 @@ +"""Add configurable memory context limit to agents. + +Background: + The memory.md context limit was previously hardcoded to 2,000 characters. + +Scope: + Add a per-agent memory_context_max_chars column with a default value of 2,000. + +Idempotent: + The column is added only when it does not already exist. + +Revision ID: f066_configurable_memory_context +Revises: f065_feishu_group_target +Create Date: 2026-09-02 +""" + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +revision: str = "f066_configurable_memory_context" +down_revision: str | Sequence[str] | None = "f065_feishu_group_target" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + + columns = { + column["name"] + for column in inspector.get_columns("agents") + } + + if "memory_context_max_chars" not in columns: + op.add_column( + "agents", + sa.Column( + "memory_context_max_chars", + sa.Integer(), + nullable=False, + server_default=sa.text("2000"), + ), + ) + op.alter_column( + "agents", + "memory_context_max_chars", + server_default=None, + ) + + +def downgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + + columns = { + column["name"] + for column in inspector.get_columns("agents") + } + + if "memory_context_max_chars" in columns: + op.drop_column("agents", "memory_context_max_chars") \ No newline at end of file diff --git a/backend/app/models/agent.py b/backend/app/models/agent.py index d2e88f7ae..7fc7d4aa8 100644 --- a/backend/app/models/agent.py +++ b/backend/app/models/agent.py @@ -96,6 +96,11 @@ class Agent(Base): cache_creation_tokens_month: Mapped[int] = mapped_column(Integer, default=0) cache_creation_tokens_total: Mapped[int] = mapped_column(Integer, default=0) context_window_size: Mapped[int] = mapped_column(Integer, default=100) + memory_context_max_chars: Mapped[int] = mapped_column( + Integer, + default=2000, + nullable=False, + ) # Historical field name: this is the maximum number of model-decision turns # allowed for one Agent Run, not the number of tools executed. max_tool_rounds: Mapped[int] = mapped_column(Integer, default=50) diff --git a/backend/app/schemas/schemas.py b/backend/app/schemas/schemas.py index 76d3af6fa..270b7b004 100644 --- a/backend/app/schemas/schemas.py +++ b/backend/app/schemas/schemas.py @@ -297,6 +297,7 @@ class AgentOut(BaseModel): created_at: datetime last_active_at: datetime | None = None deleted_at: datetime | None = None + memory_context_max_chars: int = 2000 model_config = {"from_attributes": True} @@ -322,6 +323,11 @@ class AgentUpdate(BaseModel): heartbeat_active_hours: str | None = None timezone: str | None = None expires_at: datetime | None = None # Admin only — extend agent expiry + memory_context_max_chars: int | None = Field( + default=None, + ge=0, + le=100000, + ) @field_validator("timezone") @classmethod diff --git a/backend/app/services/agent_context.py b/backend/app/services/agent_context.py index 6045863b2..18f1a9249 100644 --- a/backend/app/services/agent_context.py +++ b/backend/app/services/agent_context.py @@ -437,6 +437,7 @@ async def build_agent_context( current_user_name: str | None = None, *, allowed_tool_names: Collection[str] | None = None, + memory_context_max_chars: int = 2000, ) -> tuple[str, str]: """Build Base Prompt V1 plus bounded, explicitly low-trust context data.""" # `role_description` remains product metadata and is intentionally ignored by @@ -463,12 +464,12 @@ async def build_agent_context( memory = await _read_file_safe( normalize_storage_key(f"{agent_id}/memory/memory.md"), - 2000, + memory_context_max_chars, ) if not memory: memory = await _read_file_safe( normalize_storage_key(f"{agent_id}/memory.md"), - 2000, + memory_context_max_chars, ) if memory.startswith("# "): memory = "\n".join(memory.split("\n")[1:]).strip() diff --git a/backend/app/services/llm/caller.py b/backend/app/services/llm/caller.py index 35dbb6211..5d6c34750 100644 --- a/backend/app/services/llm/caller.py +++ b/backend/app/services/llm/caller.py @@ -225,26 +225,44 @@ def _usage_from_response_or_estimate(response, api_messages: list[LLMMessage]) - # Helper Functions # ═══════════════════════════════════════════════════════════════════════════════ -async def _get_agent_config(agent_id) -> tuple[int, str | None]: - """Get agent config: max_tool_rounds and token limit status.""" +async def _get_agent_config(agent_id) -> tuple[int, int, str | None]: + """Get agent config: max_tool_rounds, memory context limit and token limit status.""" if not agent_id: - return 50, None + return 50, 2000, None try: from app.models.agent import Agent as AgentModel + async with async_session() as _db: _ar = await _db.execute(select(AgentModel).where(AgentModel.id == agent_id)) _agent = _ar.scalar_one_or_none() if _agent: max_rounds = _agent.max_tool_rounds or 50 + memory_context_max_chars = ( + _agent.memory_context_max_chars + if _agent.memory_context_max_chars is not None + else 2000 + ) if _agent.max_tokens_per_day and _agent.tokens_used_today >= _agent.max_tokens_per_day: - return max_rounds, f"⚠️ Daily token usage has reached the limit ({_agent.tokens_used_today:,}/{_agent.max_tokens_per_day:,}). Please try again tomorrow or ask admin to increase the limit." + return ( + max_rounds, + memory_context_max_chars, + f"⚠ Daily token usage has reached the limit " + f"({_agent.tokens_used_today:,}/{_agent.max_tokens_per_day:,}). " + "Please try again tomorrow or ask admin to increase the limit.", + ) if _agent.max_tokens_per_month and _agent.tokens_used_month >= _agent.max_tokens_per_month: - return max_rounds, f"⚠️ Monthly token usage has reached the limit ({_agent.tokens_used_month:,}/{_agent.max_tokens_per_month:,}). Please ask admin to increase the limit." - return max_rounds, None + return ( + max_rounds, + memory_context_max_chars, + f"⚠ Monthly token usage has reached the limit " + f"({_agent.tokens_used_month:,}/{_agent.max_tokens_per_month:,}). " + "Please ask admin to increase the limit.", + ) + return max_rounds, memory_context_max_chars, None except Exception: pass - return 50, None + return 50, 2000, None async def _get_user_name(user_id) -> str | None: @@ -495,7 +513,7 @@ async def call_llm( ) -> str: """Call LLM via unified client with function-calling tool loop.""" # Get agent config for tool rounds - _max_tool_rounds, _token_limit_msg = await _get_agent_config(agent_id) + _max_tool_rounds, _memory_context_max_chars, _token_limit_msg = await _get_agent_config(agent_id) if _token_limit_msg: return _token_limit_msg if max_tool_rounds_override and max_tool_rounds_override < _max_tool_rounds: @@ -549,6 +567,7 @@ async def _default_on_tool_call(data: dict): "", current_user_name=_user_name, allowed_tool_names=allowed_tool_names, + memory_context_max_chars=_memory_context_max_chars, ) if system_prompt_suffix: dynamic_prompt = f"{dynamic_prompt}\n\n{system_prompt_suffix.strip()}" @@ -647,7 +666,7 @@ async def _completion_failure(code: str, message: str) -> str: if agent_id and _unsaved_usage.total_tokens > 0: await record_token_usage(agent_id, _unsaved_usage) _unsaved_usage = TokenUsage() - _, _token_limit_msg = await _get_agent_config(agent_id) + _, _, _token_limit_msg = await _get_agent_config(agent_id) if _token_limit_msg: logger.warning(f"[LLM] Token limit exceeded mid-loop: {_token_limit_msg}") await client.close() diff --git a/backend/tests/test_agent_context.py b/backend/tests/test_agent_context.py index d366e426c..e9622b380 100644 --- a/backend/tests/test_agent_context.py +++ b/backend/tests/test_agent_context.py @@ -67,6 +67,45 @@ async def test_base_prompt_starts_with_name_and_soul_and_never_injects_self_role assert "call `finish`" not in static assert "return the exact final answer as normal Assistant content" in static +@pytest.mark.asyncio +async def test_memory_context_uses_configured_character_limit(): + from app.services.agent_context import build_agent_context + + agent_id = uuid.uuid4() + observed_limits = [] + + async def fake_read_file(key, max_chars=3000): + if key.endswith("/memory/memory.md"): + observed_limits.append(max_chars) + return "memory" + return "" + + with ( + patch("app.services.agent_context._read_file_safe", side_effect=fake_read_file), + patch( + "app.services.agent_context._load_skills_index", + new_callable=AsyncMock, + return_value="", + ), + patch( + "app.services.agent_context._load_relationships_from_db", + new_callable=AsyncMock, + return_value="", + ), + patch( + "app.services.timezone_utils.get_agent_timezone", + new_callable=AsyncMock, + return_value="UTC", + ), + ): + await build_agent_context( + agent_id, + "TestAgent", + allowed_tool_names={"wait"}, + memory_context_max_chars=500, + ) + + assert observed_limits == [500] @pytest.mark.asyncio async def test_focus_mechanism_is_constant_but_tool_policy_follows_effective_tools(): @@ -228,4 +267,4 @@ async def test_experience_policy_is_short_and_only_names_enabled_operations(): assert "read_experience" in read_only assert "propose_experience_draft" not in read_only assert "现有标签" not in read_only - assert "propose_experience_draft" in with_draft + assert "propose_experience_draft" in with_draft \ No newline at end of file diff --git a/backend/tests/test_finish_protocol.py b/backend/tests/test_finish_protocol.py index b1c24296f..50bbbc08d 100644 --- a/backend/tests/test_finish_protocol.py +++ b/backend/tests/test_finish_protocol.py @@ -354,7 +354,7 @@ async def test_call_llm_returns_natural_assistant_stop_without_finish(monkeypatc _plain_response("Final answer."), ]) - monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((3, None))) + monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((3, 2000, None))) monkeypatch.setattr(caller, "_get_user_name", lambda _user_id: _async_return("Ray")) monkeypatch.setattr( "app.services.agent_context.build_agent_context", @@ -408,7 +408,7 @@ async def test_call_llm_routes_embedded_thinking_before_final_content(monkeypatc monkeypatch.setattr( caller, "_get_agent_config", - lambda _agent_id: _async_return((3, None)), + lambda _agent_id: _async_return((3, 2000, None)), ) monkeypatch.setattr( caller, @@ -471,7 +471,7 @@ async def test_call_llm_executes_exact_textual_tool_call_before_finishing( monkeypatch.setattr( caller, "_get_agent_config", - lambda _agent_id: _async_return((3, None)), + lambda _agent_id: _async_return((3, 2000, None)), ) monkeypatch.setattr( caller, @@ -553,7 +553,7 @@ async def test_call_llm_repairs_textual_result_instead_of_publishing_it(monkeypa monkeypatch.setattr( caller, "_get_agent_config", - lambda _agent_id: _async_return((3, None)), + lambda _agent_id: _async_return((3, 2000, None)), ) monkeypatch.setattr( caller, @@ -622,7 +622,7 @@ async def test_legacy_tool_loop_calls_saved_model_without_verified_tool_calling( model.supports_tool_calling = supports_tool_calling fake_client = FakeStreamClient([_finish_response("Final answer.")]) - monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((50, None))) + monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((50, 2000, None))) monkeypatch.setattr(caller, "_get_user_name", lambda _user_id: _async_return("Ray")) monkeypatch.setattr( "app.services.agent_context.build_agent_context", @@ -662,7 +662,7 @@ async def test_call_llm_truncated_output_repair_is_bounded(monkeypatch): _plain_response("Second partial response.", finish_reason="length"), ]) - monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((50, None))) + monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((50, 2000, None))) monkeypatch.setattr(caller, "_get_user_name", lambda _user_id: _async_return("Ray")) monkeypatch.setattr( "app.services.agent_context.build_agent_context", @@ -716,7 +716,7 @@ async def test_invalid_finish_does_not_stop_and_is_returned_as_tool_error(monkey _finish_response("Recovered final."), ]) - monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((3, None))) + monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((3, 2000, None))) monkeypatch.setattr(caller, "_get_user_name", lambda _user_id: _async_return("Ray")) monkeypatch.setattr( "app.services.agent_context.build_agent_context", @@ -766,7 +766,7 @@ async def test_repeated_invalid_finish_is_bounded_by_protocol_code(monkeypatch): _finish_response_with_arguments("{}"), _finish_response_with_arguments("{}"), ]) - monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((50, None))) + monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((50, 2000, None))) monkeypatch.setattr(caller, "_get_user_name", lambda _user_id: _async_return("Ray")) monkeypatch.setattr( "app.services.agent_context.build_agent_context", @@ -815,7 +815,7 @@ async def test_repeated_invalid_tool_json_is_bounded_by_protocol_code(monkeypatc ], ) fake_client = FakeStreamClient([invalid] * 11) - monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((50, None))) + monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((50, 2000, None))) monkeypatch.setattr(caller, "_get_user_name", lambda _user_id: _async_return("Ray")) monkeypatch.setattr( "app.services.agent_context.build_agent_context", @@ -867,7 +867,7 @@ async def test_invalid_write_file_json_gets_ten_bounded_repairs(monkeypatch): ], ) fake_client = FakeStreamClient([invalid] * 11) - monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((50, None))) + monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((50, 2000, None))) monkeypatch.setattr(caller, "_get_user_name", lambda _user_id: _async_return("Ray")) monkeypatch.setattr( "app.services.agent_context.build_agent_context", @@ -910,7 +910,7 @@ async def test_skip_tools_uses_natural_completion_without_any_tools(monkeypatch) fake_client = FakeStreamClient([_plain_response("Onboarding done.")]) - monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((1, None))) + monkeypatch.setattr(caller, "_get_agent_config", lambda _agent_id: _async_return((1, 2000, None))) monkeypatch.setattr(caller, "_get_user_name", lambda _user_id: _async_return("Ray")) monkeypatch.setattr( "app.services.agent_context.build_agent_context", @@ -997,8 +997,8 @@ async def mock_get_agent_config(agent_id): nonlocal configs_called configs_called += 1 if configs_called > 1: - return 50, "⚠️ Daily token usage limit exceeded" - return 50, None + return 50, 2000, "⚠️ Daily token usage limit exceeded" + return 50, 2000, None monkeypatch.setattr(caller, "_get_agent_config", mock_get_agent_config) monkeypatch.setattr(caller, "_get_user_name", lambda _user_id: _async_return("Ray")) @@ -1042,4 +1042,4 @@ async def _async_return(value): async def _async_append(items, value): - items.append(value) + items.append(value) \ No newline at end of file