feat(context): implement SemanticChunkingStrategy using embedding-based similarity #10770
@@ -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]
|
||||
Reference in New Issue
Block a user