Python: Add a HarnessAgent with available features and sample (#6041)

* Add a HarnessAgent with available features and sample * Fix formatting * Address PR comments and fix mypy error * Add web search support to HarnessAgent * Fix build warning * Apply suggestions from code review Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com> * Address PR comments * Address PR comments * Address further PR comments. * Fix markdown broken link --------- Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2026-05-27 14:54:00 +01:00
parent d5c07f2623
commit ef86fb51d5
11 changed files with 1262 additions and 5 deletions
@@ -19,6 +19,7 @@ from agent_framework import (
    ChatResponse,
    CompactionProvider,
    Content,
+    ContextWindowCompactionStrategy,
    Message,
    SelectiveToolCallCompactionStrategy,
    SlidingWindowStrategy,
@@ -952,3 +953,159 @@ async def test_in_memory_history_provider_default_loads_all() -> None:

    loaded = await provider.get_messages(session_id="test", state=state)
    assert len(loaded) == 3
+
+
+# --- ContextWindowCompactionStrategy tests ---
+
+
+async def test_context_window_strategy_noop_under_threshold() -> None:
+    """No compaction when total tokens are below 50% of input budget."""
+    # input_budget = 1000 - 200 = 800; tool eviction threshold = 50% = 400 tokens
+    # CharacterEstimatorTokenizer: 4 chars/token
+    # Each short message ~4-5 tokens, total well under 400
+    messages = [
+        Message(role="system", contents=["sys"]),
+        Message(role="user", contents=["hello"]),
+        Message(role="assistant", contents=["hi"]),
+    ]
+    strategy = ContextWindowCompactionStrategy(
+        max_context_window_tokens=1000,
+        max_output_tokens=200,
+    )
+
+    changed = await strategy(messages)
+
+    assert changed is False
+    assert len(included_messages(messages)) == 3
+
+
+async def test_context_window_strategy_tool_eviction_triggers_at_threshold() -> None:
+    """Tool eviction fires when tokens exceed 50% but truncation does not."""
+    # input_budget = 20000 - 200 = 19800
+    # tool eviction at 50% = 9900 tokens; truncation at 80% = 15840 tokens
+    # CharacterEstimatorTokenizer: 4 chars/token
+    # Each tool result: "x" * 8000 = 8000 chars = 2000 tokens
+    # 5 groups * ~2000 = ~10000+ tokens (exceeds 9900, under 15840)
+    # Tool eviction collapses older groups; truncation threshold not reached.
+    messages = [
+        Message(role="system", contents=["system prompt"]),
+        Message(role="user", contents=["u1"]),
+        _assistant_function_call("c1"),
+        _tool_result("c1", "x" * 8000),
+        Message(role="user", contents=["u2"]),
+        _assistant_function_call("c2"),
+        _tool_result("c2", "x" * 8000),
+        Message(role="user", contents=["u3"]),
+        _assistant_function_call("c3"),
+        _tool_result("c3", "x" * 8000),
+        Message(role="user", contents=["u4"]),
+        _assistant_function_call("c4"),
+        _tool_result("c4", "x" * 8000),
+        Message(role="user", contents=["u5"]),
+        _assistant_function_call("c5"),
+        _tool_result("c5", "x" * 8000),
+    ]
+    strategy = ContextWindowCompactionStrategy(
+        max_context_window_tokens=20000,
+        max_output_tokens=200,
+        keep_last_tool_call_groups=2,
+    )
+
+    changed = await strategy(messages)
+
+    assert changed is True
+    projected = included_messages(messages)
+    # Verify that tool results were compacted (summary messages present).
+    summary_msgs = [m for m in projected if m.text and "[Tool results:" in m.text]
+    assert len(summary_msgs) > 0
+    # Verify that the truncation phase did NOT fire — no messages excluded with "truncation" reason.
+    from agent_framework._compaction import EXCLUDE_REASON_KEY
+
+    truncation_excluded = [m for m in messages if m.additional_properties.get(EXCLUDE_REASON_KEY) == "truncation"]
+    assert len(truncation_excluded) == 0
+
+
+async def test_context_window_strategy_truncation_triggers_above_80_pct() -> None:
+    """Truncation fires when tokens exceed 80% of input budget."""
+    # input_budget = 1000 - 100 = 900
+    # tool eviction at 50% = 450 tokens; truncation at 80% = 720 tokens
+    # We'll create messages with no tool calls (so tool eviction does nothing)
+    # but exceeding 720 tokens total (>2880 chars)
+    messages = [
+        Message(role="system", contents=["sys"]),
+        Message(role="user", contents=["u1 " * 400]),  # ~1200 chars = 300 tokens
+        Message(role="assistant", contents=["a1 " * 400]),  # ~1200 chars = 300 tokens
+        Message(role="user", contents=["u2 " * 400]),  # ~1200 chars = 300 tokens
+        Message(role="assistant", contents=["a2 " * 400]),  # ~1200 chars = 300 tokens
+    ]
+    strategy = ContextWindowCompactionStrategy(
+        max_context_window_tokens=1000,
+        max_output_tokens=100,
+    )
+
+    changed = await strategy(messages)
+
+    assert changed is True
+    projected = included_messages(messages)
+    # System message should always be preserved
+    assert projected[0].role == "system"
+    # Some messages should have been excluded
+    assert len(projected) < 5
+
+
+async def test_context_window_strategy_keep_last_tool_call_groups_respected() -> None:
+    """The keep_last_tool_call_groups parameter controls how many groups are retained."""
+    # Create enough tokens to trigger tool eviction (>50% of input budget)
+    # input_budget = 1000 - 100 = 900; threshold = 450 tokens
+    messages = [
+        Message(role="system", contents=["sys"]),
+        Message(role="user", contents=["u1"]),
+        _assistant_function_call("c1"),
+        _tool_result("c1", "r1 " * 200),
+        Message(role="user", contents=["u2"]),
+        _assistant_function_call("c2"),
+        _tool_result("c2", "r2 " * 200),
+        Message(role="user", contents=["u3"]),
+        _assistant_function_call("c3"),
+        _tool_result("c3", "r3 " * 200),
+    ]
+    # keep_last_tool_call_groups=1: only the last group (c3) should be kept verbatim
+    strategy = ContextWindowCompactionStrategy(
+        max_context_window_tokens=1000,
+        max_output_tokens=100,
+        keep_last_tool_call_groups=1,
+    )
+
+    changed = await strategy(messages)
+
+    assert changed is True
+    projected = included_messages(messages)
+    # The last tool call group (c3) should be in the projected messages
+    has_c3 = any(
+        c.call_id == "c3" for m in projected for c in m.contents if c.type in ("function_call", "function_result")
+    )
+    assert has_c3
+
+
+def test_context_window_strategy_validates_thresholds() -> None:
+    """Invalid threshold combinations raise ValueError."""
+    import pytest
+
+    with pytest.raises(ValueError, match="max_context_window_tokens must be positive"):
+        ContextWindowCompactionStrategy(max_context_window_tokens=0, max_output_tokens=0)
+
+    with pytest.raises(ValueError, match="max_output_tokens must be >= 0"):
+        ContextWindowCompactionStrategy(max_context_window_tokens=1000, max_output_tokens=1000)
+
+    with pytest.raises(ValueError, match="tool_eviction_threshold must be in"):
+        ContextWindowCompactionStrategy(
+            max_context_window_tokens=1000, max_output_tokens=100, tool_eviction_threshold=0.0
+        )
+
+    with pytest.raises(ValueError, match="truncation_threshold must be >= tool_eviction_threshold"):
+        ContextWindowCompactionStrategy(
+            max_context_window_tokens=1000,
+            max_output_tokens=100,
+            tool_eviction_threshold=0.8,
+            truncation_threshold=0.5,
+        )
@@ -0,0 +1,396 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Mapping
+from typing import Any
+
+import pytest
+
+from agent_framework import (
+    AgentSession,
+    ChatResponse,
+    CompactionProvider,
+    InMemoryHistoryProvider,
+    Message,
+    SkillsProvider,
+    TodoProvider,
+    create_harness_agent,
+)
+from agent_framework._harness._agent import DEFAULT_HARNESS_INSTRUCTIONS, _assemble_instructions
+from agent_framework._harness._mode import AgentModeProvider
+from agent_framework._sessions import ContextProvider
+
+
+class _FakeChatClient:
+    """Minimal chat client stub for testing assembly."""
+
+    model = "test-model"
+
+    async def get_response(
+        self,
+        *,
+        messages: list[Message],
+        options: Mapping[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> ChatResponse:
+        return ChatResponse(messages=[Message(role="assistant", contents=["Hello"])])
+
+    async def get_streaming_response(
+        self,
+        *,
+        messages: list[Message],
+        options: Mapping[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[Any]:
+        yield Message(role="assistant", contents=["Hello"])  # pragma: no cover
+
+
+# --- Assembly Tests ---
+
+
+def test_create_harness_agent_with_defaults() -> None:
+    """create_harness_agent should assemble successfully with default options."""
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+    )
+    assert agent.id is not None
+
+
+def test_create_harness_agent_includes_all_default_providers() -> None:
+    """Default assembly should include history, compaction, todo, mode (no skills by default)."""
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+    )
+    providers = agent.context_providers
+    provider_types = [type(p) for p in providers]
+
+    assert InMemoryHistoryProvider in provider_types
+    assert CompactionProvider in provider_types
+    assert TodoProvider in provider_types
+    assert AgentModeProvider in provider_types
+    assert SkillsProvider not in provider_types
+
+
+def test_create_harness_agent_disable_todo() -> None:
+    """disable_todo=True should exclude TodoProvider."""
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+        disable_todo=True,
+    )
+    provider_types = [type(p) for p in agent.context_providers]
+    assert TodoProvider not in provider_types
+
+
+def test_create_harness_agent_disable_mode() -> None:
+    """disable_mode=True should exclude AgentModeProvider."""
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+        disable_mode=True,
+    )
+    provider_types = [type(p) for p in agent.context_providers]
+    assert AgentModeProvider not in provider_types
+
+
+def test_create_harness_agent_disable_memory() -> None:
+    """disable_memory=True should exclude MemoryContextProvider even when memory_store is provided."""
+    from agent_framework import MemoryContextProvider
+    from agent_framework._harness._memory import MemoryStore
+
+    class _FakeMemoryStore(MemoryStore):
+        def list_topics(self, session, *, source_id):
+            return []
+
+        def get_topic(self, session, *, source_id, topic):
+            raise NotImplementedError
+
+        def write_topic(self, session, record, *, source_id):
+            pass
+
+        def delete_topic(self, session, *, source_id, topic):
+            pass
+
+        def get_index_text(self, session, *, source_id):
+            return ""
+
+        def get_transcripts_directory(self, session, *, source_id):
+            return ""
+
+        def read_state(self, session, *, source_id):
+            return {}
+
+        def rebuild_index(self, session, *, source_id):
+            pass
+
+        def search_transcripts(self, session, *, source_id, query):
+            return []
+
+        def write_state(self, session, state, *, source_id):
+            pass
+
+    # With memory_store provided and disable_memory=False, MemoryContextProvider should be present.
+    agent_with_memory = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+        memory_store=_FakeMemoryStore(),
+    )
+    provider_types = [type(p) for p in agent_with_memory.context_providers]
+    assert MemoryContextProvider in provider_types
+
+    # With memory_store provided and disable_memory=True, MemoryContextProvider should be absent.
+    agent_disabled = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+        memory_store=_FakeMemoryStore(),
+        disable_memory=True,
+    )
+    provider_types = [type(p) for p in agent_disabled.context_providers]
+    assert MemoryContextProvider not in provider_types
+
+
+def test_create_harness_agent_skills_paths_adds_provider() -> None:
+    """skills_paths should add a SkillsProvider."""
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+        skills_paths=["./test-skills"],
+    )
+    provider_types = [type(p) for p in agent.context_providers]
+    assert SkillsProvider in provider_types
+
+
+def test_create_harness_agent_disable_compaction() -> None:
+    """disable_compaction=True should exclude CompactionProvider."""
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+        disable_compaction=True,
+    )
+    provider_types = [type(p) for p in agent.context_providers]
+    assert CompactionProvider not in provider_types
+
+
+def test_create_harness_agent_returns_full_agent() -> None:
+    """Factory should return an Agent instance (with telemetry)."""
+    from agent_framework._agents import Agent as FullAgent
+
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+    )
+    assert isinstance(agent, FullAgent)
+
+
+# --- Validation Tests ---
+
+
+def test_create_harness_agent_rejects_invalid_context_tokens() -> None:
+    """max_context_window_tokens must be positive."""
+    with pytest.raises(ValueError, match="max_context_window_tokens must be positive"):
+        create_harness_agent(
+            client=_FakeChatClient(),  # type: ignore[arg-type]
+            max_context_window_tokens=0,
+            max_output_tokens=100,
+        )
+
+
+def test_create_harness_agent_rejects_negative_output_tokens() -> None:
+    """max_output_tokens must be non-negative."""
+    with pytest.raises(ValueError, match="max_output_tokens must be non-negative"):
+        create_harness_agent(
+            client=_FakeChatClient(),  # type: ignore[arg-type]
+            max_context_window_tokens=1000,
+            max_output_tokens=-1,
+        )
+
+
+def test_create_harness_agent_rejects_output_gte_context() -> None:
+    """max_output_tokens must be less than max_context_window_tokens."""
+    with pytest.raises(ValueError, match="max_output_tokens must be less than"):
+        create_harness_agent(
+            client=_FakeChatClient(),  # type: ignore[arg-type]
+            max_context_window_tokens=1000,
+            max_output_tokens=1000,
+        )
+
+
+# --- Instructions Tests ---
+
+
+def test_default_instructions() -> None:
+    """None args should produce default harness instructions."""
+    result = _assemble_instructions(None, None)
+    assert result == DEFAULT_HARNESS_INSTRUCTIONS.strip()
+
+
+def test_custom_agent_instructions_appended() -> None:
+    """Agent instructions should be appended after harness instructions."""
+    result = _assemble_instructions(None, "Focus on code review.")
+    assert DEFAULT_HARNESS_INSTRUCTIONS in result  # type: ignore[operator]
+    assert "Focus on code review." in result  # type: ignore[operator]
+
+
+def test_empty_harness_instructions_uses_agent_only() -> None:
+    """Empty harness_instructions should return agent instructions only."""
+    result = _assemble_instructions("", "Custom only.")
+    assert result == "Custom only."
+
+
+# --- Identity Tests ---
+
+
+def test_create_harness_agent_custom_identity() -> None:
+    """Custom id, name, description should propagate."""
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+        id="my-agent-id",
+        name="my-agent",
+        description="A test agent",
+    )
+    assert agent.id == "my-agent-id"
+    assert agent.name == "my-agent"
+    assert agent.description == "A test agent"
+
+
+# --- Session Tests ---
+
+
+def test_create_harness_agent_create_session() -> None:
+    """create_session should return an AgentSession."""
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+    )
+    session = agent.create_session()
+    assert isinstance(session, AgentSession)
+
+
+def test_create_harness_agent_create_session_with_id() -> None:
+    """create_session should accept a custom session_id."""
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+    )
+    session = agent.create_session(session_id="custom-id")
+    assert session.session_id == "custom-id"
+
+
+async def test_create_harness_agent_run_returns_response() -> None:
+    """agent.run() should return a response."""
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+    )
+    session = agent.create_session()
+    response = await agent.run("hello", session=session)
+    assert response.messages
+    assert response.messages[-1].role == "assistant"
+
+
+# --- Protocol Tests ---
+
+
+def test_create_harness_agent_satisfies_protocol() -> None:
+    """Returned agent should satisfy SupportsAgentRun protocol."""
+    from agent_framework import SupportsAgentRun
+
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+    )
+    assert isinstance(agent, SupportsAgentRun)
+
+
+# --- Additional providers ---
+
+
+def test_create_harness_agent_extra_context_providers() -> None:
+    """Additional context_providers should be appended."""
+
+    class _CustomProvider(ContextProvider):
+        pass
+
+    custom = _CustomProvider("custom")
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+        context_providers=[custom],
+    )
+    assert custom in agent.context_providers
+
+
+# --- Web Search Tool Tests ---
+
+
+class _FakeWebSearchClient(_FakeChatClient):
+    """Fake client that supports web search tool."""
+
+    def get_web_search_tool(self, **kwargs: Any) -> str:
+        return "web_search_tool_instance"
+
+
+def test_create_harness_agent_auto_adds_web_search_tool() -> None:
+    """Web search tool should be auto-added when client supports it."""
+    agent = create_harness_agent(
+        client=_FakeWebSearchClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+    )
+    tools = agent.default_options.get("tools", [])
+    assert "web_search_tool_instance" in tools
+
+
+def test_create_harness_agent_disable_web_search() -> None:
+    """disable_web_search=True should skip auto-adding the web search tool."""
+    agent = create_harness_agent(
+        client=_FakeWebSearchClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+        disable_web_search=True,
+    )
+    tools = agent.default_options.get("tools", [])
+    assert "web_search_tool_instance" not in tools
+
+
+def test_create_harness_agent_no_web_search_when_unsupported() -> None:
+    """Web search tool should NOT be added when client does not support it."""
+    agent = create_harness_agent(
+        client=_FakeChatClient(),  # type: ignore[arg-type]
+        max_context_window_tokens=128_000,
+        max_output_tokens=16_384,
+    )
+    tools = agent.default_options.get("tools", [])
+    assert "web_search_tool_instance" not in tools
+
+
+def test_create_harness_agent_logs_warning_when_no_web_search(caplog: pytest.LogCaptureFixture) -> None:
+    """A warning should be logged when client doesn't support web search."""
+    import logging
+
+    with caplog.at_level(logging.WARNING, logger="agent_framework._harness._agent"):
+        create_harness_agent(
+            client=_FakeChatClient(),  # type: ignore[arg-type]
+            max_context_window_tokens=128_000,
+            max_output_tokens=16_384,
+        )
+    assert any("SupportsWebSearchTool" in msg for msg in caplog.messages)