mcp-context-guard 1.0.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.
- mcp_context_guard-1.0.0/LICENSE +17 -0
- mcp_context_guard-1.0.0/PKG-INFO +79 -0
- mcp_context_guard-1.0.0/README.md +58 -0
- mcp_context_guard-1.0.0/mcp_context_guard.egg-info/PKG-INFO +79 -0
- mcp_context_guard-1.0.0/mcp_context_guard.egg-info/SOURCES.txt +14 -0
- mcp_context_guard-1.0.0/mcp_context_guard.egg-info/dependency_links.txt +1 -0
- mcp_context_guard-1.0.0/mcp_context_guard.egg-info/top_level.txt +1 -0
- mcp_context_guard-1.0.0/pyproject.toml +30 -0
- mcp_context_guard-1.0.0/setup.cfg +4 -0
- mcp_context_guard-1.0.0/src/__init__.py +5 -0
- mcp_context_guard-1.0.0/src/context_guard_engine.py +324 -0
- mcp_context_guard-1.0.0/src/server.py +75 -0
- mcp_context_guard-1.0.0/tests/test_context_guard.py +217 -0
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
Copyright (c) 2026 AMEOBIUS
|
|
3
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
4
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
5
|
+
in the Software without restriction, including without limitation the rights
|
|
6
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
7
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
8
|
+
furnished to do so, subject to the following conditions:
|
|
9
|
+
The above copyright notice and this permission notice shall be included in all
|
|
10
|
+
copies or substantial portions of the Software.
|
|
11
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
12
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
13
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
14
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
15
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
16
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
17
|
+
SOFTWARE.
|
|
@@ -0,0 +1,79 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: mcp-context-guard
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: Context window management: compress, deduplicate, filter tool outputs
|
|
5
|
+
Author: aaameobius-crypto
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/aaameobius-crypto/darkbot-ai-templates
|
|
8
|
+
Project-URL: Repository, https://github.com/aaameobius-crypto/darkbot-ai-templates
|
|
9
|
+
Keywords: mcp,ai,agent,harness,tools
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Intended Audience :: Developers
|
|
12
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
17
|
+
Requires-Python: >=3.10
|
|
18
|
+
Description-Content-Type: text/markdown
|
|
19
|
+
License-File: LICENSE
|
|
20
|
+
Dynamic: license-file
|
|
21
|
+
|
|
22
|
+
# MCP Context Guard — Context Window Management for AI Agents
|
|
23
|
+
|
|
24
|
+
> Compress tool outputs, manage token budgets, deduplicate content, and filter by relevance. Zero dependencies, pure Python stdlib.
|
|
25
|
+
|
|
26
|
+
## The Problem
|
|
27
|
+
|
|
28
|
+
AI agents waste context window tokens on:
|
|
29
|
+
- Verbose tool outputs (file reads, search results, logs)
|
|
30
|
+
- Duplicate content across tool calls
|
|
31
|
+
- Irrelevant passages that don't match the task
|
|
32
|
+
|
|
33
|
+
## The Solution
|
|
34
|
+
|
|
35
|
+
**MCP Context Guard** sits between your tools and the LLM, compressing and filtering everything that enters the context window.
|
|
36
|
+
|
|
37
|
+
## Tools (14)
|
|
38
|
+
|
|
39
|
+
| Tool | What it does |
|
|
40
|
+
|------|-------------|
|
|
41
|
+
| `compress` | Extractive summarization to N tokens |
|
|
42
|
+
| `set_budget` | Set a total token budget |
|
|
43
|
+
| `check_budget` | Check if text fits remaining budget |
|
|
44
|
+
| `consume_budget` | Deduct tokens from budget |
|
|
45
|
+
| `deduplicate` | Remove near-duplicate texts (Jaccard similarity) |
|
|
46
|
+
| `extract_key` | Extract top-N key sentences |
|
|
47
|
+
| `truncate_smart` | Truncate at sentence boundaries |
|
|
48
|
+
| `chunk` | Split into token-sized chunks with overlap |
|
|
49
|
+
| `token_count` | Estimate token count (word-based heuristic) |
|
|
50
|
+
| `summarize_history` | Compress conversation messages |
|
|
51
|
+
| `filter_relevant` | BM25 relevance scoring, return top-K passages |
|
|
52
|
+
| `merge_context` | Combine sources with dedup + compression |
|
|
53
|
+
| `get_stats` | Context usage statistics |
|
|
54
|
+
| `reset` | Reset all state |
|
|
55
|
+
|
|
56
|
+
## Install
|
|
57
|
+
|
|
58
|
+
```bash
|
|
59
|
+
git clone https://github.com/aaameobius-crypto/mcp-context-guard.git
|
|
60
|
+
cd mcp-context-guard
|
|
61
|
+
python -m src.server --stdio
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
## Tests
|
|
65
|
+
|
|
66
|
+
```bash
|
|
67
|
+
python -m pytest tests/ -v # 36 tests, all passing
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
## Inspiration
|
|
71
|
+
|
|
72
|
+
- [headroom](https://github.com/chopratejas/headroom) — 60-95% token reduction
|
|
73
|
+
- [context-mode](https://github.com/mksglu/context-mode) — Intercept tool output
|
|
74
|
+
- [LLMLingua](https://github.com/microsoft/LLMLingua) — Prompt compression
|
|
75
|
+
- [Autonomous Context Compression](https://blog.langchain.com/autonomous-context-compression/)
|
|
76
|
+
|
|
77
|
+
## License
|
|
78
|
+
|
|
79
|
+
MIT — aaameobius-crypto
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
# MCP Context Guard — Context Window Management for AI Agents
|
|
2
|
+
|
|
3
|
+
> Compress tool outputs, manage token budgets, deduplicate content, and filter by relevance. Zero dependencies, pure Python stdlib.
|
|
4
|
+
|
|
5
|
+
## The Problem
|
|
6
|
+
|
|
7
|
+
AI agents waste context window tokens on:
|
|
8
|
+
- Verbose tool outputs (file reads, search results, logs)
|
|
9
|
+
- Duplicate content across tool calls
|
|
10
|
+
- Irrelevant passages that don't match the task
|
|
11
|
+
|
|
12
|
+
## The Solution
|
|
13
|
+
|
|
14
|
+
**MCP Context Guard** sits between your tools and the LLM, compressing and filtering everything that enters the context window.
|
|
15
|
+
|
|
16
|
+
## Tools (14)
|
|
17
|
+
|
|
18
|
+
| Tool | What it does |
|
|
19
|
+
|------|-------------|
|
|
20
|
+
| `compress` | Extractive summarization to N tokens |
|
|
21
|
+
| `set_budget` | Set a total token budget |
|
|
22
|
+
| `check_budget` | Check if text fits remaining budget |
|
|
23
|
+
| `consume_budget` | Deduct tokens from budget |
|
|
24
|
+
| `deduplicate` | Remove near-duplicate texts (Jaccard similarity) |
|
|
25
|
+
| `extract_key` | Extract top-N key sentences |
|
|
26
|
+
| `truncate_smart` | Truncate at sentence boundaries |
|
|
27
|
+
| `chunk` | Split into token-sized chunks with overlap |
|
|
28
|
+
| `token_count` | Estimate token count (word-based heuristic) |
|
|
29
|
+
| `summarize_history` | Compress conversation messages |
|
|
30
|
+
| `filter_relevant` | BM25 relevance scoring, return top-K passages |
|
|
31
|
+
| `merge_context` | Combine sources with dedup + compression |
|
|
32
|
+
| `get_stats` | Context usage statistics |
|
|
33
|
+
| `reset` | Reset all state |
|
|
34
|
+
|
|
35
|
+
## Install
|
|
36
|
+
|
|
37
|
+
```bash
|
|
38
|
+
git clone https://github.com/aaameobius-crypto/mcp-context-guard.git
|
|
39
|
+
cd mcp-context-guard
|
|
40
|
+
python -m src.server --stdio
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
## Tests
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
python -m pytest tests/ -v # 36 tests, all passing
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
## Inspiration
|
|
50
|
+
|
|
51
|
+
- [headroom](https://github.com/chopratejas/headroom) — 60-95% token reduction
|
|
52
|
+
- [context-mode](https://github.com/mksglu/context-mode) — Intercept tool output
|
|
53
|
+
- [LLMLingua](https://github.com/microsoft/LLMLingua) — Prompt compression
|
|
54
|
+
- [Autonomous Context Compression](https://blog.langchain.com/autonomous-context-compression/)
|
|
55
|
+
|
|
56
|
+
## License
|
|
57
|
+
|
|
58
|
+
MIT — aaameobius-crypto
|
|
@@ -0,0 +1,79 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: mcp-context-guard
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: Context window management: compress, deduplicate, filter tool outputs
|
|
5
|
+
Author: aaameobius-crypto
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/aaameobius-crypto/darkbot-ai-templates
|
|
8
|
+
Project-URL: Repository, https://github.com/aaameobius-crypto/darkbot-ai-templates
|
|
9
|
+
Keywords: mcp,ai,agent,harness,tools
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Intended Audience :: Developers
|
|
12
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
17
|
+
Requires-Python: >=3.10
|
|
18
|
+
Description-Content-Type: text/markdown
|
|
19
|
+
License-File: LICENSE
|
|
20
|
+
Dynamic: license-file
|
|
21
|
+
|
|
22
|
+
# MCP Context Guard — Context Window Management for AI Agents
|
|
23
|
+
|
|
24
|
+
> Compress tool outputs, manage token budgets, deduplicate content, and filter by relevance. Zero dependencies, pure Python stdlib.
|
|
25
|
+
|
|
26
|
+
## The Problem
|
|
27
|
+
|
|
28
|
+
AI agents waste context window tokens on:
|
|
29
|
+
- Verbose tool outputs (file reads, search results, logs)
|
|
30
|
+
- Duplicate content across tool calls
|
|
31
|
+
- Irrelevant passages that don't match the task
|
|
32
|
+
|
|
33
|
+
## The Solution
|
|
34
|
+
|
|
35
|
+
**MCP Context Guard** sits between your tools and the LLM, compressing and filtering everything that enters the context window.
|
|
36
|
+
|
|
37
|
+
## Tools (14)
|
|
38
|
+
|
|
39
|
+
| Tool | What it does |
|
|
40
|
+
|------|-------------|
|
|
41
|
+
| `compress` | Extractive summarization to N tokens |
|
|
42
|
+
| `set_budget` | Set a total token budget |
|
|
43
|
+
| `check_budget` | Check if text fits remaining budget |
|
|
44
|
+
| `consume_budget` | Deduct tokens from budget |
|
|
45
|
+
| `deduplicate` | Remove near-duplicate texts (Jaccard similarity) |
|
|
46
|
+
| `extract_key` | Extract top-N key sentences |
|
|
47
|
+
| `truncate_smart` | Truncate at sentence boundaries |
|
|
48
|
+
| `chunk` | Split into token-sized chunks with overlap |
|
|
49
|
+
| `token_count` | Estimate token count (word-based heuristic) |
|
|
50
|
+
| `summarize_history` | Compress conversation messages |
|
|
51
|
+
| `filter_relevant` | BM25 relevance scoring, return top-K passages |
|
|
52
|
+
| `merge_context` | Combine sources with dedup + compression |
|
|
53
|
+
| `get_stats` | Context usage statistics |
|
|
54
|
+
| `reset` | Reset all state |
|
|
55
|
+
|
|
56
|
+
## Install
|
|
57
|
+
|
|
58
|
+
```bash
|
|
59
|
+
git clone https://github.com/aaameobius-crypto/mcp-context-guard.git
|
|
60
|
+
cd mcp-context-guard
|
|
61
|
+
python -m src.server --stdio
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
## Tests
|
|
65
|
+
|
|
66
|
+
```bash
|
|
67
|
+
python -m pytest tests/ -v # 36 tests, all passing
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
## Inspiration
|
|
71
|
+
|
|
72
|
+
- [headroom](https://github.com/chopratejas/headroom) — 60-95% token reduction
|
|
73
|
+
- [context-mode](https://github.com/mksglu/context-mode) — Intercept tool output
|
|
74
|
+
- [LLMLingua](https://github.com/microsoft/LLMLingua) — Prompt compression
|
|
75
|
+
- [Autonomous Context Compression](https://blog.langchain.com/autonomous-context-compression/)
|
|
76
|
+
|
|
77
|
+
## License
|
|
78
|
+
|
|
79
|
+
MIT — aaameobius-crypto
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
./src/__init__.py
|
|
5
|
+
./src/context_guard_engine.py
|
|
6
|
+
./src/server.py
|
|
7
|
+
mcp_context_guard.egg-info/PKG-INFO
|
|
8
|
+
mcp_context_guard.egg-info/SOURCES.txt
|
|
9
|
+
mcp_context_guard.egg-info/dependency_links.txt
|
|
10
|
+
mcp_context_guard.egg-info/top_level.txt
|
|
11
|
+
src/__init__.py
|
|
12
|
+
src/context_guard_engine.py
|
|
13
|
+
src/server.py
|
|
14
|
+
tests/test_context_guard.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
src
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=64"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "mcp-context-guard"
|
|
7
|
+
version = "1.0.0"
|
|
8
|
+
description = "Context window management: compress, deduplicate, filter tool outputs"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = {text = "MIT"}
|
|
11
|
+
requires-python = ">=3.10"
|
|
12
|
+
authors = [{name = "aaameobius-crypto"}]
|
|
13
|
+
keywords = ["mcp", "ai", "agent", "harness", "tools"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Development Status :: 4 - Beta",
|
|
16
|
+
"Intended Audience :: Developers",
|
|
17
|
+
"License :: OSI Approved :: MIT License",
|
|
18
|
+
"Programming Language :: Python :: 3.10",
|
|
19
|
+
"Programming Language :: Python :: 3.11",
|
|
20
|
+
"Programming Language :: Python :: 3.12",
|
|
21
|
+
"Topic :: Software Development :: Libraries :: Python Modules",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
[project.urls]
|
|
25
|
+
Homepage = "https://github.com/aaameobius-crypto/darkbot-ai-templates"
|
|
26
|
+
Repository = "https://github.com/aaameobius-crypto/darkbot-ai-templates"
|
|
27
|
+
|
|
28
|
+
[tool.setuptools]
|
|
29
|
+
packages = ["src"]
|
|
30
|
+
package-dir = {"" = "."}
|
|
@@ -0,0 +1,324 @@
|
|
|
1
|
+
"""Context Guard engine — zero dependencies.
|
|
2
|
+
Compresses tool outputs, manages token budgets, deduplicates content,
|
|
3
|
+
extracts key passages, and provides BM25-like relevance filtering.
|
|
4
|
+
"""
|
|
5
|
+
import re, json, math, hashlib, collections
|
|
6
|
+
from typing import Any, Dict, List, Optional, Tuple, Set
|
|
7
|
+
|
|
8
|
+
class ContextGuard:
|
|
9
|
+
# Token estimation: avg English word = 1.3 tokens
|
|
10
|
+
TOKEN_RATIO = 1.3
|
|
11
|
+
|
|
12
|
+
@staticmethod
|
|
13
|
+
def create_store() -> Dict:
|
|
14
|
+
return {
|
|
15
|
+
"stats": {"compressed": 0, "tokens_saved": 0, "chunks_created": 0, "filtered": 0, "errors": 0},
|
|
16
|
+
"budget": {"limit": 0, "used": 0},
|
|
17
|
+
"seen_hashes": set(),
|
|
18
|
+
"history": [],
|
|
19
|
+
}
|
|
20
|
+
|
|
21
|
+
@staticmethod
|
|
22
|
+
def _track(store: Dict, op: str, count: int = 1):
|
|
23
|
+
store["stats"][op] = store["stats"].get(op, 0) + count
|
|
24
|
+
|
|
25
|
+
@staticmethod
|
|
26
|
+
def estimate_tokens(text: str) -> int:
|
|
27
|
+
"""Estimate token count using word-based heuristic."""
|
|
28
|
+
words = len(re.findall(r"\b\w+\b", text))
|
|
29
|
+
return max(1, int(words * ContextGuard.TOKEN_RATIO))
|
|
30
|
+
|
|
31
|
+
@staticmethod
|
|
32
|
+
def _split_sentences(text: str) -> List[str]:
|
|
33
|
+
"""Split text into sentences."""
|
|
34
|
+
sentences = re.split(r"(?<=[.!?])\s+", text.strip())
|
|
35
|
+
return [s.strip() for s in sentences if s.strip()]
|
|
36
|
+
|
|
37
|
+
@staticmethod
|
|
38
|
+
def _word_freq(text: str) -> Dict[str, int]:
|
|
39
|
+
"""Compute word frequency map."""
|
|
40
|
+
words = re.findall(r"\b\w+\b", text.lower())
|
|
41
|
+
return dict(collections.Counter(words))
|
|
42
|
+
|
|
43
|
+
@staticmethod
|
|
44
|
+
def _sentence_score(sentence: str, word_freqs: Dict[str, int], total_words: int) -> float:
|
|
45
|
+
"""Score a sentence by sum of word frequencies normalized by length."""
|
|
46
|
+
words = re.findall(r"\b\w+\b", sentence.lower())
|
|
47
|
+
if not words:
|
|
48
|
+
return 0.0
|
|
49
|
+
score = sum(word_freqs.get(w, 0) for w in words) / math.sqrt(len(words))
|
|
50
|
+
# Boost sentences with numbers, code, or keywords
|
|
51
|
+
if re.search(r"\d|def |class |import |function |error |warning ", sentence, re.I):
|
|
52
|
+
score *= 1.2
|
|
53
|
+
# Penalize very short or very long sentences
|
|
54
|
+
if len(words) < 3:
|
|
55
|
+
score *= 0.5
|
|
56
|
+
if len(words) > 40:
|
|
57
|
+
score *= 0.7
|
|
58
|
+
return score
|
|
59
|
+
|
|
60
|
+
@staticmethod
|
|
61
|
+
def compress(store: Dict, text: str, max_tokens: int = 500, strategy: str = "extractive") -> Dict:
|
|
62
|
+
"""Compress text to approximately max_tokens using extractive summarization."""
|
|
63
|
+
ContextGuard._track(store, "compressed")
|
|
64
|
+
original_tokens = ContextGuard.estimate_tokens(text)
|
|
65
|
+
if original_tokens <= max_tokens:
|
|
66
|
+
return {"success": True, "text": text, "original_tokens": original_tokens, "compressed_tokens": original_tokens, "saved": 0, "ratio": 1.0}
|
|
67
|
+
|
|
68
|
+
sentences = ContextGuard._split_sentences(text)
|
|
69
|
+
if not sentences:
|
|
70
|
+
return {"success": True, "text": text, "original_tokens": original_tokens, "compressed_tokens": original_tokens, "saved": 0, "ratio": 1.0}
|
|
71
|
+
|
|
72
|
+
word_freqs = ContextGuard._word_freq(text)
|
|
73
|
+
total_words = sum(word_freqs.values())
|
|
74
|
+
|
|
75
|
+
scored = [(ContextGuard._sentence_score(s, word_freqs, total_words), i, s) for i, s in enumerate(sentences)]
|
|
76
|
+
scored.sort(key=lambda x: (-x[0], x[1])) # Best score first, preserve order for ties
|
|
77
|
+
|
|
78
|
+
# Greedily add sentences until we hit token budget
|
|
79
|
+
selected = []
|
|
80
|
+
current_tokens = 0
|
|
81
|
+
for score, idx, sent in scored:
|
|
82
|
+
sent_tokens = ContextGuard.estimate_tokens(sent)
|
|
83
|
+
if current_tokens + sent_tokens > max_tokens:
|
|
84
|
+
continue
|
|
85
|
+
selected.append((idx, sent))
|
|
86
|
+
current_tokens += sent_tokens
|
|
87
|
+
|
|
88
|
+
selected.sort(key=lambda x: x[0]) # Restore original order
|
|
89
|
+
result = " ".join(s for _, s in selected)
|
|
90
|
+
|
|
91
|
+
compressed_tokens = ContextGuard.estimate_tokens(result)
|
|
92
|
+
saved = original_tokens - compressed_tokens
|
|
93
|
+
ContextGuard._track(store, "tokens_saved", saved)
|
|
94
|
+
|
|
95
|
+
return {
|
|
96
|
+
"success": True,
|
|
97
|
+
"text": result,
|
|
98
|
+
"original_tokens": original_tokens,
|
|
99
|
+
"compressed_tokens": compressed_tokens,
|
|
100
|
+
"saved": saved,
|
|
101
|
+
"ratio": round(compressed_tokens / max(original_tokens, 1), 4),
|
|
102
|
+
"strategy": strategy,
|
|
103
|
+
"sentences_kept": len(selected),
|
|
104
|
+
"sentences_total": len(sentences),
|
|
105
|
+
}
|
|
106
|
+
|
|
107
|
+
@staticmethod
|
|
108
|
+
def set_budget(store: Dict, limit: int) -> Dict:
|
|
109
|
+
store["budget"]["limit"] = limit
|
|
110
|
+
return {"success": True, "limit": limit, "used": store["budget"]["used"]}
|
|
111
|
+
|
|
112
|
+
@staticmethod
|
|
113
|
+
def check_budget(store: Dict, text: str) -> Dict:
|
|
114
|
+
"""Check if text fits within remaining budget."""
|
|
115
|
+
tokens = ContextGuard.estimate_tokens(text)
|
|
116
|
+
remaining = store["budget"]["limit"] - store["budget"]["used"]
|
|
117
|
+
return {
|
|
118
|
+
"success": True,
|
|
119
|
+
"tokens": tokens,
|
|
120
|
+
"remaining": remaining,
|
|
121
|
+
"fits": tokens <= remaining if store["budget"]["limit"] > 0 else True,
|
|
122
|
+
"would_overflow": tokens > remaining if store["budget"]["limit"] > 0 else False,
|
|
123
|
+
}
|
|
124
|
+
|
|
125
|
+
@staticmethod
|
|
126
|
+
def consume_budget(store: Dict, text: str) -> Dict:
|
|
127
|
+
"""Add text tokens to budget usage."""
|
|
128
|
+
tokens = ContextGuard.estimate_tokens(text)
|
|
129
|
+
store["budget"]["used"] += tokens
|
|
130
|
+
remaining = store["budget"]["limit"] - store["budget"]["used"] if store["budget"]["limit"] > 0 else -1
|
|
131
|
+
return {"success": True, "consumed": tokens, "total_used": store["budget"]["used"], "remaining": remaining}
|
|
132
|
+
|
|
133
|
+
@staticmethod
|
|
134
|
+
def deduplicate(store: Dict, texts: List[str], similarity_threshold: float = 0.8) -> Dict:
|
|
135
|
+
"""Remove near-duplicate texts using Jaccard similarity on word sets."""
|
|
136
|
+
ContextGuard._track(store, "filtered")
|
|
137
|
+
if not texts:
|
|
138
|
+
return {"success": True, "unique": [], "removed": 0, "kept": 0}
|
|
139
|
+
|
|
140
|
+
unique = []
|
|
141
|
+
removed = 0
|
|
142
|
+
for text in texts:
|
|
143
|
+
words = set(re.findall(r"\b\w+\b", text.lower()))
|
|
144
|
+
is_dup = False
|
|
145
|
+
for u in unique:
|
|
146
|
+
u_words = set(re.findall(r"\b\w+\b", u.lower()))
|
|
147
|
+
if u_words and words:
|
|
148
|
+
jaccard = len(words & u_words) / len(words | u_words)
|
|
149
|
+
if jaccard >= similarity_threshold:
|
|
150
|
+
is_dup = True
|
|
151
|
+
removed += 1
|
|
152
|
+
break
|
|
153
|
+
if not is_dup:
|
|
154
|
+
unique.append(text)
|
|
155
|
+
|
|
156
|
+
ContextGuard._track(store, "tokens_saved", ContextGuard.estimate_tokens(" ".join(texts)) - ContextGuard.estimate_tokens(" ".join(unique)))
|
|
157
|
+
return {"success": True, "unique": unique, "removed": removed, "kept": len(unique)}
|
|
158
|
+
|
|
159
|
+
@staticmethod
|
|
160
|
+
def extract_key(store: Dict, text: str, num_keys: int = 5) -> Dict:
|
|
161
|
+
"""Extract key sentences from text."""
|
|
162
|
+
ContextGuard._track(store, "compressed")
|
|
163
|
+
sentences = ContextGuard._split_sentences(text)
|
|
164
|
+
if not sentences:
|
|
165
|
+
return {"success": True, "keys": [], "count": 0}
|
|
166
|
+
|
|
167
|
+
word_freqs = ContextGuard._word_freq(text)
|
|
168
|
+
total_words = sum(word_freqs.values())
|
|
169
|
+
scored = [(ContextGuard._sentence_score(s, word_freqs, total_words), i, s) for i, s in enumerate(sentences)]
|
|
170
|
+
scored.sort(key=lambda x: (-x[0], x[1]))
|
|
171
|
+
top = scored[:num_keys]
|
|
172
|
+
top.sort(key=lambda x: x[1]) # Restore order
|
|
173
|
+
|
|
174
|
+
return {"success": True, "keys": [s for _, _, s in top], "count": len(top)}
|
|
175
|
+
|
|
176
|
+
@staticmethod
|
|
177
|
+
def truncate_smart(store: Dict, text: str, max_tokens: int = 200) -> Dict:
|
|
178
|
+
"""Smart truncation that preserves structure."""
|
|
179
|
+
ContextGuard._track(store, "compressed")
|
|
180
|
+
original_tokens = ContextGuard.estimate_tokens(text)
|
|
181
|
+
if original_tokens <= max_tokens:
|
|
182
|
+
return {"success": True, "text": text, "truncated": False, "original_tokens": original_tokens, "final_tokens": original_tokens}
|
|
183
|
+
|
|
184
|
+
# Try to cut at sentence boundary
|
|
185
|
+
sentences = ContextGuard._split_sentences(text)
|
|
186
|
+
result = ""
|
|
187
|
+
for s in sentences:
|
|
188
|
+
candidate = result + " " + s if result else s
|
|
189
|
+
if ContextGuard.estimate_tokens(candidate) > max_tokens:
|
|
190
|
+
break
|
|
191
|
+
result = candidate
|
|
192
|
+
|
|
193
|
+
if not result:
|
|
194
|
+
# Fallback: hard cut at word boundary
|
|
195
|
+
words = text.split()
|
|
196
|
+
while words and ContextGuard.estimate_tokens(" ".join(words)) > max_tokens:
|
|
197
|
+
words.pop()
|
|
198
|
+
result = " ".join(words) + "..."
|
|
199
|
+
|
|
200
|
+
final_tokens = ContextGuard.estimate_tokens(result)
|
|
201
|
+
saved = original_tokens - final_tokens
|
|
202
|
+
ContextGuard._track(store, "tokens_saved", saved)
|
|
203
|
+
|
|
204
|
+
return {"success": True, "text": result, "truncated": True, "original_tokens": original_tokens, "final_tokens": final_tokens, "saved": saved}
|
|
205
|
+
|
|
206
|
+
@staticmethod
|
|
207
|
+
def chunk(store: Dict, text: str, max_tokens: int = 500, overlap: int = 50) -> Dict:
|
|
208
|
+
"""Split text into token-sized chunks with optional overlap."""
|
|
209
|
+
ContextGuard._track(store, "chunks_created")
|
|
210
|
+
sentences = ContextGuard._split_sentences(text)
|
|
211
|
+
chunks = []
|
|
212
|
+
current = ""
|
|
213
|
+
for s in sentences:
|
|
214
|
+
candidate = current + " " + s if current else s
|
|
215
|
+
if ContextGuard.estimate_tokens(candidate) > max_tokens and current:
|
|
216
|
+
chunks.append(current.strip())
|
|
217
|
+
# Overlap: keep last few words
|
|
218
|
+
if overlap > 0:
|
|
219
|
+
words = current.split()
|
|
220
|
+
overlap_text = " ".join(words[-overlap:])
|
|
221
|
+
current = overlap_text + " " + s
|
|
222
|
+
else:
|
|
223
|
+
current = s
|
|
224
|
+
else:
|
|
225
|
+
current = candidate
|
|
226
|
+
if current.strip():
|
|
227
|
+
chunks.append(current.strip())
|
|
228
|
+
|
|
229
|
+
return {"success": True, "chunks": chunks, "count": len(chunks), "total_tokens": sum(ContextGuard.estimate_tokens(c) for c in chunks)}
|
|
230
|
+
|
|
231
|
+
@staticmethod
|
|
232
|
+
def token_count(store: Dict, text: str) -> Dict:
|
|
233
|
+
"""Estimate token count for text."""
|
|
234
|
+
tokens = ContextGuard.estimate_tokens(text)
|
|
235
|
+
words = len(re.findall(r"\b\w+\b", text))
|
|
236
|
+
chars = len(text)
|
|
237
|
+
return {"success": True, "tokens": tokens, "words": words, "chars": chars, "ratio": round(tokens / max(words, 1), 2)}
|
|
238
|
+
|
|
239
|
+
@staticmethod
|
|
240
|
+
def summarize_history(store: Dict, messages: List[Dict], max_tokens: int = 300) -> Dict:
|
|
241
|
+
"""Compress conversation history into key points."""
|
|
242
|
+
ContextGuard._track(store, "compressed")
|
|
243
|
+
# Extract text content from messages
|
|
244
|
+
texts = []
|
|
245
|
+
for msg in messages:
|
|
246
|
+
role = msg.get("role", "unknown")
|
|
247
|
+
content = msg.get("content", "")
|
|
248
|
+
if isinstance(content, str):
|
|
249
|
+
texts.append(f"[{role}] {content}")
|
|
250
|
+
elif isinstance(content, list):
|
|
251
|
+
for block in content:
|
|
252
|
+
if isinstance(block, dict) and block.get("type") == "text":
|
|
253
|
+
texts.append(f"[{role}] {block.get('text', '')}")
|
|
254
|
+
|
|
255
|
+
combined = " ".join(texts)
|
|
256
|
+
result = ContextGuard.compress(store, combined, max_tokens=max_tokens)
|
|
257
|
+
result["messages_processed"] = len(messages)
|
|
258
|
+
return result
|
|
259
|
+
|
|
260
|
+
@staticmethod
|
|
261
|
+
def filter_relevant(store: Dict, query: str, passages: List[str], top_k: int = 5) -> Dict:
|
|
262
|
+
"""BM25-like relevance scoring. Return top-K passages for query."""
|
|
263
|
+
ContextGuard._track(store, "filtered")
|
|
264
|
+
if not passages:
|
|
265
|
+
return {"success": True, "results": [], "count": 0}
|
|
266
|
+
|
|
267
|
+
query_words = set(re.findall(r"\b\w+\b", query.lower()))
|
|
268
|
+
# Build document frequency map
|
|
269
|
+
doc_freqs = collections.Counter()
|
|
270
|
+
for p in passages:
|
|
271
|
+
p_words = set(re.findall(r"\b\w+\b", p.lower()))
|
|
272
|
+
for w in p_words:
|
|
273
|
+
doc_freqs[w] += 1
|
|
274
|
+
|
|
275
|
+
N = len(passages)
|
|
276
|
+
avg_len = sum(len(re.findall(r"\b\w+\b", p)) for p in passages) / N
|
|
277
|
+
k1, b = 1.5, 0.75 # BM25 params
|
|
278
|
+
|
|
279
|
+
scored = []
|
|
280
|
+
for i, p in enumerate(passages):
|
|
281
|
+
p_words = re.findall(r"\b\w+\b", p.lower())
|
|
282
|
+
p_len = len(p_words)
|
|
283
|
+
p_word_counts = collections.Counter(p_words)
|
|
284
|
+
score = 0.0
|
|
285
|
+
for qw in query_words:
|
|
286
|
+
if qw not in p_word_counts:
|
|
287
|
+
continue
|
|
288
|
+
tf = p_word_counts[qw]
|
|
289
|
+
df = doc_freqs.get(qw, 0)
|
|
290
|
+
idf = math.log((N - df + 0.5) / (df + 0.5) + 1)
|
|
291
|
+
score += idf * (tf * (k1 + 1)) / (tf + k1 * (1 - b + b * p_len / max(avg_len, 1)))
|
|
292
|
+
scored.append((score, i, p))
|
|
293
|
+
|
|
294
|
+
scored.sort(key=lambda x: -x[0])
|
|
295
|
+
top = scored[:top_k]
|
|
296
|
+
results = [{"passage": p, "score": round(s, 4), "index": i} for s, i, p in top if s > 0]
|
|
297
|
+
|
|
298
|
+
return {"success": True, "results": results, "count": len(results), "total_passages": len(passages)}
|
|
299
|
+
|
|
300
|
+
@staticmethod
|
|
301
|
+
def merge_context(store: Dict, sources: List[str], max_tokens: int = 1000, deduplicate: bool = True) -> Dict:
|
|
302
|
+
"""Merge multiple text sources, deduplicate, sort by relevance, truncate."""
|
|
303
|
+
ContextGuard._track(store, "compressed")
|
|
304
|
+
if deduplicate:
|
|
305
|
+
dedup = ContextGuard.deduplicate(store, sources)
|
|
306
|
+
sources = dedup["unique"]
|
|
307
|
+
|
|
308
|
+
combined = " ".join(sources)
|
|
309
|
+
result = ContextGuard.compress(store, combined, max_tokens=max_tokens)
|
|
310
|
+
result["sources_merged"] = len(sources)
|
|
311
|
+
return result
|
|
312
|
+
|
|
313
|
+
@staticmethod
|
|
314
|
+
def get_stats(store: Dict) -> Dict:
|
|
315
|
+
return {"success": True, **store["stats"], "budget": store["budget"]}
|
|
316
|
+
|
|
317
|
+
@staticmethod
|
|
318
|
+
def reset(store: Dict) -> Dict:
|
|
319
|
+
old = ContextGuard.get_stats(store)
|
|
320
|
+
store["stats"] = {"compressed": 0, "tokens_saved": 0, "chunks_created": 0, "filtered": 0, "errors": 0}
|
|
321
|
+
store["budget"] = {"limit": 0, "used": 0}
|
|
322
|
+
store["seen_hashes"] = set()
|
|
323
|
+
store["history"] = []
|
|
324
|
+
return {"success": True, "reset": old}
|
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
"""MCP Server for Context Guard — compress, budget, deduplicate, filter."""
|
|
2
|
+
import json, sys, argparse
|
|
3
|
+
from .context_guard_engine import ContextGuard
|
|
4
|
+
|
|
5
|
+
_store = ContextGuard.create_store()
|
|
6
|
+
|
|
7
|
+
TOOL_DEFS = [
|
|
8
|
+
{"name":"compress","description":"Compress text to a token budget using extractive summarization.","inputSchema":{"type":"object","properties":{"text":{"type":"string"},"max_tokens":{"type":"integer","default":500},"strategy":{"type":"string","default":"extractive"}},"required":["text"]}},
|
|
9
|
+
{"name":"set_budget","description":"Set a token budget limit.","inputSchema":{"type":"object","properties":{"limit":{"type":"integer"}},"required":["limit"]}},
|
|
10
|
+
{"name":"check_budget","description":"Check if text fits within remaining token budget.","inputSchema":{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}},
|
|
11
|
+
{"name":"consume_budget","description":"Add text tokens to budget usage.","inputSchema":{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}},
|
|
12
|
+
{"name":"deduplicate","description":"Remove near-duplicate texts using Jaccard similarity.","inputSchema":{"type":"object","properties":{"texts":{"type":"array","items":{"type":"string"}},"similarity_threshold":{"type":"number","default":0.8}},"required":["texts"]}},
|
|
13
|
+
{"name":"extract_key","description":"Extract key sentences from text.","inputSchema":{"type":"object","properties":{"text":{"type":"string"},"num_keys":{"type":"integer","default":5}},"required":["text"]}},
|
|
14
|
+
{"name":"truncate_smart","description":"Truncate text at sentence boundaries to fit token budget.","inputSchema":{"type":"object","properties":{"text":{"type":"string"},"max_tokens":{"type":"integer","default":200}},"required":["text"]}},
|
|
15
|
+
{"name":"chunk","description":"Split text into token-sized chunks with overlap.","inputSchema":{"type":"object","properties":{"text":{"type":"string"},"max_tokens":{"type":"integer","default":500},"overlap":{"type":"integer","default":50}},"required":["text"]}},
|
|
16
|
+
{"name":"token_count","description":"Estimate token count for text.","inputSchema":{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}},
|
|
17
|
+
{"name":"summarize_history","description":"Compress conversation history into key points.","inputSchema":{"type":"object","properties":{"messages":{"type":"array","items":{"type":"object"}},"max_tokens":{"type":"integer","default":300}},"required":["messages"]}},
|
|
18
|
+
{"name":"filter_relevant","description":"BM25 relevance scoring. Return top-K passages for a query.","inputSchema":{"type":"object","properties":{"query":{"type":"string"},"passages":{"type":"array","items":{"type":"string"}},"top_k":{"type":"integer","default":5}},"required":["query","passages"]}},
|
|
19
|
+
{"name":"merge_context","description":"Merge multiple text sources with dedup and compression.","inputSchema":{"type":"object","properties":{"sources":{"type":"array","items":{"type":"string"}},"max_tokens":{"type":"integer","default":1000},"deduplicate":{"type":"boolean","default":True}},"required":["sources"]}},
|
|
20
|
+
{"name":"get_stats","description":"Get context guard statistics.","inputSchema":{"type":"object","properties":{},"required":[]}},
|
|
21
|
+
{"name":"reset","description":"Reset all stats, budgets, and history.","inputSchema":{"type":"object","properties":{},"required":[]}},
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
class MCPContextGuardServer:
|
|
25
|
+
def __init__(self,name="mcp-context-guard",version="1.0.0"):
|
|
26
|
+
self.name=name;self.version=version
|
|
27
|
+
def list_tools(self):return TOOL_DEFS
|
|
28
|
+
def manifest(self):return{"server":{"name":self.name,"version":self.version},"capabilities":{"tools":{"listChanged":False}},"tools":self.list_tools()}
|
|
29
|
+
def handle_tool_call(self,name,args):
|
|
30
|
+
try:
|
|
31
|
+
if name=="compress":return json.dumps(ContextGuard.compress(_store,args["text"],args.get("max_tokens",500),args.get("strategy","extractive")))
|
|
32
|
+
elif name=="set_budget":return json.dumps(ContextGuard.set_budget(_store,args["limit"]))
|
|
33
|
+
elif name=="check_budget":return json.dumps(ContextGuard.check_budget(_store,args["text"]))
|
|
34
|
+
elif name=="consume_budget":return json.dumps(ContextGuard.consume_budget(_store,args["text"]))
|
|
35
|
+
elif name=="deduplicate":return json.dumps(ContextGuard.deduplicate(_store,args["texts"],args.get("similarity_threshold",0.8)))
|
|
36
|
+
elif name=="extract_key":return json.dumps(ContextGuard.extract_key(_store,args["text"],args.get("num_keys",5)))
|
|
37
|
+
elif name=="truncate_smart":return json.dumps(ContextGuard.truncate_smart(_store,args["text"],args.get("max_tokens",200)))
|
|
38
|
+
elif name=="chunk":return json.dumps(ContextGuard.chunk(_store,args["text"],args.get("max_tokens",500),args.get("overlap",50)))
|
|
39
|
+
elif name=="token_count":return json.dumps(ContextGuard.token_count(_store,args["text"]))
|
|
40
|
+
elif name=="summarize_history":return json.dumps(ContextGuard.summarize_history(_store,args["messages"],args.get("max_tokens",300)))
|
|
41
|
+
elif name=="filter_relevant":return json.dumps(ContextGuard.filter_relevant(_store,args["query"],args["passages"],args.get("top_k",5)))
|
|
42
|
+
elif name=="merge_context":return json.dumps(ContextGuard.merge_context(_store,args["sources"],args.get("max_tokens",1000),args.get("deduplicate",True)))
|
|
43
|
+
elif name=="get_stats":return json.dumps(ContextGuard.get_stats(_store))
|
|
44
|
+
elif name=="reset":return json.dumps(ContextGuard.reset(_store))
|
|
45
|
+
else:return json.dumps({"error":f"Unknown tool: {name}"})
|
|
46
|
+
except KeyError as e:return json.dumps({"error":f"Missing required parameter: {e}","tool":name})
|
|
47
|
+
except Exception as e:return json.dumps({"error":str(e),"tool":name})
|
|
48
|
+
|
|
49
|
+
def _run_stdio():
|
|
50
|
+
server=MCPContextGuardServer()
|
|
51
|
+
for line in sys.stdin:
|
|
52
|
+
line=line.strip()
|
|
53
|
+
if not line:continue
|
|
54
|
+
try:request=json.loads(line)
|
|
55
|
+
except json.JSONDecodeError:print(json.dumps({"jsonrpc":"2.0","error":{"code":-32700,"message":"Parse error"}}),flush=True);continue
|
|
56
|
+
method=request.get("method","");req_id=request.get("id");params=request.get("params",{})
|
|
57
|
+
if method=="initialize":response={"jsonrpc":"2.0","id":req_id,"result":{"server":server.name,"version":server.version}}
|
|
58
|
+
elif method=="tools/list":response={"jsonrpc":"2.0","id":req_id,"result":{"tools":server.list_tools()}}
|
|
59
|
+
elif method=="tools/call":
|
|
60
|
+
result=server.handle_tool_call(params.get("name",""),params.get("arguments",{}))
|
|
61
|
+
response={"jsonrpc":"2.0","id":req_id,"result":{"content":[{"type":"text","text":result}]}}
|
|
62
|
+
elif method=="shutdown":response={"jsonrpc":"2.0","id":req_id,"result":{}};print(json.dumps(response),flush=True);break
|
|
63
|
+
else:response={"jsonrpc":"2.0","id":req_id,"error":{"code":-32601,"message":f"Method not found: {method}"}}
|
|
64
|
+
print(json.dumps(response),flush=True)
|
|
65
|
+
|
|
66
|
+
def main():
|
|
67
|
+
parser=argparse.ArgumentParser(description="MCP Context Guard Server")
|
|
68
|
+
parser.add_argument("--stdio",action="store_true")
|
|
69
|
+
parser.add_argument("--manifest",action="store_true")
|
|
70
|
+
args=parser.parse_args()
|
|
71
|
+
if args.manifest:print(json.dumps(MCPContextGuardServer().manifest(),indent=2))
|
|
72
|
+
elif args.stdio:_run_stdio()
|
|
73
|
+
else:parser.print_help()
|
|
74
|
+
|
|
75
|
+
if __name__=="__main__":main()
|
|
@@ -0,0 +1,217 @@
|
|
|
1
|
+
"""Tests for MCP Context Guard — compress, budget, deduplicate, filter, chunk."""
|
|
2
|
+
import json, pytest, os, sys
|
|
3
|
+
from unittest.mock import patch
|
|
4
|
+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
|
5
|
+
from src.server import MCPContextGuardServer, TOOL_DEFS
|
|
6
|
+
from src.context_guard_engine import ContextGuard
|
|
7
|
+
|
|
8
|
+
class TestToolDefs:
|
|
9
|
+
def test_names(self):
|
|
10
|
+
for t in TOOL_DEFS: assert "name" in t and len(t["name"])>0
|
|
11
|
+
def test_descs(self):
|
|
12
|
+
for t in TOOL_DEFS: assert "description" in t and len(t["description"])>15
|
|
13
|
+
def test_schema(self):
|
|
14
|
+
for t in TOOL_DEFS: assert "inputSchema" in t and t["inputSchema"]["type"]=="object"
|
|
15
|
+
def test_count(self):
|
|
16
|
+
assert len(TOOL_DEFS)==14
|
|
17
|
+
def test_required(self):
|
|
18
|
+
names={t["name"] for t in TOOL_DEFS}
|
|
19
|
+
expected={"compress","set_budget","check_budget","consume_budget","deduplicate","extract_key","truncate_smart","chunk","token_count","summarize_history","filter_relevant","merge_context","get_stats","reset"}
|
|
20
|
+
assert names==expected
|
|
21
|
+
|
|
22
|
+
class TestManifest:
|
|
23
|
+
def test_manifest(self):
|
|
24
|
+
s=MCPContextGuardServer();m=s.manifest()
|
|
25
|
+
assert m["server"]["name"]=="mcp-context-guard"
|
|
26
|
+
assert len(m["tools"])==14
|
|
27
|
+
|
|
28
|
+
class TestTokenCount:
|
|
29
|
+
def test_basic(self):
|
|
30
|
+
s=ContextGuard.create_store()
|
|
31
|
+
r=ContextGuard.token_count(s,"hello world test")
|
|
32
|
+
assert r["tokens"]>0
|
|
33
|
+
assert r["words"]==3
|
|
34
|
+
def test_empty(self):
|
|
35
|
+
s=ContextGuard.create_store()
|
|
36
|
+
r=ContextGuard.token_count(s,"")
|
|
37
|
+
assert r["tokens"]==1 # max(1, 0)
|
|
38
|
+
|
|
39
|
+
class TestCompress:
|
|
40
|
+
def test_short_text(self):
|
|
41
|
+
s=ContextGuard.create_store()
|
|
42
|
+
r=ContextGuard.compress(s,"Hello world.",max_tokens=100)
|
|
43
|
+
assert r["ratio"]==1.0
|
|
44
|
+
assert r["saved"]==0
|
|
45
|
+
def test_long_text(self):
|
|
46
|
+
s=ContextGuard.create_store()
|
|
47
|
+
text = ". ".join(f"Sentence number {i} about topic {i%3}" for i in range(50))
|
|
48
|
+
r=ContextGuard.compress(s,text,max_tokens=50)
|
|
49
|
+
assert r["compressed_tokens"]<=60 # approximate
|
|
50
|
+
assert r["saved"]>0
|
|
51
|
+
assert r["ratio"]<1.0
|
|
52
|
+
def test_preserves_structure(self):
|
|
53
|
+
s=ContextGuard.create_store()
|
|
54
|
+
text = "First important sentence. Second sentence here. Third one is also key. Fourth is filler about nothing. Fifth concludes everything."
|
|
55
|
+
r=ContextGuard.compress(s,text,max_tokens=20)
|
|
56
|
+
assert len(r["text"])>0
|
|
57
|
+
|
|
58
|
+
class TestBudget:
|
|
59
|
+
def test_set_and_check(self):
|
|
60
|
+
s=ContextGuard.create_store()
|
|
61
|
+
ContextGuard.set_budget(s, 100)
|
|
62
|
+
r=ContextGuard.check_budget(s, "hello world")
|
|
63
|
+
assert r["fits"] is True
|
|
64
|
+
def test_overflow(self):
|
|
65
|
+
s=ContextGuard.create_store()
|
|
66
|
+
ContextGuard.set_budget(s, 5)
|
|
67
|
+
r=ContextGuard.check_budget(s, "this is a very long text that exceeds the budget")
|
|
68
|
+
assert r["would_overflow"] is True
|
|
69
|
+
def test_consume(self):
|
|
70
|
+
s=ContextGuard.create_store()
|
|
71
|
+
ContextGuard.set_budget(s, 100)
|
|
72
|
+
r=ContextGuard.consume_budget(s, "hello world test")
|
|
73
|
+
assert r["consumed"]>0
|
|
74
|
+
assert r["total_used"]>0
|
|
75
|
+
|
|
76
|
+
class TestDeduplicate:
|
|
77
|
+
def test_removes_dup(self):
|
|
78
|
+
s=ContextGuard.create_store()
|
|
79
|
+
texts = ["hello world", "hello world", "different text"]
|
|
80
|
+
r=ContextGuard.deduplicate(s, texts)
|
|
81
|
+
assert r["kept"]==2
|
|
82
|
+
assert r["removed"]==1
|
|
83
|
+
def test_near_dup(self):
|
|
84
|
+
s=ContextGuard.create_store()
|
|
85
|
+
texts = ["The quick brown fox jumps over the lazy dog", "The quick brown fox jumps over the lazy dog today"]
|
|
86
|
+
r=ContextGuard.deduplicate(s, texts, similarity_threshold=0.7)
|
|
87
|
+
assert r["kept"]==1
|
|
88
|
+
def test_no_dup(self):
|
|
89
|
+
s=ContextGuard.create_store()
|
|
90
|
+
texts = ["completely different", "totally unrelated", "no similarity here"]
|
|
91
|
+
r=ContextGuard.deduplicate(s, texts)
|
|
92
|
+
assert r["kept"]==3
|
|
93
|
+
def test_empty(self):
|
|
94
|
+
s=ContextGuard.create_store()
|
|
95
|
+
r=ContextGuard.deduplicate(s, [])
|
|
96
|
+
assert r["kept"]==0
|
|
97
|
+
|
|
98
|
+
class TestExtractKey:
|
|
99
|
+
def test_basic(self):
|
|
100
|
+
s=ContextGuard.create_store()
|
|
101
|
+
text = "Python is great. The weather is nice. Machine learning uses Python. I like pizza. Data science requires Python skills."
|
|
102
|
+
r=ContextGuard.extract_key(s, text, num_keys=2)
|
|
103
|
+
assert r["count"]==2
|
|
104
|
+
# Python sentences should score higher due to frequency
|
|
105
|
+
def test_empty(self):
|
|
106
|
+
s=ContextGuard.create_store()
|
|
107
|
+
r=ContextGuard.extract_key(s, "", num_keys=3)
|
|
108
|
+
assert r["count"]==0
|
|
109
|
+
|
|
110
|
+
class TestTruncateSmart:
|
|
111
|
+
def test_short(self):
|
|
112
|
+
s=ContextGuard.create_store()
|
|
113
|
+
r=ContextGuard.truncate_smart(s, "Hello world.", max_tokens=100)
|
|
114
|
+
assert r["truncated"] is False
|
|
115
|
+
def test_long(self):
|
|
116
|
+
s=ContextGuard.create_store()
|
|
117
|
+
text = ". ".join(f"Sentence {i}" for i in range(100))
|
|
118
|
+
r=ContextGuard.truncate_smart(s, text, max_tokens=30)
|
|
119
|
+
assert r["truncated"] is True
|
|
120
|
+
assert r["saved"]>0
|
|
121
|
+
|
|
122
|
+
class TestChunk:
|
|
123
|
+
def test_basic(self):
|
|
124
|
+
s=ContextGuard.create_store()
|
|
125
|
+
text = ". ".join(f"Sentence {i}" for i in range(20))
|
|
126
|
+
r=ContextGuard.chunk(s, text, max_tokens=30, overlap=0)
|
|
127
|
+
assert r["count"]>1
|
|
128
|
+
assert all(c for c in r["chunks"])
|
|
129
|
+
def test_overlap(self):
|
|
130
|
+
s=ContextGuard.create_store()
|
|
131
|
+
text = ". ".join(f"Sentence {i}" for i in range(20))
|
|
132
|
+
r=ContextGuard.chunk(s, text, max_tokens=30, overlap=5)
|
|
133
|
+
assert r["count"]>1
|
|
134
|
+
|
|
135
|
+
class TestSummarizeHistory:
|
|
136
|
+
def test_basic(self):
|
|
137
|
+
s=ContextGuard.create_store()
|
|
138
|
+
msgs = [
|
|
139
|
+
{"role": "user", "content": "How do I write a Python function?"},
|
|
140
|
+
{"role": "assistant", "content": "You can use def keyword to define a function in Python."},
|
|
141
|
+
{"role": "user", "content": "What about classes?"},
|
|
142
|
+
{"role": "assistant", "content": "Use class keyword to define a class."},
|
|
143
|
+
]
|
|
144
|
+
r=ContextGuard.summarize_history(s, msgs, max_tokens=50)
|
|
145
|
+
assert r["messages_processed"]==4
|
|
146
|
+
assert len(r["text"])>0
|
|
147
|
+
|
|
148
|
+
class TestFilterRelevant:
|
|
149
|
+
def test_basic(self):
|
|
150
|
+
s=ContextGuard.create_store()
|
|
151
|
+
query = "Python function"
|
|
152
|
+
passages = [
|
|
153
|
+
"Python functions are defined with def",
|
|
154
|
+
"Java classes use different syntax",
|
|
155
|
+
"def is a Python keyword for functions",
|
|
156
|
+
"The weather is nice today",
|
|
157
|
+
"Python also has lambda functions",
|
|
158
|
+
]
|
|
159
|
+
r=ContextGuard.filter_relevant(s, query, passages, top_k=3)
|
|
160
|
+
assert r["count"]<=3
|
|
161
|
+
assert r["results"][0]["score"]>0
|
|
162
|
+
# Python passages should rank higher
|
|
163
|
+
def test_no_match(self):
|
|
164
|
+
s=ContextGuard.create_store()
|
|
165
|
+
r=ContextGuard.filter_relevant(s, "xyzabc", ["hello world", "foo bar"], top_k=5)
|
|
166
|
+
assert r["count"]==0
|
|
167
|
+
def test_empty(self):
|
|
168
|
+
s=ContextGuard.create_store()
|
|
169
|
+
r=ContextGuard.filter_relevant(s, "test", [], top_k=5)
|
|
170
|
+
assert r["count"]==0
|
|
171
|
+
|
|
172
|
+
class TestMergeContext:
|
|
173
|
+
def test_basic(self):
|
|
174
|
+
s=ContextGuard.create_store()
|
|
175
|
+
sources = [
|
|
176
|
+
"Python is a programming language.",
|
|
177
|
+
"Python is used for data science.",
|
|
178
|
+
"Java is also a programming language.",
|
|
179
|
+
]
|
|
180
|
+
r=ContextGuard.merge_context(s, sources, max_tokens=50)
|
|
181
|
+
assert r["sources_merged"]>=1
|
|
182
|
+
assert len(r["text"])>0
|
|
183
|
+
def test_dedup(self):
|
|
184
|
+
s=ContextGuard.create_store()
|
|
185
|
+
sources = ["hello world", "hello world"]
|
|
186
|
+
r=ContextGuard.merge_context(s, sources, max_tokens=100, deduplicate=True)
|
|
187
|
+
assert r["sources_merged"]==1
|
|
188
|
+
|
|
189
|
+
class TestStatsReset:
|
|
190
|
+
def test_stats(self):
|
|
191
|
+
s=ContextGuard.create_store()
|
|
192
|
+
ContextGuard.compress(s, "test text here now")
|
|
193
|
+
r=ContextGuard.get_stats(s)
|
|
194
|
+
assert r["compressed"]==1
|
|
195
|
+
def test_reset(self):
|
|
196
|
+
s=ContextGuard.create_store()
|
|
197
|
+
ContextGuard.compress(s, "test text here now")
|
|
198
|
+
r=ContextGuard.reset(s)
|
|
199
|
+
assert r["reset"]["compressed"]==1
|
|
200
|
+
assert ContextGuard.get_stats(s)["compressed"]==0
|
|
201
|
+
|
|
202
|
+
class TestDispatch:
|
|
203
|
+
def test_unknown(self):
|
|
204
|
+
s=MCPContextGuardServer();assert "error" in json.loads(s.handle_tool_call("nope",{}))
|
|
205
|
+
def test_missing(self):
|
|
206
|
+
s=MCPContextGuardServer();assert "error" in json.loads(s.handle_tool_call("compress",{}))
|
|
207
|
+
def test_compress_dispatch(self):
|
|
208
|
+
s=MCPContextGuardServer()
|
|
209
|
+
r=json.loads(s.handle_tool_call("compress",{"text":"hello world"}))
|
|
210
|
+
assert r["success"] is True
|
|
211
|
+
|
|
212
|
+
class TestSTDIO:
|
|
213
|
+
def test_manifest_flag(self,capsys):
|
|
214
|
+
from src.server import main
|
|
215
|
+
with patch("sys.argv",["server","--manifest"]):main()
|
|
216
|
+
parsed=json.loads(capsys.readouterr().out.strip())
|
|
217
|
+
assert parsed["server"]["name"]=="mcp-context-guard"
|