pyagent-compress 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,18 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *$py.class
4
+ *.egg-info/
5
+ dist/
6
+ build/
7
+ .eggs/
8
+ *.egg
9
+ .venv/
10
+ venv/
11
+ .mypy_cache/
12
+ .ruff_cache/
13
+ .pytest_cache/
14
+ .coverage
15
+ htmlcov/
16
+ site/
17
+ .env
18
+ *.log
@@ -0,0 +1,53 @@
1
+ Metadata-Version: 2.4
2
+ Name: pyagent-compress
3
+ Version: 0.1.0
4
+ Summary: Inter-agent message compression and token budget management for multi-agent LLM systems
5
+ License: MIT
6
+ Keywords: LLM,agents,compression,efficiency,tokens
7
+ Classifier: Development Status :: 3 - Alpha
8
+ Classifier: Intended Audience :: Developers
9
+ Classifier: License :: OSI Approved :: MIT License
10
+ Classifier: Programming Language :: Python :: 3.11
11
+ Classifier: Programming Language :: Python :: 3.12
12
+ Classifier: Programming Language :: Python :: 3.13
13
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
14
+ Classifier: Typing :: Typed
15
+ Requires-Python: >=3.11
16
+ Requires-Dist: pyagent-patterns>=0.1.0
17
+ Provides-Extra: dev
18
+ Requires-Dist: mypy>=1.10; extra == 'dev'
19
+ Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
20
+ Requires-Dist: pytest>=8.0; extra == 'dev'
21
+ Requires-Dist: ruff>=0.5; extra == 'dev'
22
+ Description-Content-Type: text/markdown
23
+
24
+ # pyagent-compress
25
+
26
+ **Inter-agent message compression and token budget management** for multi-agent LLM systems. Reduce token costs without losing key information.
27
+
28
+ ## Install
29
+
30
+ ```bash
31
+ pip install pyagent-compress
32
+ ```
33
+
34
+ ## Components
35
+
36
+ - **MessageCompressor** — Reduce message size by removing filler and ranking sentences
37
+ - **AgentPruner** — Detect and remove non-contributing agents
38
+ - **InteractionPruner** — Detect consensus and prune redundant rounds
39
+ - **TokenBudget** — Enforce per-agent and per-workflow token limits
40
+ - **CompressMiddleware** — Auto-compress agent outputs
41
+
42
+ ## Quick Example
43
+
44
+ ```python
45
+ from pyagent_compress import MessageCompressor, TokenBudget
46
+
47
+ compressor = MessageCompressor(target_ratio=0.5)
48
+ result = compressor.compress("Let me think about this... Basically, revenue grew 15%.")
49
+ print(f"Savings: {result.savings_pct:.0%}")
50
+
51
+ budget = TokenBudget(workflow_limit=50_000, per_agent_limit=10_000)
52
+ budget.consume("agent_a", 3000)
53
+ ```
@@ -0,0 +1,30 @@
1
+ # pyagent-compress
2
+
3
+ **Inter-agent message compression and token budget management** for multi-agent LLM systems. Reduce token costs without losing key information.
4
+
5
+ ## Install
6
+
7
+ ```bash
8
+ pip install pyagent-compress
9
+ ```
10
+
11
+ ## Components
12
+
13
+ - **MessageCompressor** — Reduce message size by removing filler and ranking sentences
14
+ - **AgentPruner** — Detect and remove non-contributing agents
15
+ - **InteractionPruner** — Detect consensus and prune redundant rounds
16
+ - **TokenBudget** — Enforce per-agent and per-workflow token limits
17
+ - **CompressMiddleware** — Auto-compress agent outputs
18
+
19
+ ## Quick Example
20
+
21
+ ```python
22
+ from pyagent_compress import MessageCompressor, TokenBudget
23
+
24
+ compressor = MessageCompressor(target_ratio=0.5)
25
+ result = compressor.compress("Let me think about this... Basically, revenue grew 15%.")
26
+ print(f"Savings: {result.savings_pct:.0%}")
27
+
28
+ budget = TokenBudget(workflow_limit=50_000, per_agent_limit=10_000)
29
+ budget.consume("agent_a", 3000)
30
+ ```
@@ -0,0 +1,29 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "pyagent-compress"
7
+ version = "0.1.0"
8
+ description = "Inter-agent message compression and token budget management for multi-agent LLM systems"
9
+ readme = "README.md"
10
+ requires-python = ">=3.11"
11
+ license = {text = "MIT"}
12
+ keywords = ["agents", "compression", "LLM", "tokens", "efficiency"]
13
+ classifiers = [
14
+ "Development Status :: 3 - Alpha",
15
+ "Intended Audience :: Developers",
16
+ "License :: OSI Approved :: MIT License",
17
+ "Programming Language :: Python :: 3.11",
18
+ "Programming Language :: Python :: 3.12",
19
+ "Programming Language :: Python :: 3.13",
20
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
21
+ "Typing :: Typed",
22
+ ]
23
+ dependencies = ["pyagent-patterns>=0.1.0"]
24
+
25
+ [project.optional-dependencies]
26
+ dev = ["pytest>=8.0", "pytest-asyncio>=0.23", "ruff>=0.5", "mypy>=1.10"]
27
+
28
+ [tool.hatch.build.targets.wheel]
29
+ packages = ["src/pyagent_compress"]
@@ -0,0 +1,15 @@
1
+ """PyAgent Compress — inter-agent message compression and token budget management."""
2
+
3
+ from pyagent_compress.budget import TokenBudget
4
+ from pyagent_compress.compressor import MessageCompressor
5
+ from pyagent_compress.middleware import CompressMiddleware
6
+ from pyagent_compress.pruner import AgentPruner, InteractionPruner
7
+
8
+ __all__ = [
9
+ "MessageCompressor",
10
+ "AgentPruner",
11
+ "InteractionPruner",
12
+ "TokenBudget",
13
+ "CompressMiddleware",
14
+ ]
15
+ __version__ = "0.1.0"
@@ -0,0 +1,118 @@
1
+ """TokenBudget: enforce per-agent and per-workflow token limits.
2
+
3
+ Tracks token consumption across a workflow and raises BudgetExceeded
4
+ when limits are hit, enabling graceful degradation.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ from dataclasses import dataclass, field
10
+
11
+
12
+ class BudgetExceeded(Exception):
13
+ """Raised when a token budget is exceeded."""
14
+
15
+ def __init__(self, agent: str, used: int, limit: int) -> None:
16
+ self.agent = agent
17
+ self.used = used
18
+ self.limit = limit
19
+ super().__init__(f"Agent '{agent}' exceeded budget: {used}/{limit} tokens")
20
+
21
+
22
+ @dataclass
23
+ class AgentBudget:
24
+ """Budget tracking for a single agent."""
25
+
26
+ name: str
27
+ limit: int
28
+ used: int = 0
29
+
30
+ @property
31
+ def remaining(self) -> int:
32
+ return max(0, self.limit - self.used)
33
+
34
+ @property
35
+ def utilization(self) -> float:
36
+ return self.used / self.limit if self.limit > 0 else 0.0
37
+
38
+ def consume(self, tokens: int, strict: bool = True) -> None:
39
+ """Consume tokens from budget.
40
+
41
+ Args:
42
+ tokens: Number of tokens consumed.
43
+ strict: If True, raises BudgetExceeded. If False, just tracks.
44
+ """
45
+ self.used += tokens
46
+ if strict and self.used > self.limit:
47
+ raise BudgetExceeded(self.name, self.used, self.limit)
48
+
49
+
50
+ @dataclass
51
+ class TokenBudget:
52
+ """Manage token budgets across a multi-agent workflow.
53
+
54
+ Args:
55
+ workflow_limit: Total token limit across all agents.
56
+ per_agent_limit: Default per-agent token limit.
57
+ strict: If True, raises BudgetExceeded on limit violations.
58
+ """
59
+
60
+ workflow_limit: int = 100_000
61
+ per_agent_limit: int = 20_000
62
+ strict: bool = True
63
+ _agents: dict[str, AgentBudget] = field(default_factory=dict)
64
+ _total_used: int = 0
65
+
66
+ def register_agent(self, name: str, limit: int | None = None) -> None:
67
+ """Register an agent with an optional custom limit."""
68
+ self._agents[name] = AgentBudget(name=name, limit=limit or self.per_agent_limit)
69
+
70
+ def consume(self, agent_name: str, tokens: int) -> None:
71
+ """Record token consumption for an agent.
72
+
73
+ Args:
74
+ agent_name: The consuming agent's name.
75
+ tokens: Number of tokens consumed.
76
+ """
77
+ if agent_name not in self._agents:
78
+ self.register_agent(agent_name)
79
+
80
+ self._agents[agent_name].consume(tokens, strict=self.strict)
81
+ self._total_used += tokens
82
+
83
+ if self.strict and self._total_used > self.workflow_limit:
84
+ raise BudgetExceeded("workflow", self._total_used, self.workflow_limit)
85
+
86
+ def remaining(self, agent_name: str | None = None) -> int:
87
+ """Get remaining tokens for an agent or the whole workflow."""
88
+ if agent_name:
89
+ budget = self._agents.get(agent_name)
90
+ return budget.remaining if budget else self.per_agent_limit
91
+ return max(0, self.workflow_limit - self._total_used)
92
+
93
+ @property
94
+ def total_used(self) -> int:
95
+ return self._total_used
96
+
97
+ @property
98
+ def workflow_utilization(self) -> float:
99
+ return self._total_used / self.workflow_limit if self.workflow_limit > 0 else 0.0
100
+
101
+ def summary(self) -> dict[str, dict[str, int | float]]:
102
+ """Return a summary of all agent budgets."""
103
+ result: dict[str, dict[str, int | float]] = {
104
+ "workflow": {
105
+ "limit": self.workflow_limit,
106
+ "used": self._total_used,
107
+ "remaining": self.remaining(),
108
+ "utilization": self.workflow_utilization,
109
+ }
110
+ }
111
+ for name, budget in self._agents.items():
112
+ result[name] = {
113
+ "limit": budget.limit,
114
+ "used": budget.used,
115
+ "remaining": budget.remaining,
116
+ "utilization": budget.utilization,
117
+ }
118
+ return result
@@ -0,0 +1,144 @@
1
+ """MessageCompressor: reduce inter-agent message size while preserving key information.
2
+
3
+ Implements extractive compression: keeps the most important sentences
4
+ from verbose LLM reasoning traces, stripping filler and repetition.
5
+
6
+ Based on: "Cut the Crap: Economical Communication Pipeline for LLM-based MAS" (2024)
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import re
12
+ from dataclasses import dataclass
13
+
14
+
15
+ @dataclass(frozen=True)
16
+ class CompressionResult:
17
+ """Result of message compression.
18
+
19
+ Attributes:
20
+ original: The original text.
21
+ compressed: The compressed text.
22
+ original_tokens: Estimated original token count.
23
+ compressed_tokens: Estimated compressed token count.
24
+ savings_pct: Percentage of tokens saved (0-1).
25
+ """
26
+
27
+ original: str
28
+ compressed: str
29
+ original_tokens: int
30
+ compressed_tokens: int
31
+
32
+ @property
33
+ def savings_pct(self) -> float:
34
+ if self.original_tokens == 0:
35
+ return 0.0
36
+ return 1.0 - (self.compressed_tokens / self.original_tokens)
37
+
38
+
39
+ # Filler phrases commonly found in LLM reasoning traces
40
+ _FILLER_PATTERNS = [
41
+ r"(?i)\b(let me think|let's think|okay so|well|basically|essentially|in other words)\b[,.]?\s*",
42
+ r"(?i)\b(as I mentioned|as we discussed|to summarize so far|in summary)\b[,.]?\s*",
43
+ r"(?i)\b(it's worth noting that|it's important to note that|I should mention that)\b\s*",
44
+ r"(?i)\b(first of all|secondly|thirdly|finally|in conclusion)\b[,.]?\s*",
45
+ r"(?i)\b(I think|I believe|in my opinion|from my perspective)\b[,.]?\s*",
46
+ ]
47
+
48
+
49
+ class MessageCompressor:
50
+ """Compress inter-agent messages by removing filler and extracting key sentences.
51
+
52
+ Args:
53
+ target_ratio: Target compression ratio (0.3 = keep 30% of tokens).
54
+ min_sentence_length: Minimum characters for a sentence to be kept.
55
+ remove_filler: Whether to strip filler phrases.
56
+ """
57
+
58
+ def __init__(
59
+ self,
60
+ target_ratio: float = 0.5,
61
+ min_sentence_length: int = 20,
62
+ remove_filler: bool = True,
63
+ ) -> None:
64
+ self._target_ratio = target_ratio
65
+ self._min_sentence_length = min_sentence_length
66
+ self._remove_filler = remove_filler
67
+
68
+ def compress(self, text: str) -> CompressionResult:
69
+ """Compress a message text.
70
+
71
+ Strategy:
72
+ 1. Remove filler phrases
73
+ 2. Split into sentences
74
+ 3. Score sentences by information density
75
+ 4. Keep top sentences until target ratio met
76
+ """
77
+ original_tokens = len(text) // 4
78
+
79
+ if original_tokens <= 20:
80
+ # Too short to compress meaningfully
81
+ return CompressionResult(text, text, original_tokens, original_tokens)
82
+
83
+ working = text
84
+
85
+ # Step 1: Remove filler
86
+ if self._remove_filler:
87
+ for pattern in _FILLER_PATTERNS:
88
+ working = re.sub(pattern, "", working)
89
+
90
+ # Step 2: Split into sentences
91
+ sentences = re.split(r"(?<=[.!?])\s+", working)
92
+ sentences = [s.strip() for s in sentences if len(s.strip()) >= self._min_sentence_length]
93
+
94
+ if not sentences:
95
+ compressed_tokens = len(working) // 4
96
+ return CompressionResult(text, working.strip(), original_tokens, compressed_tokens)
97
+
98
+ # Step 3: Score sentences
99
+ scored = [(self._score_sentence(s), s) for s in sentences]
100
+ scored.sort(key=lambda x: x[0], reverse=True)
101
+
102
+ # Step 4: Keep top sentences until target ratio
103
+ target_tokens = int(original_tokens * self._target_ratio)
104
+ kept: list[str] = []
105
+ running_tokens = 0
106
+
107
+ for score, sentence in scored:
108
+ sent_tokens = len(sentence) // 4
109
+ if running_tokens + sent_tokens > target_tokens and kept:
110
+ break
111
+ kept.append(sentence)
112
+ running_tokens += sent_tokens
113
+
114
+ # Restore original order
115
+ ordered = [s for s in sentences if s in kept]
116
+ compressed = " ".join(ordered)
117
+ compressed_tokens = len(compressed) // 4
118
+
119
+ return CompressionResult(text, compressed, original_tokens, compressed_tokens)
120
+
121
+ @staticmethod
122
+ def _score_sentence(sentence: str) -> float:
123
+ """Score a sentence by information density (higher = more informative)."""
124
+ score = 0.0
125
+
126
+ # Longer sentences tend to carry more information (diminishing returns)
127
+ score += min(len(sentence.split()) / 20, 1.0) * 0.3
128
+
129
+ # Sentences with numbers/data are more informative
130
+ numbers = len(re.findall(r"\d+\.?\d*", sentence))
131
+ score += min(numbers / 3, 1.0) * 0.3
132
+
133
+ # Sentences with technical terms
134
+ technical = len(re.findall(
135
+ r"(?i)\b(result|conclusion|therefore|because|shows|indicates|found|"
136
+ r"evidence|data|percent|increase|decrease|significant)\b", sentence
137
+ ))
138
+ score += min(technical / 3, 1.0) * 0.2
139
+
140
+ # Shorter sentences with substance are preferred
141
+ if 10 <= len(sentence.split()) <= 25:
142
+ score += 0.2
143
+
144
+ return score
@@ -0,0 +1,102 @@
1
+ """CompressMiddleware: inject compression between pattern stages.
2
+
3
+ Wraps agents so that their outputs are automatically compressed
4
+ before being passed to the next agent, reducing inter-agent token transfer.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ from pyagent_patterns.base import Agent, Message
10
+ from pyagent_compress.budget import TokenBudget
11
+ from pyagent_compress.compressor import MessageCompressor
12
+
13
+
14
+ class CompressedAgent(Agent):
15
+ """Agent wrapper that compresses output messages.
16
+
17
+ Args:
18
+ agent: The original agent.
19
+ compressor: MessageCompressor instance.
20
+ budget: Optional TokenBudget for tracking.
21
+ """
22
+
23
+ def __init__(
24
+ self,
25
+ agent: Agent,
26
+ compressor: MessageCompressor,
27
+ budget: TokenBudget | None = None,
28
+ ) -> None:
29
+ super().__init__(
30
+ name=agent.name,
31
+ llm=agent.llm,
32
+ system_prompt=agent.system_prompt,
33
+ description=agent.description,
34
+ )
35
+ self._original = agent
36
+ self._compressor = compressor
37
+ self._budget = budget
38
+ self.compression_log: list[dict[str, int | float]] = []
39
+
40
+ async def run(self, messages: list[Message]) -> Message:
41
+ """Run the agent and compress the output."""
42
+ result = await self._original.run(messages)
43
+
44
+ # Compress the output
45
+ compressed = self._compressor.compress(result.content)
46
+ self.compression_log.append({
47
+ "original_tokens": compressed.original_tokens,
48
+ "compressed_tokens": compressed.compressed_tokens,
49
+ "savings_pct": compressed.savings_pct,
50
+ })
51
+
52
+ # Track budget if available
53
+ if self._budget:
54
+ self._budget.consume(self._original.name, compressed.compressed_tokens)
55
+
56
+ return Message(
57
+ role=result.role,
58
+ content=compressed.compressed,
59
+ name=result.name,
60
+ metadata={
61
+ **result.metadata,
62
+ "compressed": True,
63
+ "original_tokens": compressed.original_tokens,
64
+ "compressed_tokens": compressed.compressed_tokens,
65
+ "savings_pct": compressed.savings_pct,
66
+ },
67
+ )
68
+
69
+
70
+ class CompressMiddleware:
71
+ """Middleware that wraps agents with automatic output compression.
72
+
73
+ Usage:
74
+ middleware = CompressMiddleware(target_ratio=0.5)
75
+ compressed_agent = middleware.wrap(my_agent)
76
+
77
+ Args:
78
+ compressor: Optional MessageCompressor. Created with defaults if None.
79
+ budget: Optional TokenBudget for workflow-wide tracking.
80
+ target_ratio: Target compression ratio (used if compressor not provided).
81
+ """
82
+
83
+ def __init__(
84
+ self,
85
+ compressor: MessageCompressor | None = None,
86
+ budget: TokenBudget | None = None,
87
+ target_ratio: float = 0.5,
88
+ ) -> None:
89
+ self._compressor = compressor or MessageCompressor(target_ratio=target_ratio)
90
+ self._budget = budget
91
+
92
+ def wrap(self, agent: Agent) -> CompressedAgent:
93
+ """Wrap an agent with compression."""
94
+ return CompressedAgent(
95
+ agent=agent,
96
+ compressor=self._compressor,
97
+ budget=self._budget,
98
+ )
99
+
100
+ def wrap_all(self, agents: list[Agent]) -> list[CompressedAgent]:
101
+ """Wrap multiple agents."""
102
+ return [self.wrap(a) for a in agents]
@@ -0,0 +1,161 @@
1
+ """AgentPruner and InteractionPruner: reduce multi-agent overhead.
2
+
3
+ AgentPruner: detect and eliminate non-contributing agents mid-execution.
4
+ InteractionPruner: skip communication rounds when consensus is reached early.
5
+
6
+ Based on: arxiv:2503.18891 "AgentDropout: Dynamic Agent Elimination for Token-Efficient MAS"
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ from dataclasses import dataclass
12
+
13
+ from pyagent_patterns.base import Message
14
+
15
+
16
+ @dataclass(frozen=True)
17
+ class ContributionScore:
18
+ """How much an agent is contributing to the conversation."""
19
+
20
+ agent_name: str
21
+ score: float # 0-1
22
+ unique_info: float # 0-1 — how much unique information vs repetition
23
+ message_count: int
24
+
25
+
26
+ class AgentPruner:
27
+ """Detect and flag non-contributing agents for removal.
28
+
29
+ Scores each agent's contribution based on:
30
+ - Unique information (not repeated from other agents)
31
+ - Response diversity (different from own previous responses)
32
+ - Task relevance (overlap with original task keywords)
33
+
34
+ Args:
35
+ min_contribution: Minimum contribution score (0-1) to keep an agent.
36
+ window_size: Number of recent messages to analyze per agent.
37
+ """
38
+
39
+ def __init__(self, min_contribution: float = 0.3, window_size: int = 5) -> None:
40
+ self._min_contribution = min_contribution
41
+ self._window_size = window_size
42
+
43
+ def score_agents(
44
+ self,
45
+ messages: list[Message],
46
+ task: str,
47
+ ) -> list[ContributionScore]:
48
+ """Score each agent's contribution from message history."""
49
+ # Group messages by agent name
50
+ by_agent: dict[str, list[str]] = {}
51
+ for msg in messages:
52
+ name = msg.name or "unknown"
53
+ by_agent.setdefault(name, []).append(msg.content)
54
+
55
+ all_contents = [m.content for m in messages]
56
+ task_words = set(task.lower().split())
57
+
58
+ scores: list[ContributionScore] = []
59
+ for agent_name, contents in by_agent.items():
60
+ recent = contents[-self._window_size :]
61
+
62
+ # Unique info: how much of this agent's content is NOT in other agents' content
63
+ other_text = " ".join(
64
+ c for name, cs in by_agent.items() if name != agent_name for c in cs
65
+ ).lower()
66
+ unique_words = set()
67
+ for content in recent:
68
+ for word in content.lower().split():
69
+ if word not in other_text and len(word) > 3:
70
+ unique_words.add(word)
71
+
72
+ total_words = sum(len(c.split()) for c in recent)
73
+ unique_ratio = len(unique_words) / max(total_words, 1)
74
+
75
+ # Task relevance
76
+ agent_words = set(" ".join(recent).lower().split())
77
+ relevance = len(agent_words & task_words) / max(len(task_words), 1)
78
+
79
+ # Self-diversity (are recent messages different from each other?)
80
+ if len(recent) >= 2:
81
+ diversity = 1.0 - self._similarity(recent[-1], recent[-2])
82
+ else:
83
+ diversity = 1.0
84
+
85
+ score = 0.4 * unique_ratio + 0.3 * relevance + 0.3 * diversity
86
+ scores.append(ContributionScore(
87
+ agent_name=agent_name,
88
+ score=min(score, 1.0),
89
+ unique_info=unique_ratio,
90
+ message_count=len(contents),
91
+ ))
92
+
93
+ return scores
94
+
95
+ def should_prune(self, scores: list[ContributionScore]) -> list[str]:
96
+ """Return agent names that should be pruned (below min contribution)."""
97
+ return [s.agent_name for s in scores if s.score < self._min_contribution]
98
+
99
+ @staticmethod
100
+ def _similarity(a: str, b: str) -> float:
101
+ """Simple word-overlap similarity between two texts."""
102
+ words_a = set(a.lower().split())
103
+ words_b = set(b.lower().split())
104
+ if not words_a or not words_b:
105
+ return 0.0
106
+ return len(words_a & words_b) / len(words_a | words_b)
107
+
108
+
109
+ class InteractionPruner:
110
+ """Detect early consensus and skip remaining communication rounds.
111
+
112
+ Analyzes agent outputs to determine if consensus has been reached,
113
+ allowing patterns like Debate or Voting to terminate early.
114
+
115
+ Args:
116
+ consensus_threshold: Minimum similarity between agent outputs
117
+ to consider consensus reached (0-1).
118
+ min_rounds: Minimum rounds before early termination is allowed.
119
+ """
120
+
121
+ def __init__(
122
+ self,
123
+ consensus_threshold: float = 0.7,
124
+ min_rounds: int = 1,
125
+ ) -> None:
126
+ self._threshold = consensus_threshold
127
+ self._min_rounds = min_rounds
128
+
129
+ def has_consensus(self, outputs: list[str], current_round: int) -> bool:
130
+ """Check if agents have reached consensus.
131
+
132
+ Args:
133
+ outputs: Current round's outputs from all agents.
134
+ current_round: The current round number (1-indexed).
135
+
136
+ Returns:
137
+ True if consensus is reached and we can stop early.
138
+ """
139
+ if current_round < self._min_rounds:
140
+ return False
141
+
142
+ if len(outputs) < 2:
143
+ return True
144
+
145
+ # Check pairwise similarity
146
+ similarities: list[float] = []
147
+ for i in range(len(outputs)):
148
+ for j in range(i + 1, len(outputs)):
149
+ sim = self._similarity(outputs[i], outputs[j])
150
+ similarities.append(sim)
151
+
152
+ avg_similarity = sum(similarities) / len(similarities) if similarities else 0.0
153
+ return avg_similarity >= self._threshold
154
+
155
+ @staticmethod
156
+ def _similarity(a: str, b: str) -> float:
157
+ words_a = set(a.lower().split())
158
+ words_b = set(b.lower().split())
159
+ if not words_a or not words_b:
160
+ return 0.0
161
+ return len(words_a & words_b) / len(words_a | words_b)
File without changes
File without changes
@@ -0,0 +1,98 @@
1
+ """Tests for pyagent-compress."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import pytest
6
+
7
+ from pyagent_patterns.base import Agent, Message, MockLLM, Role
8
+ from pyagent_compress.budget import BudgetExceeded, TokenBudget
9
+ from pyagent_compress.compressor import MessageCompressor
10
+ from pyagent_compress.middleware import CompressMiddleware
11
+ from pyagent_compress.pruner import AgentPruner, InteractionPruner
12
+
13
+
14
+ def test_compressor_reduces_tokens():
15
+ compressor = MessageCompressor(target_ratio=0.5)
16
+ text = (
17
+ "Let me think about this carefully. Basically, the analysis shows that "
18
+ "revenue increased by 15% year-over-year. The data indicates a strong "
19
+ "upward trend in Q3 earnings. In other words, the company is performing "
20
+ "well above market expectations. It's worth noting that the profit margin "
21
+ "expanded to 23%, which is significant compared to the industry average of 18%. "
22
+ "The conclusion is that this represents a solid buy opportunity."
23
+ )
24
+ result = compressor.compress(text)
25
+ assert result.compressed_tokens < result.original_tokens
26
+ assert result.savings_pct > 0
27
+
28
+
29
+ def test_compressor_short_text_unchanged():
30
+ compressor = MessageCompressor()
31
+ result = compressor.compress("Short text")
32
+ assert result.compressed == result.original
33
+
34
+
35
+ def test_agent_pruner_scores():
36
+ pruner = AgentPruner(min_contribution=0.3)
37
+ messages = [
38
+ Message(role=Role.ASSISTANT, content="Unique analysis of market trends", name="analyst"),
39
+ Message(role=Role.ASSISTANT, content="Unique analysis of market trends", name="copycat"),
40
+ Message(role=Role.ASSISTANT, content="Different risk assessment with new data", name="risk"),
41
+ ]
42
+ scores = pruner.score_agents(messages, "analyze market trends")
43
+ assert len(scores) == 3
44
+
45
+
46
+ def test_interaction_pruner_consensus():
47
+ pruner = InteractionPruner(consensus_threshold=0.5, min_rounds=1)
48
+
49
+ # Same outputs → consensus
50
+ assert pruner.has_consensus(["The answer is 42", "The answer is 42"], current_round=1) is True
51
+
52
+ # Very different outputs → no consensus
53
+ assert pruner.has_consensus(
54
+ ["Completely different analysis about stocks",
55
+ "An unrelated discussion about weather patterns"],
56
+ current_round=1,
57
+ ) is False
58
+
59
+
60
+ def test_interaction_pruner_min_rounds():
61
+ pruner = InteractionPruner(consensus_threshold=0.5, min_rounds=2)
62
+ # Round 1: should not trigger even if consensus
63
+ assert pruner.has_consensus(["Same", "Same"], current_round=1) is False
64
+ # Round 2: now can trigger
65
+ assert pruner.has_consensus(["Same answer", "Same answer"], current_round=2) is True
66
+
67
+
68
+ def test_token_budget_tracking():
69
+ budget = TokenBudget(workflow_limit=10000, per_agent_limit=5000, strict=False)
70
+ budget.consume("agent_a", 1000)
71
+ budget.consume("agent_b", 2000)
72
+ assert budget.total_used == 3000
73
+ assert budget.remaining() == 7000
74
+ assert budget.remaining("agent_a") == 4000
75
+
76
+
77
+ def test_token_budget_strict_raises():
78
+ budget = TokenBudget(workflow_limit=100, per_agent_limit=50, strict=True)
79
+ budget.consume("agent_a", 40)
80
+ with pytest.raises(BudgetExceeded):
81
+ budget.consume("agent_a", 20) # Exceeds per-agent limit of 50
82
+
83
+
84
+ @pytest.mark.asyncio
85
+ async def test_compress_middleware():
86
+ llm = MockLLM(responses=[
87
+ "Let me think about this carefully. Basically, the revenue data shows "
88
+ "a 15% increase year-over-year, which is significant. The profit margin "
89
+ "expanded to 23%. In conclusion, this is a buy signal."
90
+ ])
91
+ agent = Agent("analyst", llm)
92
+
93
+ middleware = CompressMiddleware(target_ratio=0.5)
94
+ compressed = middleware.wrap(agent)
95
+
96
+ result = await compressed.run([Message.user("Analyze AAPL")])
97
+ assert result.metadata.get("compressed") is True
98
+ assert result.metadata.get("savings_pct", 0) >= 0