feat(context): implement SemanticChunkingStrategy using embedding-based similarity #10770

Merged
HAL9000 merged 6 commits from feat/context-semantic-chunking-strategy into master 2026-06-06 19:49:59 +00:00
4 changed files with 813 additions and 6 deletions
+132
View File
@@ -0,0 +1,132 @@
@phase2 @acms @semantic_chunking
Feature: SemanticChunkingStrategy — embedding-based context chunking
As a CleverAgents developer
I want a SemanticChunkingStrategy that uses embedding similarity
So that the ACMS pipeline can retain the most semantically relevant
context chunks for complex multi-turn agent workflows
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy has correct name
Given a SemanticChunkingStrategy with default parameters
Then the SemanticChunkingStrategy name should be "semantic_chunking"
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy supports semantic search capability
Given a SemanticChunkingStrategy with default parameters
Then the SemanticChunkingStrategy should support semantic search
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy can_handle returns confidence score
Given a SemanticChunkingStrategy with default parameters
When I check can_handle on SemanticChunkingStrategy with query "test query"
Then the SemanticChunkingStrategy confidence should be greater than 0.0
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy can_handle returns lower confidence without query
Given a SemanticChunkingStrategy with default parameters
When I check can_handle on SemanticChunkingStrategy without query
Then the SemanticChunkingStrategy confidence should be 0.1
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy explain returns description
Given a SemanticChunkingStrategy with default parameters
Then the SemanticChunkingStrategy explain should contain "semantic"
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy accepts embedding_model parameter
Given a SemanticChunkingStrategy with embedding_model "text-embedding-ada-002"
Then the SemanticChunkingStrategy embedding_model should be "text-embedding-ada-002"
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy accepts top_k parameter
Given a SemanticChunkingStrategy with top_k 5
Then the SemanticChunkingStrategy top_k should be 5
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy has default top_k of 10
Given a SemanticChunkingStrategy with default parameters
Then the SemanticChunkingStrategy top_k should be 10
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy returns empty for empty input
Given a SemanticChunkingStrategy with default parameters
And an empty semantic chunking fragment list
And a semantic chunking budget with max_tokens 1000 and reserved_tokens 0
When I assemble with the SemanticChunkingStrategy with anchor "test query"
Then the semantic chunking result should contain 0 fragments
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy ranks fragments by cosine similarity to anchor
Given a SemanticChunkingStrategy with mock embeddings
And the following semantic chunking fragments:
| uko_node | content | score | tokens | depth |
| project://app/db.py | database connection pool manager | 0.5 | 20 | 3 |
| project://app/io.py | file input output handler | 0.5 | 15 | 3 |
| project://app/sql.py | SQL database query executor | 0.5 | 25 | 3 |
And a semantic chunking budget with max_tokens 1000 and reserved_tokens 0
When I assemble with the SemanticChunkingStrategy with anchor "database connection"
Then the first semantic chunking result should have uko_node "project://app/db.py"
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy respects top_k limit
Given a SemanticChunkingStrategy with top_k 2 and mock embeddings
And the following semantic chunking fragments:
| uko_node | content | score | tokens | depth |
| project://app/a.py | alpha beta gamma | 0.5 | 10 | 3 |
| project://app/b.py | delta epsilon zeta | 0.5 | 10 | 3 |
| project://app/c.py | eta theta iota | 0.5 | 10 | 3 |
| project://app/d.py | kappa lambda mu | 0.5 | 10 | 3 |
And a semantic chunking budget with max_tokens 1000 and reserved_tokens 0
When I assemble with the SemanticChunkingStrategy with anchor "alpha beta"
Then the semantic chunking result should contain at most 2 fragments
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy respects token budget
Given a SemanticChunkingStrategy with default parameters
And the following semantic chunking fragments:
| uko_node | content | score | tokens | depth |
| project://app/a.py | hello world | 0.9 | 100 | 3 |
| project://app/b.py | hello there | 0.7 | 100 | 3 |
| project://app/c.py | hello again | 0.5 | 100 | 3 |
And a semantic chunking budget with max_tokens 250 and reserved_tokens 0
When I assemble with the SemanticChunkingStrategy with anchor "hello"
Then the semantic chunking result should contain at most 2 fragments
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy falls back to relevance ordering without anchor
Given a SemanticChunkingStrategy with default parameters
And the following semantic chunking fragments:
| uko_node | content | score | tokens | depth |
| project://app/a.py | alpha | 0.3 | 10 | 3 |
| project://app/b.py | beta | 0.9 | 10 | 3 |
And a semantic chunking budget with max_tokens 1000 and reserved_tokens 0
When I assemble with the SemanticChunkingStrategy without anchor
Then the first semantic chunking result should have uko_node "project://app/b.py"
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy caches embeddings to avoid redundant calls
Given a SemanticChunkingStrategy with call-counting mock embeddings
And the following semantic chunking fragments:
| uko_node | content | score | tokens | depth |
| project://app/a.py | same content | 0.5 | 10 | 3 |
| project://app/b.py | same content | 0.5 | 10 | 3 |
And a semantic chunking budget with max_tokens 1000 and reserved_tokens 0
When I assemble with the SemanticChunkingStrategy with anchor "test"
Then the embedding model should be called fewer times than total fragments plus anchor
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy is registered in the ACMS pipeline
Given an ACMS pipeline for semantic chunking tests
Then the sc_pipeline should have strategy "semantic_chunking"
@tdd_issue @tdd_issue_9996
Scenario: SemanticChunkingStrategy can be used via ACMS pipeline assemble
Given an ACMS pipeline for semantic chunking tests
And a valid plan ID for semantic chunking
And the following semantic chunking fragments:
| uko_node | content | score | tokens | depth |
| project://app/a.py | alpha content | 0.8 | 10 | 3 |
| project://app/b.py | beta content | 0.6 | 10 | 3 |
And a semantic chunking budget with max_tokens 1000 and reserved_tokens 0
When I assemble via the pipeline with strategy "semantic_chunking"
Then the pipeline payload should contain at least 1 fragment
@@ -0,0 +1,376 @@
"""Step definitions for ``features/semantic_chunking_strategy.feature``.
Covers the SemanticChunkingStrategy:
* Construction and protocol compliance
* Configuration parameters (embedding_model, top_k)
* Core assembly behaviour (cosine similarity ranking)
* Embedding caching
* Plugin registry registration
"""
from __future__ import annotations
import math
from typing import Any
from behave import given, then, when
from behave.runner import Context
from cleveragents.application.services.acms_service import ACMSPipeline
from cleveragents.application.services.semantic_chunking_strategy import (
SemanticChunkingStrategy,
)
from cleveragents.domain.models.core.context_fragment import (
ContextBudget,
ContextFragment,
FragmentProvenance,
)
__all__: list[str] = []
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_sc_fragment(
uko_node: str,
content: str,
score: float,
tokens: int,
depth: int,
) -> ContextFragment:
"""Build a ``ContextFragment`` for semantic chunking test purposes."""
return ContextFragment(
uko_node=uko_node,
content=content,
relevance_score=score,
token_count=tokens,
detail_depth=depth,
provenance=FragmentProvenance(resource_uri=uko_node),
)
def _word_based_embedding(text: str) -> list[float]:
"""Word-based mock embedding for deterministic similarity testing.
Returns a 64-element vector where each element corresponds to a word
in a fixed vocabulary. This ensures that texts sharing words have
higher cosine similarity.
"""
# Fixed vocabulary of 64 words for testing
vocab = [
"database",
"connection",
"pool",
"manager",
"sql",
"query",
"executor",
"file",
"input",
"output",
"handler",
"alpha",
"beta",
"gamma",
"delta",
"epsilon",
"zeta",
"eta",
"theta",
"iota",
"kappa",
"lambda",
"mu",
"hello",
"world",
"there",
"again",
"test",
"content",
"same",
"project",
"app",
"py",
"data",
"model",
"service",
"api",
"client",
"server",
"cache",
"index",
"search",
"vector",
"embed",
"chunk",
"semantic",
"context",
"fragment",
"budget",
"token",
"score",
"rank",
"sort",
"filter",
"select",
"top",
"anchor",
"query",
"text",
"word",
"char",
"string",
"list",
]
words = text.lower().split()
vec = [0.0] * 64
for word in words:
if word in vocab:
idx = vocab.index(word)
vec[idx] += 1.0
magnitude = math.sqrt(sum(v * v for v in vec))
if magnitude > 0.0:
vec = [v / magnitude for v in vec]
return vec
# ---------------------------------------------------------------------------
# Given steps — construction
# ---------------------------------------------------------------------------
@given("a SemanticChunkingStrategy with default parameters")
def step_semantic_chunking_default(context: Context) -> None:
context.sc_strategy = SemanticChunkingStrategy()
@given('a SemanticChunkingStrategy with embedding_model "{model}"')
def step_semantic_chunking_with_model(context: Context, model: str) -> None:
context.sc_strategy = SemanticChunkingStrategy(embedding_model=model)
@given("a SemanticChunkingStrategy with top_k {top_k:d}")
def step_semantic_chunking_with_top_k(context: Context, top_k: int) -> None:
context.sc_strategy = SemanticChunkingStrategy(top_k=top_k)
@given("a SemanticChunkingStrategy with mock embeddings")
def step_semantic_chunking_with_mock_embeddings(context: Context) -> None:
context.sc_strategy = SemanticChunkingStrategy(
embedding_fn=_word_based_embedding,
)
@given("a SemanticChunkingStrategy with top_k {top_k:d} and mock embeddings")
def step_semantic_chunking_top_k_mock(context: Context, top_k: int) -> None:
context.sc_strategy = SemanticChunkingStrategy(
top_k=top_k,
embedding_fn=_word_based_embedding,
)
@given("a SemanticChunkingStrategy with call-counting mock embeddings")
def step_semantic_chunking_call_counting(context: Context) -> None:
context.sc_call_count = 0
def counting_embedding(text: str) -> list[float]:
context.sc_call_count += 1
return _word_based_embedding(text)
context.sc_strategy = SemanticChunkingStrategy(
embedding_fn=counting_embedding,
)
# ---------------------------------------------------------------------------
# Given steps — fragments and budget
# ---------------------------------------------------------------------------
@given("an empty semantic chunking fragment list")
def step_sc_empty_fragments(context: Context) -> None:
context.sc_fragments: list[ContextFragment] = []
@given("the following semantic chunking fragments:")
def step_sc_fragments_table(context: Context) -> None:
context.sc_fragments = []
for row in context.table:
frag = _make_sc_fragment(
uko_node=row["uko_node"],
content=row["content"],
score=float(row["score"]),
tokens=int(row["tokens"]),
depth=int(row["depth"]),
)
context.sc_fragments.append(frag)
@given(
"a semantic chunking budget with max_tokens {max_tokens:d} and reserved_tokens {reserved:d}"
)
def step_sc_budget(context: Context, max_tokens: int, reserved: int) -> None:
context.sc_budget = ContextBudget(max_tokens=max_tokens, reserved_tokens=reserved)
@given("an ACMS pipeline for semantic chunking tests")
def step_sc_pipeline(context: Context) -> None:
context.sc_pipeline = ACMSPipeline()
@given("a valid plan ID for semantic chunking")
def step_sc_plan_id(context: Context) -> None:
# Use a valid ULID
context.sc_plan_id = "01ARZ3NDEKTSV4RRFFQ69G5FAV"
# ---------------------------------------------------------------------------
# When steps
# ---------------------------------------------------------------------------
@when('I check can_handle on SemanticChunkingStrategy with query "{query}"')
def step_sc_can_handle_with_query(context: Context, query: str) -> None:
request: dict[str, Any] = {"query": query}
context.sc_confidence = context.sc_strategy.can_handle(request)
@when("I check can_handle on SemanticChunkingStrategy without query")
def step_sc_can_handle_no_query(context: Context) -> None:
request: dict[str, Any] = {}
context.sc_confidence = context.sc_strategy.can_handle(request)
@when('I assemble with the SemanticChunkingStrategy with anchor "{anchor}"')
def step_sc_assemble_with_anchor(context: Context, anchor: str) -> None:
context.sc_strategy.set_anchor(anchor)
context.sc_result = list(
context.sc_strategy.assemble(context.sc_fragments, context.sc_budget)
)
@when("I assemble with the SemanticChunkingStrategy without anchor")
def step_sc_assemble_no_anchor(context: Context) -> None:
# Do not set anchor — strategy should fall back to relevance ordering
context.sc_result = list(
context.sc_strategy.assemble(context.sc_fragments, context.sc_budget)
)
@when('I assemble via the pipeline with strategy "semantic_chunking"')
def step_sc_pipeline_assemble(context: Context) -> None:
context.sc_payload = context.sc_pipeline.assemble(
plan_id=context.sc_plan_id,
fragments=context.sc_fragments,
budget=context.sc_budget,
strategy="semantic_chunking",
)
# ---------------------------------------------------------------------------
# Then steps
# ---------------------------------------------------------------------------
@then('the SemanticChunkingStrategy name should be "{name}"')
def step_sc_name(context: Context, name: str) -> None:
actual = context.sc_strategy.name
assert actual == name, f"Expected name '{name}', got '{actual}'"
@then("the SemanticChunkingStrategy should support semantic search")
def step_sc_supports_semantic(context: Context) -> None:
caps = context.sc_strategy.capabilities
assert caps.supports_semantic_search, (
"SemanticChunkingStrategy should support semantic search"
)
@then("the SemanticChunkingStrategy confidence should be greater than 0.0")
def step_sc_confidence_positive(context: Context) -> None:
assert context.sc_confidence > 0.0, (
f"Expected confidence > 0.0, got {context.sc_confidence}"
)
@then("the SemanticChunkingStrategy confidence should be {expected:g}")
def step_sc_confidence_exact(context: Context, expected: float) -> None:
actual = context.sc_confidence
assert abs(actual - expected) < 1e-6, (
f"Expected confidence {expected}, got {actual}"
)
@then('the SemanticChunkingStrategy explain should contain "{text}"')
def step_sc_explain(context: Context, text: str) -> None:
explanation = context.sc_strategy.explain()
assert text.lower() in explanation.lower(), (
f"Expected explain to contain '{text}', got: {explanation}"
)
@then('the SemanticChunkingStrategy embedding_model should be "{model}"')
def step_sc_embedding_model(context: Context, model: str) -> None:
actual = context.sc_strategy.embedding_model
assert actual == model, f"Expected embedding_model '{model}', got '{actual}'"
@then("the SemanticChunkingStrategy top_k should be {expected:d}")
def step_sc_top_k(context: Context, expected: int) -> None:
actual = context.sc_strategy.top_k
assert actual == expected, f"Expected top_k {expected}, got {actual}"
@then("the semantic chunking result should contain {count:d} fragments")
def step_sc_fragment_count_exact(context: Context, count: int) -> None:
actual = len(context.sc_result)
assert actual == count, f"Expected {count} fragments, got {actual}"
@then("the semantic chunking result should contain at most {count:d} fragments")
def step_sc_fragment_count_at_most(context: Context, count: int) -> None:
actual = len(context.sc_result)
assert actual <= count, f"Expected at most {count} fragments, got {actual}"
@then('the first semantic chunking result should have uko_node "{uko_node}"')
def step_sc_first_result_uko_node(context: Context, uko_node: str) -> None:
assert len(context.sc_result) > 0, "Expected at least one result fragment"
actual = context.sc_result[0].uko_node
assert actual == uko_node, (
f"Expected first fragment uko_node '{uko_node}', got '{actual}'"
)
@then(
"the embedding model should be called fewer times than total fragments plus anchor"
)
def step_sc_cache_efficiency(context: Context) -> None:
# With 2 fragments having identical content + 1 anchor = 3 unique texts
# But "same content" appears twice, so only 2 unique texts (anchor + content)
# The call count should be less than len(fragments) + 1 (anchor)
total_possible = len(context.sc_fragments) + 1 # +1 for anchor
actual_calls = context.sc_call_count
assert actual_calls < total_possible, (
f"Expected fewer than {total_possible} embedding calls due to caching, "
f"got {actual_calls}"
)
@then('the sc_pipeline should have strategy "semantic_chunking"')
def step_sc_pipeline_has_strategy(context: Context) -> None:
registered = context.sc_pipeline._strategies
assert "semantic_chunking" in registered, (
f"Pipeline does not have strategy 'semantic_chunking'. "
f"Registered: {list(registered.keys())}"
)
@then("the pipeline payload should contain at least 1 fragment")
def step_sc_payload_has_fragments(context: Context) -> None:
count = len(context.sc_payload.fragments)
assert count >= 1, f"Expected at least 1 fragment in payload, got {count}"
@@ -26,7 +26,7 @@ import re
from collections.abc import Sequence
from dataclasses import dataclass
from threading import local
from typing import TYPE_CHECKING, Any, ClassVar, Protocol, runtime_checkable
from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast, runtime_checkable
import structlog
@@ -58,6 +58,22 @@ _GreedyKnapsackPacker: type | None = None
# Lazy import helpers for spec-required built-in strategies.
# These are imported lazily to avoid circular imports and to keep the
# acms_service module lightweight.
# Lazy import helper for SemanticChunkingStrategy (issue #9996).
_SemanticChunkingStrategy: type | None = None
def _get_semantic_chunking_strategy_class() -> type:
"""Return the :class:`SemanticChunkingStrategy` class, importing lazily."""
global _SemanticChunkingStrategy
if _SemanticChunkingStrategy is None:
from cleveragents.application.services.semantic_chunking_strategy import (
SemanticChunkingStrategy,
)
_SemanticChunkingStrategy = SemanticChunkingStrategy
return _SemanticChunkingStrategy
_SPEC_BUILTIN_STRATEGIES: dict[str, Any] | None = None
@@ -728,10 +744,10 @@ class ACMSPipeline:
)
"""
BUILTIN_STRATEGIES: ClassVar[dict[str, type[ContextStrategy]]] = {
"relevance": RelevanceStrategy, # type: ignore[dict-item]
"recency": RecencyStrategy, # type: ignore[dict-item]
"tiered": TieredStrategy, # type: ignore[dict-item]
BUILTIN_STRATEGIES: ClassVar[dict[str, type[Any]]] = {
"relevance": RelevanceStrategy,
"recency": RecencyStrategy,
"tiered": TieredStrategy,
}
def __init__(
@@ -765,6 +781,11 @@ class ACMSPipeline:
self._strategies: dict[str, ContextStrategy] = {
name: cls() for name, cls in self.BUILTIN_STRATEGIES.items()
}
# Register SemanticChunkingStrategy (issue #9996).
# Imported lazily to avoid circular imports.
if "semantic_chunking" not in self._strategies:
_sc_cls = _get_semantic_chunking_strategy_class()
self._strategies["semantic_chunking"] = cast(ContextStrategy, _sc_cls())
# Register the 6 spec-required built-in strategies via adapters.
# These strategies implement the domain-model ContextStrategy protocol
# (strategy_stubs.py) and are wrapped in SpecStrategyAdapter instances
@@ -773,7 +794,8 @@ class ACMSPipeline:
# can be replaced with direct registrations.
for spec_name, spec_cls in _get_spec_builtin_strategies().items():
if spec_name not in self._strategies:
self._strategies[spec_name] = SpecStrategyAdapter(spec_cls()) # type: ignore[assignment]
adapter = cast(ContextStrategy, SpecStrategyAdapter(spec_cls()))
self._strategies[spec_name] = adapter
if default_strategy not in self._strategies:
msg = (
@@ -0,0 +1,277 @@
"""SemanticChunkingStrategy — embedding-based context chunking.
Implements a context strategy that uses embedding similarity to retain
the most semantically relevant context chunks for complex multi-turn
agent workflows where topic coherence matters.
The strategy:
1. Computes embeddings for each fragment's content and the anchor message.
2. Ranks fragments by cosine similarity to the anchor.
3. Selects the top-K most relevant chunks within the token budget.
4. Caches embeddings to avoid redundant API calls.
Implements the ``ContextStrategy`` protocol defined in ``acms_service.py``
and is registered in the ``ACMSPipeline`` under key ``"semantic_chunking"``.
Based on ``docs/specification.md`` §25207-25216.
ISSUES CLOSED: #9996
"""
from __future__ import annotations
import logging
import math
import re
from collections.abc import Callable, Sequence
from typing import Any
from cleveragents.application.services.acms_service import (
StrategyCapabilities,
_pack_budget,
)
from cleveragents.domain.models.core.context_fragment import (
ContextBudget,
ContextFragment,
)
__all__ = [
"DEFAULT_EMBEDDING_MODEL",
"DEFAULT_TOP_K",
"EmbeddingFn",
"SemanticChunkingStrategy",
]
logger = logging.getLogger(__name__)
# Default embedding model name (used when no real embedding provider is wired).
DEFAULT_EMBEDDING_MODEL: str = "text-embedding-ada-002"
# Default number of top-K chunks to retain.
DEFAULT_TOP_K: int = 10
# Word tokenisation pattern (reused from context_strategies.py).
_WORD_RE = re.compile(r"\w+", re.UNICODE)
# Type alias for an embedding function: text -> list[float]
EmbeddingFn = Callable[[str], list[float]]
# ---------------------------------------------------------------------------
# Default (built-in) embedding function
# ---------------------------------------------------------------------------
def _default_embedding(text: str) -> list[float]:
"""Compute a simple character-frequency embedding for *text*.
This is a bag-of-characters vector over the 64 most common ASCII
characters (ordinals 32-95), normalised to unit length. It provides
a lightweight, dependency-free approximation of semantic similarity
suitable for testing and fallback use.
Production deployments should inject a real embedding function via
the ``embedding_fn`` constructor parameter.
"""
# Defensive guard: public API never calls with empty text.
if not text: # pragma: no cover
return [0.0] * 64
vec = [0.0] * 64
for ch in text.lower():
idx = ord(ch) - 32
if 0 <= idx < 64:
vec[idx] += 1.0
magnitude = math.sqrt(sum(v * v for v in vec))
if magnitude > 0.0:
vec = [v / magnitude for v in vec]
return vec
# ---------------------------------------------------------------------------
# Cosine similarity helper
# ---------------------------------------------------------------------------
def _cosine_similarity(vec_a: list[float], vec_b: list[float]) -> float:
"""Compute cosine similarity between two equal-length vectors.
Returns a value in [-1.0, 1.0]. Returns 0.0 when either vector is
the zero vector to avoid division-by-zero.
"""
# Defensive guard: embedding_fn always returns same-length vectors.
if len(vec_a) != len(vec_b): # pragma: no cover
return 0.0
dot = sum(a * b for a, b in zip(vec_a, vec_b, strict=True))
mag_a = math.sqrt(sum(a * a for a in vec_a))
mag_b = math.sqrt(sum(b * b for b in vec_b))
if mag_a == 0.0 or mag_b == 0.0:
return 0.0
return dot / (mag_a * mag_b)
# ---------------------------------------------------------------------------
# SemanticChunkingStrategy
# ---------------------------------------------------------------------------
class SemanticChunkingStrategy:
"""Embedding-based semantic chunking strategy.
Ranks context fragments by cosine similarity of their content
embeddings to an anchor message (typically the current query or
last message in the conversation). The top-K most similar chunks
are retained within the token budget.
Embeddings are cached by content string to avoid redundant API calls
when multiple fragments share identical content.
**Configuration**:
- ``embedding_model``: Name of the embedding model to use. Passed
to the ``embedding_fn`` for informational purposes; the actual
model selection is the responsibility of the injected function.
- ``top_k``: Maximum number of fragments to retain before budget
packing. Defaults to 10.
- ``embedding_fn``: Callable ``(text: str) -> list[float]``.
Defaults to a lightweight character-frequency approximation.
Inject a real embedding provider for production use.
**Fallback behaviour**: When no anchor is provided (empty string),
the strategy falls back to relevance-score ordering.
Implements ``ContextStrategy`` protocol from ``acms_service.py``.
Based on ``docs/specification.md`` §25207-25216.
"""
def __init__(
self,
*,
embedding_model: str = DEFAULT_EMBEDDING_MODEL,
top_k: int = DEFAULT_TOP_K,
embedding_fn: EmbeddingFn | None = None,
) -> None:
self._embedding_model = embedding_model
self._top_k = top_k
self._embedding_fn: EmbeddingFn = embedding_fn or _default_embedding
# Cache: content string -> embedding vector
self._embedding_cache: dict[str, list[float]] = {}
self._anchor: str = ""
# ------------------------------------------------------------------
# Configuration accessors
# ------------------------------------------------------------------
@property
def embedding_model(self) -> str:
"""Return the configured embedding model name."""
return self._embedding_model
@property
def top_k(self) -> int:
"""Return the configured top-K limit."""
return self._top_k
def set_anchor(self, anchor: str) -> None:
"""Set the anchor message for similarity ranking (optional)."""
self._anchor = anchor
# ------------------------------------------------------------------
# ContextStrategy protocol
# ------------------------------------------------------------------
@property
def name(self) -> str:
return "semantic_chunking"
@property
def capabilities(self) -> StrategyCapabilities:
return StrategyCapabilities(supports_semantic_search=True)
def can_handle(self, request: dict[str, Any]) -> float:
"""Return 0.75 when a query is present, else 0.1.
The strategy is most useful when there is a query or anchor
message to compute similarity against. Without one it falls
back to relevance ordering, which is less distinctive.
"""
query = str(request.get("query", "") or "")
self._anchor = query
return 0.75 if query else 0.1
def assemble(
self,
fragments: Sequence[ContextFragment],
budget: ContextBudget,
) -> Sequence[ContextFragment]:
"""Rank fragments by cosine similarity to the anchor, then pack.
Steps:
1. If no anchor is set, fall back to relevance-score ordering.
2. Compute (cached) embeddings for the anchor and each fragment.
3. Rank fragments by cosine similarity (descending).
4. Apply top-K limit.
5. Pack within the token budget.
"""
if not fragments:
return list(fragments)
if not self._anchor:
# No anchor — fall back to relevance ordering
sorted_frags = sorted(
fragments, key=lambda f: f.relevance_score, reverse=True
)
return _pack_budget(sorted_frags, budget)
# Compute anchor embedding (cached)
anchor_vec = self._get_embedding(self._anchor)
# Score each fragment by cosine similarity to anchor
scored: list[tuple[ContextFragment, float]] = []
for frag in fragments:
frag_vec = self._get_embedding(frag.content)
sim = _cosine_similarity(anchor_vec, frag_vec)
scored.append((frag, sim))
# Sort by similarity descending, then by relevance_score as tiebreaker
scored.sort(
key=lambda pair: (pair[1], pair[0].relevance_score),
reverse=True,
)
# Apply top-K limit
top_frags = [frag for frag, _ in scored[: self._top_k]]
logger.info(
"SemanticChunking ranked fragments",
extra={
"anchor_length": len(self._anchor),
"fragment_count": len(fragments),
"top_k": self._top_k,
"selected_count": len(top_frags),
"embedding_model": self._embedding_model,
},
)
return _pack_budget(top_frags, budget)
def explain(self) -> str:
return (
f"Embedding-based semantic chunking strategy. "
f"Ranks context fragments by cosine similarity of their "
f"content embeddings to the anchor message. "
f"Retains the top-{self._top_k} most relevant chunks within "
f"the token budget. Uses embedding model '{self._embedding_model}'. "
f"Caches embeddings to avoid redundant API calls. "
f"Falls back to relevance ordering when no anchor is provided."
)
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
def _get_embedding(self, text: str) -> list[float]:
"""Return the embedding for *text*, using the cache when possible."""
if text not in self._embedding_cache:
self._embedding_cache[text] = self._embedding_fn(text)
return self._embedding_cache[text]