agent-vault-mcp 0.2.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.
- agent_vault_mcp-0.2.0/PKG-INFO +131 -0
- agent_vault_mcp-0.2.0/README.md +123 -0
- agent_vault_mcp-0.2.0/agent_vault/__init__.py +0 -0
- agent_vault_mcp-0.2.0/agent_vault/core/__init__.py +0 -0
- agent_vault_mcp-0.2.0/agent_vault/core/dsa.py +71 -0
- agent_vault_mcp-0.2.0/agent_vault/core/storage.py +343 -0
- agent_vault_mcp-0.2.0/agent_vault/server.py +144 -0
- agent_vault_mcp-0.2.0/agent_vault_mcp.egg-info/PKG-INFO +131 -0
- agent_vault_mcp-0.2.0/agent_vault_mcp.egg-info/SOURCES.txt +14 -0
- agent_vault_mcp-0.2.0/agent_vault_mcp.egg-info/dependency_links.txt +1 -0
- agent_vault_mcp-0.2.0/agent_vault_mcp.egg-info/entry_points.txt +2 -0
- agent_vault_mcp-0.2.0/agent_vault_mcp.egg-info/requires.txt +1 -0
- agent_vault_mcp-0.2.0/agent_vault_mcp.egg-info/top_level.txt +1 -0
- agent_vault_mcp-0.2.0/pyproject.toml +12 -0
- agent_vault_mcp-0.2.0/setup.cfg +4 -0
- agent_vault_mcp-0.2.0/tests/test_prompt_cache.py +53 -0
|
@@ -0,0 +1,131 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: agent-vault-mcp
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: DSA-Optimized AI Knowledge Bank MCP Server
|
|
5
|
+
Requires-Python: >=3.10
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
Requires-Dist: mcp<3,>=2
|
|
8
|
+
|
|
9
|
+
# Agent Vault MCP
|
|
10
|
+
|
|
11
|
+
A blazing-fast, token-efficient AI Knowledge Bank and Memory Server built on the **Model Context Protocol (MCP)**.
|
|
12
|
+
|
|
13
|
+
Agent Vault is designed to give AI coding agents (like Claude Code, Antigravity, and Cursor) persistent, cross-session memory without blowing up token limits. It acts as an intelligent file digest cache and semantic knowledge base.
|
|
14
|
+
|
|
15
|
+
## Overview
|
|
16
|
+
AI agents typically start every session with amnesia, forcing them to re-read thousands of lines of code. Agent Vault solves this using a **Two-Step Funnel**:
|
|
17
|
+
1. **Discovery**: The agent searches the Vault to find out *which* files matter.
|
|
18
|
+
2. **Execution**: The agent reads only the raw code of those specific files, edits them, and updates the vault.
|
|
19
|
+
|
|
20
|
+
When the agent wants to check a file, the Vault hashes it (SHA-256). If it hasn't changed, the Vault instantly returns a cached AST/summary instead of making the agent read the entire file.
|
|
21
|
+
|
|
22
|
+
## Features
|
|
23
|
+
- **$O(1)$ Fast Rejection**: Uses Bloom Filters to instantly know if a file is unindexed.
|
|
24
|
+
- **Token-Bounded MinHeap**: Ranks the best context snippets and strictly cuts off when the maximum token limit is reached, protecting the context window.
|
|
25
|
+
- **SQLite FTS5 (BM25)**: Fast lexical and semantic search for symbols, errors, and flows.
|
|
26
|
+
- **ROI Telemetry**: Natively calculates and tracks how many tokens and hours of inference time are saved by skipping raw file reads.
|
|
27
|
+
- **100% Portable**: Caches are stored using relative paths, meaning you can move or rename your project folder without breaking the vault.
|
|
28
|
+
|
|
29
|
+
## Dependencies
|
|
30
|
+
- Python 3.10+
|
|
31
|
+
- `mcp` (Official Model Context Protocol SDK, `mcp>=2,<3`)
|
|
32
|
+
- Standard library components (`sqlite3`, `hashlib`, `heapq`)
|
|
33
|
+
|
|
34
|
+
## Installation
|
|
35
|
+
|
|
36
|
+
### Option 1: Global Python (3.10+)
|
|
37
|
+
```bash
|
|
38
|
+
# Clone or navigate to the repository
|
|
39
|
+
cd agent-vault-mcp
|
|
40
|
+
|
|
41
|
+
# Install the package
|
|
42
|
+
pip install -e .
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
### Option 2: Using `uv` (Virtual Environment)
|
|
46
|
+
```bash
|
|
47
|
+
uv venv --python 3.12
|
|
48
|
+
.venv\Scripts\activate # On Windows
|
|
49
|
+
uv pip install -e .
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
## How to use with Antigravity / Claude Code
|
|
53
|
+
|
|
54
|
+
To add this to an MCP client like Antigravity, add the following to your `mcp_config.json` (e.g., `~/.gemini/config/mcp_config.json`):
|
|
55
|
+
|
|
56
|
+
```json
|
|
57
|
+
{
|
|
58
|
+
"mcpServers": {
|
|
59
|
+
"agent-vault": {
|
|
60
|
+
"command": "python",
|
|
61
|
+
"args": ["-m", "agent_vault.server"],
|
|
62
|
+
"env": {
|
|
63
|
+
"PYTHONPATH": "/absolute/path/to/agent-vault-mcp"
|
|
64
|
+
}
|
|
65
|
+
}
|
|
66
|
+
}
|
|
67
|
+
}
|
|
68
|
+
```
|
|
69
|
+
*Note: If you used a virtual environment, change `"command": "python"` to the absolute path of the python executable inside `.venv`.*
|
|
70
|
+
|
|
71
|
+
### Adoption (Forcing the AI to use it)
|
|
72
|
+
The Vault only saves tokens if the AI remembers to use it! Add this snippet to your project's `CLAUDE.md`, `.cursorrules`, or `GEMINI.md`:
|
|
73
|
+
|
|
74
|
+
```markdown
|
|
75
|
+
# 🧠Agent Vault & Memory Protocol
|
|
76
|
+
You are equipped with the Agent Vault MCP Server. To protect the user's token limits and eliminate hallucination, you MUST strictly adhere to the following workflow:
|
|
77
|
+
|
|
78
|
+
### 1. The "Think Before You Read" Rule (Semantic Cache)
|
|
79
|
+
Before you spend time reading files to understand an architecture, flow, or system (e.g., "How does auth work?"):
|
|
80
|
+
* **ALWAYS** call `vault_search_questions(query)` using keywords/tags to see if a previous agent already solved this.
|
|
81
|
+
* If you find a match, call `vault_search_answer(prompt)` to retrieve the pre-computed answer instantly.
|
|
82
|
+
|
|
83
|
+
### 2. The "Read Before You Write" Rule (File Cache)
|
|
84
|
+
Before you execute commands to read raw code files:
|
|
85
|
+
* **ALWAYS** call `vault_check_file(filepath)` first.
|
|
86
|
+
* If the Vault returns a Cache Hit, trust the summary/AST and DO NOT read the raw file unless you explicitly need to edit it.
|
|
87
|
+
|
|
88
|
+
### 3. The "Leave It Better Than You Found It" Rule (Updating Cache)
|
|
89
|
+
Your memory is only as good as what you save. After you complete a task:
|
|
90
|
+
* **Cache Modified Files:** If you edited a file, ALWAYS call `vault_cache_file(filepath, summary)` to update its digest and AST.
|
|
91
|
+
* **Cache New Knowledge:** If you just spent time analyzing a complex architecture or debugging a hard issue, ALWAYS call `vault_cache_answer(prompt, response, dependencies, tags)`.
|
|
92
|
+
* *Dependencies:* You MUST provide the exact file paths your answer relies on so the Vault can auto-invalidate your answer if those files change.
|
|
93
|
+
* *Tags:* Provide 5-6 broad keyword tags (e.g., "auth, login, jwt") so future agents can easily discover your answer via `vault_search_questions`.
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
## Local vs Global Vaults
|
|
97
|
+
|
|
98
|
+
By default, the MCP server creates `agent_vault.db` inside your current active project workspace. Paths are stored **relatively**. This means if you ask the agent about a file in an external project (e.g., `../Project_B/main.py`), that cross-project memory is stored locally inside your current project's database.
|
|
99
|
+
|
|
100
|
+
If you prefer a **"Global Brain"** that shares all memories and file caches across every single project on your computer, simply add `AGENT_VAULT_DB_PATH` to the `env` variables in your `mcp_config.json`:
|
|
101
|
+
|
|
102
|
+
```json
|
|
103
|
+
"env": {
|
|
104
|
+
"PYTHONPATH": "/absolute/path/to/agent-vault-mcp",
|
|
105
|
+
"AGENT_VAULT_DB_PATH": "/absolute/path/to/.global_agent_vault.db"
|
|
106
|
+
}
|
|
107
|
+
```
|
|
108
|
+
|
|
109
|
+
## Available MCP Tools
|
|
110
|
+
- `vault_store_memory(key, content, tags)`: Save arbitrary architectural notes or debugging insights.
|
|
111
|
+
- `vault_search(query, max_tokens)`: Search the vault using BM25 ranking.
|
|
112
|
+
- `vault_cache_file(filepath, summary)`: Hash a file and cache its summary.
|
|
113
|
+
- `vault_check_file(filepath)`: Verify a file's hash and return its cached summary + telemetry metrics.
|
|
114
|
+
- `vault_stats()`: View the ROI dashboard of tokens and time saved.
|
|
115
|
+
- `vault_delete_memory(key)`: Delete a stored memory.
|
|
116
|
+
- `vault_evict_file(filepath)`: Evict a file from the vault cache.- `vault_cache_answer(prompt, response, dependencies, tags)`: Save an AI-generated answer linked to specific files and keywords.
|
|
117
|
+
- `vault_search_questions(query, max_results)`: Keyword-search an FTS5 index to find exactly how previous cached questions were phrased based on tags.
|
|
118
|
+
- `vault_search_answer(prompt)`: Retrieve a cached AI answer (automatically invalidates if dependent files have changed).
|
|
119
|
+
|
|
120
|
+
## Semantic Prompt Caching (Intent Caching)
|
|
121
|
+
|
|
122
|
+
In `v0.2.0`, Agent Vault introduced **Semantic Prompt Caching**. Instead of forcing AI agents to repeatedly read files and re-reason through complex architectural questions (e.g., *"How does the auth flow work?"*), agents can now cache their reasoning.
|
|
123
|
+
|
|
124
|
+
### Dependency Invalidation
|
|
125
|
+
To solve the classic LLM problem of "stale context hallucination," Agent Vault uses **Deterministic Dependency Invalidation**. When an agent caches an answer, it explicitly lists the files that answer depends on.
|
|
126
|
+
|
|
127
|
+
When a future agent asks the same question, the Vault calculates the real-time SHA-256 digest of those dependencies. If any file has changed, the cached answer is instantly evicted, forcing the AI to generate a fresh, accurate response.
|
|
128
|
+
|
|
129
|
+
### N-to-1 Tag Mapping
|
|
130
|
+
To solve the problem of "brittle exact matching" (where *"How does login work?"* misses a cache for *"How does the login work?"*), agents can assign **tags** to cached answers.
|
|
131
|
+
Agents can use `vault_search_questions("login")` to hit the FTS5 index, discover the exact phrasing of the cached question, and then fetch the answer—bypassing the need for heavy vector databases!
|
|
@@ -0,0 +1,123 @@
|
|
|
1
|
+
# Agent Vault MCP
|
|
2
|
+
|
|
3
|
+
A blazing-fast, token-efficient AI Knowledge Bank and Memory Server built on the **Model Context Protocol (MCP)**.
|
|
4
|
+
|
|
5
|
+
Agent Vault is designed to give AI coding agents (like Claude Code, Antigravity, and Cursor) persistent, cross-session memory without blowing up token limits. It acts as an intelligent file digest cache and semantic knowledge base.
|
|
6
|
+
|
|
7
|
+
## Overview
|
|
8
|
+
AI agents typically start every session with amnesia, forcing them to re-read thousands of lines of code. Agent Vault solves this using a **Two-Step Funnel**:
|
|
9
|
+
1. **Discovery**: The agent searches the Vault to find out *which* files matter.
|
|
10
|
+
2. **Execution**: The agent reads only the raw code of those specific files, edits them, and updates the vault.
|
|
11
|
+
|
|
12
|
+
When the agent wants to check a file, the Vault hashes it (SHA-256). If it hasn't changed, the Vault instantly returns a cached AST/summary instead of making the agent read the entire file.
|
|
13
|
+
|
|
14
|
+
## Features
|
|
15
|
+
- **$O(1)$ Fast Rejection**: Uses Bloom Filters to instantly know if a file is unindexed.
|
|
16
|
+
- **Token-Bounded MinHeap**: Ranks the best context snippets and strictly cuts off when the maximum token limit is reached, protecting the context window.
|
|
17
|
+
- **SQLite FTS5 (BM25)**: Fast lexical and semantic search for symbols, errors, and flows.
|
|
18
|
+
- **ROI Telemetry**: Natively calculates and tracks how many tokens and hours of inference time are saved by skipping raw file reads.
|
|
19
|
+
- **100% Portable**: Caches are stored using relative paths, meaning you can move or rename your project folder without breaking the vault.
|
|
20
|
+
|
|
21
|
+
## Dependencies
|
|
22
|
+
- Python 3.10+
|
|
23
|
+
- `mcp` (Official Model Context Protocol SDK, `mcp>=2,<3`)
|
|
24
|
+
- Standard library components (`sqlite3`, `hashlib`, `heapq`)
|
|
25
|
+
|
|
26
|
+
## Installation
|
|
27
|
+
|
|
28
|
+
### Option 1: Global Python (3.10+)
|
|
29
|
+
```bash
|
|
30
|
+
# Clone or navigate to the repository
|
|
31
|
+
cd agent-vault-mcp
|
|
32
|
+
|
|
33
|
+
# Install the package
|
|
34
|
+
pip install -e .
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
### Option 2: Using `uv` (Virtual Environment)
|
|
38
|
+
```bash
|
|
39
|
+
uv venv --python 3.12
|
|
40
|
+
.venv\Scripts\activate # On Windows
|
|
41
|
+
uv pip install -e .
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
## How to use with Antigravity / Claude Code
|
|
45
|
+
|
|
46
|
+
To add this to an MCP client like Antigravity, add the following to your `mcp_config.json` (e.g., `~/.gemini/config/mcp_config.json`):
|
|
47
|
+
|
|
48
|
+
```json
|
|
49
|
+
{
|
|
50
|
+
"mcpServers": {
|
|
51
|
+
"agent-vault": {
|
|
52
|
+
"command": "python",
|
|
53
|
+
"args": ["-m", "agent_vault.server"],
|
|
54
|
+
"env": {
|
|
55
|
+
"PYTHONPATH": "/absolute/path/to/agent-vault-mcp"
|
|
56
|
+
}
|
|
57
|
+
}
|
|
58
|
+
}
|
|
59
|
+
}
|
|
60
|
+
```
|
|
61
|
+
*Note: If you used a virtual environment, change `"command": "python"` to the absolute path of the python executable inside `.venv`.*
|
|
62
|
+
|
|
63
|
+
### Adoption (Forcing the AI to use it)
|
|
64
|
+
The Vault only saves tokens if the AI remembers to use it! Add this snippet to your project's `CLAUDE.md`, `.cursorrules`, or `GEMINI.md`:
|
|
65
|
+
|
|
66
|
+
```markdown
|
|
67
|
+
# 🧠Agent Vault & Memory Protocol
|
|
68
|
+
You are equipped with the Agent Vault MCP Server. To protect the user's token limits and eliminate hallucination, you MUST strictly adhere to the following workflow:
|
|
69
|
+
|
|
70
|
+
### 1. The "Think Before You Read" Rule (Semantic Cache)
|
|
71
|
+
Before you spend time reading files to understand an architecture, flow, or system (e.g., "How does auth work?"):
|
|
72
|
+
* **ALWAYS** call `vault_search_questions(query)` using keywords/tags to see if a previous agent already solved this.
|
|
73
|
+
* If you find a match, call `vault_search_answer(prompt)` to retrieve the pre-computed answer instantly.
|
|
74
|
+
|
|
75
|
+
### 2. The "Read Before You Write" Rule (File Cache)
|
|
76
|
+
Before you execute commands to read raw code files:
|
|
77
|
+
* **ALWAYS** call `vault_check_file(filepath)` first.
|
|
78
|
+
* If the Vault returns a Cache Hit, trust the summary/AST and DO NOT read the raw file unless you explicitly need to edit it.
|
|
79
|
+
|
|
80
|
+
### 3. The "Leave It Better Than You Found It" Rule (Updating Cache)
|
|
81
|
+
Your memory is only as good as what you save. After you complete a task:
|
|
82
|
+
* **Cache Modified Files:** If you edited a file, ALWAYS call `vault_cache_file(filepath, summary)` to update its digest and AST.
|
|
83
|
+
* **Cache New Knowledge:** If you just spent time analyzing a complex architecture or debugging a hard issue, ALWAYS call `vault_cache_answer(prompt, response, dependencies, tags)`.
|
|
84
|
+
* *Dependencies:* You MUST provide the exact file paths your answer relies on so the Vault can auto-invalidate your answer if those files change.
|
|
85
|
+
* *Tags:* Provide 5-6 broad keyword tags (e.g., "auth, login, jwt") so future agents can easily discover your answer via `vault_search_questions`.
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
## Local vs Global Vaults
|
|
89
|
+
|
|
90
|
+
By default, the MCP server creates `agent_vault.db` inside your current active project workspace. Paths are stored **relatively**. This means if you ask the agent about a file in an external project (e.g., `../Project_B/main.py`), that cross-project memory is stored locally inside your current project's database.
|
|
91
|
+
|
|
92
|
+
If you prefer a **"Global Brain"** that shares all memories and file caches across every single project on your computer, simply add `AGENT_VAULT_DB_PATH` to the `env` variables in your `mcp_config.json`:
|
|
93
|
+
|
|
94
|
+
```json
|
|
95
|
+
"env": {
|
|
96
|
+
"PYTHONPATH": "/absolute/path/to/agent-vault-mcp",
|
|
97
|
+
"AGENT_VAULT_DB_PATH": "/absolute/path/to/.global_agent_vault.db"
|
|
98
|
+
}
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
## Available MCP Tools
|
|
102
|
+
- `vault_store_memory(key, content, tags)`: Save arbitrary architectural notes or debugging insights.
|
|
103
|
+
- `vault_search(query, max_tokens)`: Search the vault using BM25 ranking.
|
|
104
|
+
- `vault_cache_file(filepath, summary)`: Hash a file and cache its summary.
|
|
105
|
+
- `vault_check_file(filepath)`: Verify a file's hash and return its cached summary + telemetry metrics.
|
|
106
|
+
- `vault_stats()`: View the ROI dashboard of tokens and time saved.
|
|
107
|
+
- `vault_delete_memory(key)`: Delete a stored memory.
|
|
108
|
+
- `vault_evict_file(filepath)`: Evict a file from the vault cache.- `vault_cache_answer(prompt, response, dependencies, tags)`: Save an AI-generated answer linked to specific files and keywords.
|
|
109
|
+
- `vault_search_questions(query, max_results)`: Keyword-search an FTS5 index to find exactly how previous cached questions were phrased based on tags.
|
|
110
|
+
- `vault_search_answer(prompt)`: Retrieve a cached AI answer (automatically invalidates if dependent files have changed).
|
|
111
|
+
|
|
112
|
+
## Semantic Prompt Caching (Intent Caching)
|
|
113
|
+
|
|
114
|
+
In `v0.2.0`, Agent Vault introduced **Semantic Prompt Caching**. Instead of forcing AI agents to repeatedly read files and re-reason through complex architectural questions (e.g., *"How does the auth flow work?"*), agents can now cache their reasoning.
|
|
115
|
+
|
|
116
|
+
### Dependency Invalidation
|
|
117
|
+
To solve the classic LLM problem of "stale context hallucination," Agent Vault uses **Deterministic Dependency Invalidation**. When an agent caches an answer, it explicitly lists the files that answer depends on.
|
|
118
|
+
|
|
119
|
+
When a future agent asks the same question, the Vault calculates the real-time SHA-256 digest of those dependencies. If any file has changed, the cached answer is instantly evicted, forcing the AI to generate a fresh, accurate response.
|
|
120
|
+
|
|
121
|
+
### N-to-1 Tag Mapping
|
|
122
|
+
To solve the problem of "brittle exact matching" (where *"How does login work?"* misses a cache for *"How does the login work?"*), agents can assign **tags** to cached answers.
|
|
123
|
+
Agents can use `vault_search_questions("login")` to hit the FTS5 index, discover the exact phrasing of the cached question, and then fetch the answer—bypassing the need for heavy vector databases!
|
|
File without changes
|
|
File without changes
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
import hashlib
|
|
2
|
+
import os
|
|
3
|
+
import heapq
|
|
4
|
+
|
|
5
|
+
class BloomFilter:
|
|
6
|
+
def __init__(self, size=1000000, hash_count=5):
|
|
7
|
+
self.size = size
|
|
8
|
+
self.hash_count = hash_count
|
|
9
|
+
self.bit_array = [False] * size
|
|
10
|
+
|
|
11
|
+
def _hashes(self, item):
|
|
12
|
+
# Simple string-based hashing for simplicity and to avoid external dependencies
|
|
13
|
+
hashes = []
|
|
14
|
+
for i in range(self.hash_count):
|
|
15
|
+
# Using md5 for fast hashing across multiple seeds
|
|
16
|
+
h = int(hashlib.md5(f"{item}:{i}".encode('utf-8')).hexdigest(), 16)
|
|
17
|
+
hashes.append(h % self.size)
|
|
18
|
+
return hashes
|
|
19
|
+
|
|
20
|
+
def add(self, item):
|
|
21
|
+
for h in self._hashes(item):
|
|
22
|
+
self.bit_array[h] = True
|
|
23
|
+
|
|
24
|
+
def check(self, item):
|
|
25
|
+
for h in self._hashes(item):
|
|
26
|
+
if not self.bit_array[h]:
|
|
27
|
+
return False
|
|
28
|
+
return True
|
|
29
|
+
|
|
30
|
+
class TokenBoundedMinHeap:
|
|
31
|
+
def __init__(self, max_tokens):
|
|
32
|
+
self.max_tokens = max_tokens
|
|
33
|
+
self.current_tokens = 0
|
|
34
|
+
self.heap = [] # stores tuples of (score, tiebreaker, tokens, item)
|
|
35
|
+
self.tiebreaker = 0
|
|
36
|
+
|
|
37
|
+
def add(self, item, score, tokens):
|
|
38
|
+
if tokens > self.max_tokens:
|
|
39
|
+
return
|
|
40
|
+
|
|
41
|
+
# We want to keep items with HIGHEST score.
|
|
42
|
+
# Min-heap pops the SMALLEST score first.
|
|
43
|
+
# So we push (score, ...), and pop when current_tokens > max_tokens.
|
|
44
|
+
# This will remove the lowest-scored items.
|
|
45
|
+
heapq.heappush(self.heap, (score, self.tiebreaker, tokens, item))
|
|
46
|
+
self.tiebreaker += 1
|
|
47
|
+
self.current_tokens += tokens
|
|
48
|
+
|
|
49
|
+
while self.current_tokens > self.max_tokens and self.heap:
|
|
50
|
+
popped = heapq.heappop(self.heap)
|
|
51
|
+
self.current_tokens -= popped[2]
|
|
52
|
+
|
|
53
|
+
def get_items(self):
|
|
54
|
+
# Return items sorted by score descending (highest score first)
|
|
55
|
+
sorted_items = sorted(self.heap, key=lambda x: x[0], reverse=True)
|
|
56
|
+
return [x[3] for x in sorted_items]
|
|
57
|
+
|
|
58
|
+
def get_file_digest(filepath):
|
|
59
|
+
"""Returns SHA-256 digest of a file, and its size, or (None, 0) if it doesn't exist."""
|
|
60
|
+
if not os.path.exists(filepath):
|
|
61
|
+
return None, 0
|
|
62
|
+
hasher = hashlib.sha256()
|
|
63
|
+
size_in_bytes = 0
|
|
64
|
+
try:
|
|
65
|
+
with open(filepath, 'rb') as f:
|
|
66
|
+
for chunk in iter(lambda: f.read(65536), b''):
|
|
67
|
+
hasher.update(chunk)
|
|
68
|
+
size_in_bytes += len(chunk)
|
|
69
|
+
return hasher.hexdigest(), size_in_bytes
|
|
70
|
+
except Exception:
|
|
71
|
+
return None, 0
|
|
@@ -0,0 +1,343 @@
|
|
|
1
|
+
import sqlite3
|
|
2
|
+
import os
|
|
3
|
+
import json
|
|
4
|
+
from .dsa import BloomFilter, get_file_digest, TokenBoundedMinHeap
|
|
5
|
+
|
|
6
|
+
class VaultStorage:
|
|
7
|
+
def __init__(self, db_path="vault.db"):
|
|
8
|
+
self.db_path = db_path
|
|
9
|
+
self.conn = sqlite3.connect(db_path, check_same_thread=False)
|
|
10
|
+
try:
|
|
11
|
+
os.chmod(self.db_path, 0o600)
|
|
12
|
+
except OSError:
|
|
13
|
+
pass
|
|
14
|
+
self.conn.execute("PRAGMA journal_mode=WAL;")
|
|
15
|
+
self.conn.execute("PRAGMA busy_timeout=5000;")
|
|
16
|
+
self.bloom_filter = BloomFilter()
|
|
17
|
+
self._init_db()
|
|
18
|
+
|
|
19
|
+
def _init_db(self):
|
|
20
|
+
with self.conn:
|
|
21
|
+
# FTS5 table for memory snippets
|
|
22
|
+
self.conn.execute('''
|
|
23
|
+
CREATE VIRTUAL TABLE IF NOT EXISTS memories USING fts5(
|
|
24
|
+
key, content, tags, tokens
|
|
25
|
+
)
|
|
26
|
+
''')
|
|
27
|
+
# Table for file cache
|
|
28
|
+
self.conn.execute('''
|
|
29
|
+
CREATE TABLE IF NOT EXISTS file_cache (
|
|
30
|
+
filepath TEXT PRIMARY KEY,
|
|
31
|
+
digest TEXT,
|
|
32
|
+
summary TEXT,
|
|
33
|
+
tokens_saved_per_hit INTEGER DEFAULT 0
|
|
34
|
+
)
|
|
35
|
+
''')
|
|
36
|
+
# Handle schema migration for file_cache safely
|
|
37
|
+
cursor = self.conn.execute("PRAGMA table_info(file_cache)")
|
|
38
|
+
columns = [col[1] for col in cursor.fetchall()]
|
|
39
|
+
if 'tokens_saved_per_hit' not in columns:
|
|
40
|
+
self.conn.execute('ALTER TABLE file_cache ADD COLUMN tokens_saved_per_hit INTEGER DEFAULT 0')
|
|
41
|
+
|
|
42
|
+
# Metrics table
|
|
43
|
+
self.conn.execute('''
|
|
44
|
+
CREATE TABLE IF NOT EXISTS metrics (
|
|
45
|
+
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
46
|
+
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
47
|
+
filepath TEXT,
|
|
48
|
+
is_hit BOOLEAN,
|
|
49
|
+
tokens_saved INTEGER,
|
|
50
|
+
reason TEXT
|
|
51
|
+
)
|
|
52
|
+
''')
|
|
53
|
+
# Handle schema migration for metrics safely
|
|
54
|
+
cursor = self.conn.execute("PRAGMA table_info(metrics)")
|
|
55
|
+
columns = [col[1] for col in cursor.fetchall()]
|
|
56
|
+
if 'reason' not in columns:
|
|
57
|
+
self.conn.execute('ALTER TABLE metrics ADD COLUMN reason TEXT')
|
|
58
|
+
|
|
59
|
+
# Prompt cache table
|
|
60
|
+
self.conn.execute('''
|
|
61
|
+
CREATE TABLE IF NOT EXISTS prompt_cache (
|
|
62
|
+
id INTEGER PRIMARY KEY,
|
|
63
|
+
query_hash TEXT UNIQUE,
|
|
64
|
+
response TEXT,
|
|
65
|
+
dependency_files TEXT
|
|
66
|
+
)
|
|
67
|
+
''')
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
self.conn.execute('''
|
|
71
|
+
CREATE TABLE IF NOT EXISTS prompt_cache (
|
|
72
|
+
id INTEGER PRIMARY KEY,
|
|
73
|
+
query_hash TEXT UNIQUE,
|
|
74
|
+
prompt TEXT,
|
|
75
|
+
response TEXT,
|
|
76
|
+
dependency_files TEXT,
|
|
77
|
+
tags TEXT
|
|
78
|
+
)
|
|
79
|
+
''')
|
|
80
|
+
# Handle migration
|
|
81
|
+
cursor = self.conn.execute("PRAGMA table_info(prompt_cache)")
|
|
82
|
+
columns = [col[1] for col in cursor.fetchall()]
|
|
83
|
+
if 'prompt' not in columns:
|
|
84
|
+
self.conn.execute('ALTER TABLE prompt_cache ADD COLUMN prompt TEXT')
|
|
85
|
+
if 'tags' not in columns:
|
|
86
|
+
self.conn.execute('ALTER TABLE prompt_cache ADD COLUMN tags TEXT')
|
|
87
|
+
|
|
88
|
+
self.conn.execute('''
|
|
89
|
+
CREATE VIRTUAL TABLE IF NOT EXISTS prompt_search USING fts5(
|
|
90
|
+
prompt, tags, query_hash UNINDEXED
|
|
91
|
+
)
|
|
92
|
+
''')
|
|
93
|
+
|
|
94
|
+
# Hydrate bloom filter with existing files
|
|
95
|
+
try:
|
|
96
|
+
cursor = self.conn.execute("SELECT filepath FROM file_cache")
|
|
97
|
+
for row in cursor:
|
|
98
|
+
self.bloom_filter.add(row[0])
|
|
99
|
+
except Exception:
|
|
100
|
+
pass
|
|
101
|
+
|
|
102
|
+
def store_memory(self, key, content, tags, tokens=None):
|
|
103
|
+
if tokens is None:
|
|
104
|
+
# simple token estimation (1 token ~= 4 chars)
|
|
105
|
+
tokens = len(content) // 4 + 1
|
|
106
|
+
|
|
107
|
+
with self.conn:
|
|
108
|
+
# Delete if exists to update
|
|
109
|
+
self.conn.execute("DELETE FROM memories WHERE key = ?", (key,))
|
|
110
|
+
self.conn.execute(
|
|
111
|
+
"INSERT INTO memories (key, content, tags, tokens) VALUES (?, ?, ?, ?)",
|
|
112
|
+
(key, content, tags, tokens)
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
def search_memory(self, query, max_tokens=2000):
|
|
116
|
+
# We use BM25 which is built into FTS5 via ORDER BY rank
|
|
117
|
+
# Standard FTS5 match query
|
|
118
|
+
|
|
119
|
+
# Sanitize query to avoid FTS syntax errors on special chars
|
|
120
|
+
safe_query = ''.join(c if c.isalnum() or c.isspace() else ' ' for c in query).strip()
|
|
121
|
+
if not safe_query:
|
|
122
|
+
return []
|
|
123
|
+
|
|
124
|
+
try:
|
|
125
|
+
cursor = self.conn.execute('''
|
|
126
|
+
SELECT key, content, tags, tokens, rank
|
|
127
|
+
FROM memories
|
|
128
|
+
WHERE memories MATCH ?
|
|
129
|
+
ORDER BY rank
|
|
130
|
+
''', (safe_query,))
|
|
131
|
+
except sqlite3.OperationalError:
|
|
132
|
+
# fallback if query is still invalid
|
|
133
|
+
return []
|
|
134
|
+
|
|
135
|
+
# Rank in FTS5 is more negative for better matches.
|
|
136
|
+
# Let's negate it so higher is better for our max-heap logic.
|
|
137
|
+
heap = TokenBoundedMinHeap(max_tokens=max_tokens)
|
|
138
|
+
|
|
139
|
+
for row in cursor:
|
|
140
|
+
key, content, tags, tokens, rank = row
|
|
141
|
+
# FTS5 rank is typically a negative value, where lower (more negative) means better.
|
|
142
|
+
# So a score of -rank means a higher positive value is better.
|
|
143
|
+
score = -rank
|
|
144
|
+
item = {
|
|
145
|
+
"key": key,
|
|
146
|
+
"content": content,
|
|
147
|
+
"tags": tags,
|
|
148
|
+
"tokens": tokens
|
|
149
|
+
}
|
|
150
|
+
# We assume tokens is an integer
|
|
151
|
+
try:
|
|
152
|
+
tokens_int = int(tokens)
|
|
153
|
+
except (ValueError, TypeError):
|
|
154
|
+
tokens_int = len(content) // 4 + 1
|
|
155
|
+
|
|
156
|
+
heap.add(item, score, tokens_int)
|
|
157
|
+
|
|
158
|
+
return heap.get_items()
|
|
159
|
+
|
|
160
|
+
def close(self):
|
|
161
|
+
if self.conn:
|
|
162
|
+
self.conn.close()
|
|
163
|
+
|
|
164
|
+
def delete_memory(self, key):
|
|
165
|
+
with self.conn:
|
|
166
|
+
self.conn.execute("DELETE FROM memories WHERE key = ?", (key,))
|
|
167
|
+
|
|
168
|
+
def evict_file(self, filepath):
|
|
169
|
+
with self.conn:
|
|
170
|
+
self.conn.execute("DELETE FROM file_cache WHERE filepath = ?", (filepath,))
|
|
171
|
+
|
|
172
|
+
def cache_file(self, filepath, summary):
|
|
173
|
+
digest, raw_file_size = get_file_digest(filepath)
|
|
174
|
+
if not digest:
|
|
175
|
+
return False
|
|
176
|
+
|
|
177
|
+
raw_tokens = raw_file_size // 4
|
|
178
|
+
summary_tokens = len(summary) // 4
|
|
179
|
+
tokens_saved_per_hit = max(0, raw_tokens - summary_tokens)
|
|
180
|
+
|
|
181
|
+
with self.conn:
|
|
182
|
+
self.conn.execute('''
|
|
183
|
+
INSERT INTO file_cache (filepath, digest, summary, tokens_saved_per_hit)
|
|
184
|
+
VALUES (?, ?, ?, ?)
|
|
185
|
+
ON CONFLICT(filepath) DO UPDATE SET
|
|
186
|
+
digest=excluded.digest,
|
|
187
|
+
summary=excluded.summary,
|
|
188
|
+
tokens_saved_per_hit=excluded.tokens_saved_per_hit
|
|
189
|
+
''', (filepath, digest, summary, tokens_saved_per_hit))
|
|
190
|
+
|
|
191
|
+
self.bloom_filter.add(filepath)
|
|
192
|
+
return True
|
|
193
|
+
|
|
194
|
+
def check_file(self, filepath):
|
|
195
|
+
# 1. Fast rejection
|
|
196
|
+
if not self.bloom_filter.check(filepath):
|
|
197
|
+
with self.conn:
|
|
198
|
+
self.conn.execute("INSERT INTO metrics (filepath, is_hit, tokens_saved, reason) VALUES (?, 0, 0, ?)", (filepath, "not_in_bloom"))
|
|
199
|
+
return {"cached": False, "reason": "Not in bloom filter"}
|
|
200
|
+
|
|
201
|
+
# 2. Check digest
|
|
202
|
+
current_digest, _ = get_file_digest(filepath)
|
|
203
|
+
if not current_digest:
|
|
204
|
+
with self.conn:
|
|
205
|
+
self.conn.execute("INSERT INTO metrics (filepath, is_hit, tokens_saved, reason) VALUES (?, 0, 0, ?)", (filepath, "not_found"))
|
|
206
|
+
return {"cached": False, "reason": "File not found or unreadable"}
|
|
207
|
+
|
|
208
|
+
cursor = self.conn.execute("SELECT digest, summary, tokens_saved_per_hit FROM file_cache WHERE filepath = ?", (filepath,))
|
|
209
|
+
row = cursor.fetchone()
|
|
210
|
+
|
|
211
|
+
if not row:
|
|
212
|
+
with self.conn:
|
|
213
|
+
self.conn.execute("INSERT INTO metrics (filepath, is_hit, tokens_saved, reason) VALUES (?, 0, 0, ?)", (filepath, "not_in_db"))
|
|
214
|
+
return {"cached": False, "reason": "Not in database"}
|
|
215
|
+
|
|
216
|
+
cached_digest, summary, tokens_saved_per_hit = row
|
|
217
|
+
if current_digest == cached_digest:
|
|
218
|
+
with self.conn:
|
|
219
|
+
self.conn.execute("INSERT INTO metrics (filepath, is_hit, tokens_saved, reason) VALUES (?, 1, ?, ?)", (filepath, tokens_saved_per_hit, "hit"))
|
|
220
|
+
return {"cached": True, "summary": summary, "tokens_saved": tokens_saved_per_hit}
|
|
221
|
+
else:
|
|
222
|
+
with self.conn:
|
|
223
|
+
self.conn.execute("INSERT INTO metrics (filepath, is_hit, tokens_saved, reason) VALUES (?, 0, 0, ?)", (filepath, "digest_mismatch"))
|
|
224
|
+
return {"cached": False, "reason": "Digest mismatch"}
|
|
225
|
+
|
|
226
|
+
def get_metrics_dashboard(self):
|
|
227
|
+
cursor = self.conn.execute("SELECT COUNT(*), SUM(is_hit), SUM(tokens_saved) FROM metrics")
|
|
228
|
+
row = cursor.fetchone()
|
|
229
|
+
|
|
230
|
+
total_checks = row[0] if row else 0
|
|
231
|
+
total_hits = row[1] if row and row[1] is not None else 0
|
|
232
|
+
total_tokens_saved = row[2] if row and row[2] is not None else 0
|
|
233
|
+
|
|
234
|
+
hit_ratio = (total_hits / total_checks) * 100 if total_checks > 0 else 0
|
|
235
|
+
time_saved_seconds = total_tokens_saved / 50.0 # Assuming 50 tokens/sec
|
|
236
|
+
|
|
237
|
+
if time_saved_seconds > 3600:
|
|
238
|
+
time_saved_str = f"{time_saved_seconds / 3600:.2f} hours"
|
|
239
|
+
elif time_saved_seconds > 60:
|
|
240
|
+
time_saved_str = f"{time_saved_seconds / 60:.2f} minutes"
|
|
241
|
+
else:
|
|
242
|
+
time_saved_str = f"{time_saved_seconds:.2f} seconds"
|
|
243
|
+
|
|
244
|
+
return (
|
|
245
|
+
f"=== Vault ROI Dashboard ===\n"
|
|
246
|
+
f"Total File Checks: {total_checks}\n"
|
|
247
|
+
f"Cache Hits: {total_hits} ({hit_ratio:.1f}% hit ratio)\n"
|
|
248
|
+
f"Total Tokens Saved: {total_tokens_saved:,}\n"
|
|
249
|
+
f"Estimated Time Saved: {time_saved_str}\n"
|
|
250
|
+
f"==========================="
|
|
251
|
+
)
|
|
252
|
+
|
|
253
|
+
def cache_answer(self, prompt, response, dependencies, tags=""):
|
|
254
|
+
import hashlib
|
|
255
|
+
import json
|
|
256
|
+
query_hash = hashlib.sha256(prompt.encode('utf-8')).hexdigest()
|
|
257
|
+
deps_json = json.dumps([os.path.relpath(d) for d in dependencies])
|
|
258
|
+
with self.conn:
|
|
259
|
+
self.conn.execute("DELETE FROM prompt_cache WHERE query_hash = ?", (query_hash,))
|
|
260
|
+
self.conn.execute("DELETE FROM prompt_search WHERE query_hash = ?", (query_hash,))
|
|
261
|
+
self.conn.execute('''
|
|
262
|
+
INSERT INTO prompt_cache (query_hash, prompt, response, dependency_files, tags)
|
|
263
|
+
VALUES (?, ?, ?, ?, ?)
|
|
264
|
+
''', (query_hash, prompt, response, deps_json, tags))
|
|
265
|
+
self.conn.execute('''
|
|
266
|
+
INSERT INTO prompt_search (prompt, tags, query_hash)
|
|
267
|
+
VALUES (?, ?, ?)
|
|
268
|
+
''', (prompt, tags, query_hash))
|
|
269
|
+
return True
|
|
270
|
+
|
|
271
|
+
def search_questions(self, query, max_results=5):
|
|
272
|
+
# Sanitize query
|
|
273
|
+
safe_query = ''.join(c if c.isalnum() or c.isspace() else ' ' for c in query).strip()
|
|
274
|
+
if not safe_query:
|
|
275
|
+
return "No valid search terms provided."
|
|
276
|
+
|
|
277
|
+
# Prepare FTS exact match string format: "word1" "word2"
|
|
278
|
+
fts_query = ' '.join(f'"{word}"' for word in safe_query.split())
|
|
279
|
+
|
|
280
|
+
try:
|
|
281
|
+
cursor = self.conn.execute('''
|
|
282
|
+
SELECT prompt, tags, rank
|
|
283
|
+
FROM prompt_search
|
|
284
|
+
WHERE prompt_search MATCH ?
|
|
285
|
+
ORDER BY rank LIMIT ?
|
|
286
|
+
''', (fts_query, max_results))
|
|
287
|
+
except Exception as e:
|
|
288
|
+
return f"Search failed: {e}"
|
|
289
|
+
|
|
290
|
+
rows = cursor.fetchall()
|
|
291
|
+
if not rows:
|
|
292
|
+
return "No matching questions found in the cache."
|
|
293
|
+
|
|
294
|
+
out = [f"Found {len(rows)} matching cached questions:"]
|
|
295
|
+
for i, (p, t, r) in enumerate(rows, 1):
|
|
296
|
+
out.append(f"--- Option {i} ---")
|
|
297
|
+
out.append(f"Prompt: {p}")
|
|
298
|
+
out.append(f"Tags: {t if t else 'None'}")
|
|
299
|
+
out.append("")
|
|
300
|
+
out.append("Use vault_search_answer with the exact Prompt string if one matches your intent.")
|
|
301
|
+
return "\n".join(out)
|
|
302
|
+
|
|
303
|
+
|
|
304
|
+
def search_answer(self, prompt):
|
|
305
|
+
import hashlib
|
|
306
|
+
import json
|
|
307
|
+
query_hash = hashlib.sha256(prompt.encode('utf-8')).hexdigest()
|
|
308
|
+
|
|
309
|
+
cursor = self.conn.execute(
|
|
310
|
+
"SELECT id, response, dependency_files FROM prompt_cache WHERE query_hash = ?",
|
|
311
|
+
(query_hash,)
|
|
312
|
+
)
|
|
313
|
+
rows = cursor.fetchall()
|
|
314
|
+
|
|
315
|
+
for row in rows:
|
|
316
|
+
row_id, response, deps_json = row
|
|
317
|
+
dependencies = json.loads(deps_json)
|
|
318
|
+
|
|
319
|
+
is_valid = True
|
|
320
|
+
for filepath in dependencies:
|
|
321
|
+
# Get current digest from disk
|
|
322
|
+
current_digest, _ = get_file_digest(filepath)
|
|
323
|
+
if not current_digest:
|
|
324
|
+
is_valid = False
|
|
325
|
+
break
|
|
326
|
+
|
|
327
|
+
# Get cached digest from file_cache
|
|
328
|
+
c = self.conn.execute("SELECT digest FROM file_cache WHERE filepath = ?", (filepath,))
|
|
329
|
+
cached_row = c.fetchone()
|
|
330
|
+
|
|
331
|
+
if not cached_row or cached_row[0] != current_digest:
|
|
332
|
+
is_valid = False
|
|
333
|
+
break
|
|
334
|
+
|
|
335
|
+
if is_valid:
|
|
336
|
+
return response
|
|
337
|
+
else:
|
|
338
|
+
# Evict stale cache entry
|
|
339
|
+
with self.conn:
|
|
340
|
+
self.conn.execute("DELETE FROM prompt_cache WHERE id = ?", (row_id,))
|
|
341
|
+
|
|
342
|
+
return None
|
|
343
|
+
|
|
@@ -0,0 +1,144 @@
|
|
|
1
|
+
from mcp.server.mcpserver import MCPServer
|
|
2
|
+
from .core.storage import VaultStorage
|
|
3
|
+
import os
|
|
4
|
+
|
|
5
|
+
import sys
|
|
6
|
+
|
|
7
|
+
# Initialize FastMCP server
|
|
8
|
+
mcp = MCPServer("AgentVault")
|
|
9
|
+
|
|
10
|
+
# Initialize storage
|
|
11
|
+
# Using an environment variable or default local db
|
|
12
|
+
db_path = os.environ.get("AGENT_VAULT_DB_PATH", "agent_vault.db")
|
|
13
|
+
try:
|
|
14
|
+
storage = VaultStorage(db_path)
|
|
15
|
+
except Exception as e:
|
|
16
|
+
print(f"Failed to initialize VaultStorage at {db_path}: {e}", file=sys.stderr)
|
|
17
|
+
sys.exit(1)
|
|
18
|
+
|
|
19
|
+
@mcp.tool()
|
|
20
|
+
def vault_store_memory(key: str, content: str, tags: str) -> str:
|
|
21
|
+
"""Store a snippet of memory in the vault with FTS5 indexing.
|
|
22
|
+
|
|
23
|
+
Args:
|
|
24
|
+
key: A unique identifier for this memory.
|
|
25
|
+
content: The actual knowledge or code snippet.
|
|
26
|
+
tags: Comma separated tags for the memory.
|
|
27
|
+
"""
|
|
28
|
+
storage.store_memory(key, content, tags)
|
|
29
|
+
return f"Memory '{key}' stored successfully."
|
|
30
|
+
|
|
31
|
+
@mcp.tool()
|
|
32
|
+
def vault_search(query: str, max_tokens: int = 2000) -> str:
|
|
33
|
+
"""Search the vault using BM25, bounded by a token limit.
|
|
34
|
+
|
|
35
|
+
Args:
|
|
36
|
+
query: The search query.
|
|
37
|
+
max_tokens: The maximum number of tokens to return to fit in context.
|
|
38
|
+
"""
|
|
39
|
+
max_tokens = max(50, max_tokens)
|
|
40
|
+
results = storage.search_memory(query, max_tokens=max_tokens)
|
|
41
|
+
if not results:
|
|
42
|
+
return "No memories found matching the query."
|
|
43
|
+
|
|
44
|
+
output = [f"Found {len(results)} results within {max_tokens} tokens:\n"]
|
|
45
|
+
for i, res in enumerate(results, 1):
|
|
46
|
+
output.append(f"--- Result {i} (Key: {res['key']}, Tags: {res['tags']}) ---")
|
|
47
|
+
output.append(res['content'])
|
|
48
|
+
output.append("")
|
|
49
|
+
|
|
50
|
+
return "\n".join(output)
|
|
51
|
+
|
|
52
|
+
@mcp.tool()
|
|
53
|
+
def vault_cache_file(filepath: str, summary: str) -> str:
|
|
54
|
+
"""Digest a file (SHA-256) and cache its summary/AST.
|
|
55
|
+
|
|
56
|
+
Args:
|
|
57
|
+
filepath: Absolute path to the file.
|
|
58
|
+
summary: The summary or AST of the file.
|
|
59
|
+
"""
|
|
60
|
+
filepath = os.path.relpath(filepath)
|
|
61
|
+
success = storage.cache_file(filepath, summary)
|
|
62
|
+
if success:
|
|
63
|
+
return f"File '{filepath}' cached successfully."
|
|
64
|
+
else:
|
|
65
|
+
return f"Failed to cache '{filepath}'. Does it exist?"
|
|
66
|
+
|
|
67
|
+
@mcp.tool()
|
|
68
|
+
def vault_check_file(filepath: str) -> str:
|
|
69
|
+
"""Check if a file has been modified since it was last cached.
|
|
70
|
+
Uses O(1) Bloom Filter for fast rejection, then SHA-256 for exact match.
|
|
71
|
+
|
|
72
|
+
Args:
|
|
73
|
+
filepath: Absolute path to the file.
|
|
74
|
+
"""
|
|
75
|
+
filepath = os.path.relpath(filepath)
|
|
76
|
+
result = storage.check_file(filepath)
|
|
77
|
+
if result["cached"]:
|
|
78
|
+
tokens_saved = result.get("tokens_saved", 0)
|
|
79
|
+
telemetry = f"\n[Vault telemetry: Saved {tokens_saved:,} tokens by skipping raw file read]"
|
|
80
|
+
return f"File '{filepath}' is unchanged. Summary:\n{result['summary']}{telemetry}"
|
|
81
|
+
else:
|
|
82
|
+
return f"File '{filepath}' needs analysis. Reason: {result['reason']}"
|
|
83
|
+
|
|
84
|
+
@mcp.tool()
|
|
85
|
+
def vault_stats() -> str:
|
|
86
|
+
"""Get the Vault ROI Dashboard showing all-time tokens and time saved."""
|
|
87
|
+
return storage.get_metrics_dashboard()
|
|
88
|
+
|
|
89
|
+
@mcp.tool()
|
|
90
|
+
def vault_delete_memory(key: str) -> str:
|
|
91
|
+
"""Delete a memory from the vault by key."""
|
|
92
|
+
storage.delete_memory(key)
|
|
93
|
+
return f"Memory '{key}' deleted successfully."
|
|
94
|
+
|
|
95
|
+
@mcp.tool()
|
|
96
|
+
def vault_evict_file(filepath: str) -> str:
|
|
97
|
+
"""Evict a file from the vault cache."""
|
|
98
|
+
filepath = os.path.relpath(filepath)
|
|
99
|
+
storage.evict_file(filepath)
|
|
100
|
+
return f"File '{filepath}' evicted successfully."
|
|
101
|
+
|
|
102
|
+
@mcp.tool()
|
|
103
|
+
def vault_cache_answer(prompt: str, response: str, dependencies: list[str], tags: str = "") -> str:
|
|
104
|
+
"""Cache an AI response with its file dependencies and tags.
|
|
105
|
+
|
|
106
|
+
Args:
|
|
107
|
+
prompt: The exact prompt/question.
|
|
108
|
+
response: The AI's generated response.
|
|
109
|
+
dependencies: A list of absolute or relative file paths this response depends on.
|
|
110
|
+
tags: Comma separated tags (e.g. 'auth, login, jwt') to help future agents find this prompt.
|
|
111
|
+
"""
|
|
112
|
+
storage.cache_answer(prompt, response, dependencies, tags)
|
|
113
|
+
return "Answer cached successfully."
|
|
114
|
+
|
|
115
|
+
@mcp.tool()
|
|
116
|
+
def vault_search_questions(query: str, max_results: int = 5) -> str:
|
|
117
|
+
"""Discover cached questions using a keyword search.
|
|
118
|
+
Returns a list of exact cached prompts. Once you find a match, use vault_search_answer with the exact prompt.
|
|
119
|
+
|
|
120
|
+
Args:
|
|
121
|
+
query: The semantic intent or keywords (e.g., 'auth login flow').
|
|
122
|
+
max_results: Max number of question options to return.
|
|
123
|
+
"""
|
|
124
|
+
return storage.search_questions(query, max_results)
|
|
125
|
+
|
|
126
|
+
@mcp.tool()
|
|
127
|
+
def vault_search_answer(prompt: str) -> str:
|
|
128
|
+
"""Search for a cached AI response based on a prompt.
|
|
129
|
+
Checks if dependency files have changed since caching.
|
|
130
|
+
|
|
131
|
+
Args:
|
|
132
|
+
prompt: The user's prompt.
|
|
133
|
+
"""
|
|
134
|
+
result = storage.search_answer(prompt)
|
|
135
|
+
if result:
|
|
136
|
+
return f"Cached Answer:\n{result}"
|
|
137
|
+
else:
|
|
138
|
+
return "No valid cached answer found."
|
|
139
|
+
|
|
140
|
+
def main():
|
|
141
|
+
mcp.run()
|
|
142
|
+
|
|
143
|
+
if __name__ == "__main__":
|
|
144
|
+
main()
|
|
@@ -0,0 +1,131 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: agent-vault-mcp
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: DSA-Optimized AI Knowledge Bank MCP Server
|
|
5
|
+
Requires-Python: >=3.10
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
Requires-Dist: mcp<3,>=2
|
|
8
|
+
|
|
9
|
+
# Agent Vault MCP
|
|
10
|
+
|
|
11
|
+
A blazing-fast, token-efficient AI Knowledge Bank and Memory Server built on the **Model Context Protocol (MCP)**.
|
|
12
|
+
|
|
13
|
+
Agent Vault is designed to give AI coding agents (like Claude Code, Antigravity, and Cursor) persistent, cross-session memory without blowing up token limits. It acts as an intelligent file digest cache and semantic knowledge base.
|
|
14
|
+
|
|
15
|
+
## Overview
|
|
16
|
+
AI agents typically start every session with amnesia, forcing them to re-read thousands of lines of code. Agent Vault solves this using a **Two-Step Funnel**:
|
|
17
|
+
1. **Discovery**: The agent searches the Vault to find out *which* files matter.
|
|
18
|
+
2. **Execution**: The agent reads only the raw code of those specific files, edits them, and updates the vault.
|
|
19
|
+
|
|
20
|
+
When the agent wants to check a file, the Vault hashes it (SHA-256). If it hasn't changed, the Vault instantly returns a cached AST/summary instead of making the agent read the entire file.
|
|
21
|
+
|
|
22
|
+
## Features
|
|
23
|
+
- **$O(1)$ Fast Rejection**: Uses Bloom Filters to instantly know if a file is unindexed.
|
|
24
|
+
- **Token-Bounded MinHeap**: Ranks the best context snippets and strictly cuts off when the maximum token limit is reached, protecting the context window.
|
|
25
|
+
- **SQLite FTS5 (BM25)**: Fast lexical and semantic search for symbols, errors, and flows.
|
|
26
|
+
- **ROI Telemetry**: Natively calculates and tracks how many tokens and hours of inference time are saved by skipping raw file reads.
|
|
27
|
+
- **100% Portable**: Caches are stored using relative paths, meaning you can move or rename your project folder without breaking the vault.
|
|
28
|
+
|
|
29
|
+
## Dependencies
|
|
30
|
+
- Python 3.10+
|
|
31
|
+
- `mcp` (Official Model Context Protocol SDK, `mcp>=2,<3`)
|
|
32
|
+
- Standard library components (`sqlite3`, `hashlib`, `heapq`)
|
|
33
|
+
|
|
34
|
+
## Installation
|
|
35
|
+
|
|
36
|
+
### Option 1: Global Python (3.10+)
|
|
37
|
+
```bash
|
|
38
|
+
# Clone or navigate to the repository
|
|
39
|
+
cd agent-vault-mcp
|
|
40
|
+
|
|
41
|
+
# Install the package
|
|
42
|
+
pip install -e .
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
### Option 2: Using `uv` (Virtual Environment)
|
|
46
|
+
```bash
|
|
47
|
+
uv venv --python 3.12
|
|
48
|
+
.venv\Scripts\activate # On Windows
|
|
49
|
+
uv pip install -e .
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
## How to use with Antigravity / Claude Code
|
|
53
|
+
|
|
54
|
+
To add this to an MCP client like Antigravity, add the following to your `mcp_config.json` (e.g., `~/.gemini/config/mcp_config.json`):
|
|
55
|
+
|
|
56
|
+
```json
|
|
57
|
+
{
|
|
58
|
+
"mcpServers": {
|
|
59
|
+
"agent-vault": {
|
|
60
|
+
"command": "python",
|
|
61
|
+
"args": ["-m", "agent_vault.server"],
|
|
62
|
+
"env": {
|
|
63
|
+
"PYTHONPATH": "/absolute/path/to/agent-vault-mcp"
|
|
64
|
+
}
|
|
65
|
+
}
|
|
66
|
+
}
|
|
67
|
+
}
|
|
68
|
+
```
|
|
69
|
+
*Note: If you used a virtual environment, change `"command": "python"` to the absolute path of the python executable inside `.venv`.*
|
|
70
|
+
|
|
71
|
+
### Adoption (Forcing the AI to use it)
|
|
72
|
+
The Vault only saves tokens if the AI remembers to use it! Add this snippet to your project's `CLAUDE.md`, `.cursorrules`, or `GEMINI.md`:
|
|
73
|
+
|
|
74
|
+
```markdown
|
|
75
|
+
# 🧠Agent Vault & Memory Protocol
|
|
76
|
+
You are equipped with the Agent Vault MCP Server. To protect the user's token limits and eliminate hallucination, you MUST strictly adhere to the following workflow:
|
|
77
|
+
|
|
78
|
+
### 1. The "Think Before You Read" Rule (Semantic Cache)
|
|
79
|
+
Before you spend time reading files to understand an architecture, flow, or system (e.g., "How does auth work?"):
|
|
80
|
+
* **ALWAYS** call `vault_search_questions(query)` using keywords/tags to see if a previous agent already solved this.
|
|
81
|
+
* If you find a match, call `vault_search_answer(prompt)` to retrieve the pre-computed answer instantly.
|
|
82
|
+
|
|
83
|
+
### 2. The "Read Before You Write" Rule (File Cache)
|
|
84
|
+
Before you execute commands to read raw code files:
|
|
85
|
+
* **ALWAYS** call `vault_check_file(filepath)` first.
|
|
86
|
+
* If the Vault returns a Cache Hit, trust the summary/AST and DO NOT read the raw file unless you explicitly need to edit it.
|
|
87
|
+
|
|
88
|
+
### 3. The "Leave It Better Than You Found It" Rule (Updating Cache)
|
|
89
|
+
Your memory is only as good as what you save. After you complete a task:
|
|
90
|
+
* **Cache Modified Files:** If you edited a file, ALWAYS call `vault_cache_file(filepath, summary)` to update its digest and AST.
|
|
91
|
+
* **Cache New Knowledge:** If you just spent time analyzing a complex architecture or debugging a hard issue, ALWAYS call `vault_cache_answer(prompt, response, dependencies, tags)`.
|
|
92
|
+
* *Dependencies:* You MUST provide the exact file paths your answer relies on so the Vault can auto-invalidate your answer if those files change.
|
|
93
|
+
* *Tags:* Provide 5-6 broad keyword tags (e.g., "auth, login, jwt") so future agents can easily discover your answer via `vault_search_questions`.
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
## Local vs Global Vaults
|
|
97
|
+
|
|
98
|
+
By default, the MCP server creates `agent_vault.db` inside your current active project workspace. Paths are stored **relatively**. This means if you ask the agent about a file in an external project (e.g., `../Project_B/main.py`), that cross-project memory is stored locally inside your current project's database.
|
|
99
|
+
|
|
100
|
+
If you prefer a **"Global Brain"** that shares all memories and file caches across every single project on your computer, simply add `AGENT_VAULT_DB_PATH` to the `env` variables in your `mcp_config.json`:
|
|
101
|
+
|
|
102
|
+
```json
|
|
103
|
+
"env": {
|
|
104
|
+
"PYTHONPATH": "/absolute/path/to/agent-vault-mcp",
|
|
105
|
+
"AGENT_VAULT_DB_PATH": "/absolute/path/to/.global_agent_vault.db"
|
|
106
|
+
}
|
|
107
|
+
```
|
|
108
|
+
|
|
109
|
+
## Available MCP Tools
|
|
110
|
+
- `vault_store_memory(key, content, tags)`: Save arbitrary architectural notes or debugging insights.
|
|
111
|
+
- `vault_search(query, max_tokens)`: Search the vault using BM25 ranking.
|
|
112
|
+
- `vault_cache_file(filepath, summary)`: Hash a file and cache its summary.
|
|
113
|
+
- `vault_check_file(filepath)`: Verify a file's hash and return its cached summary + telemetry metrics.
|
|
114
|
+
- `vault_stats()`: View the ROI dashboard of tokens and time saved.
|
|
115
|
+
- `vault_delete_memory(key)`: Delete a stored memory.
|
|
116
|
+
- `vault_evict_file(filepath)`: Evict a file from the vault cache.- `vault_cache_answer(prompt, response, dependencies, tags)`: Save an AI-generated answer linked to specific files and keywords.
|
|
117
|
+
- `vault_search_questions(query, max_results)`: Keyword-search an FTS5 index to find exactly how previous cached questions were phrased based on tags.
|
|
118
|
+
- `vault_search_answer(prompt)`: Retrieve a cached AI answer (automatically invalidates if dependent files have changed).
|
|
119
|
+
|
|
120
|
+
## Semantic Prompt Caching (Intent Caching)
|
|
121
|
+
|
|
122
|
+
In `v0.2.0`, Agent Vault introduced **Semantic Prompt Caching**. Instead of forcing AI agents to repeatedly read files and re-reason through complex architectural questions (e.g., *"How does the auth flow work?"*), agents can now cache their reasoning.
|
|
123
|
+
|
|
124
|
+
### Dependency Invalidation
|
|
125
|
+
To solve the classic LLM problem of "stale context hallucination," Agent Vault uses **Deterministic Dependency Invalidation**. When an agent caches an answer, it explicitly lists the files that answer depends on.
|
|
126
|
+
|
|
127
|
+
When a future agent asks the same question, the Vault calculates the real-time SHA-256 digest of those dependencies. If any file has changed, the cached answer is instantly evicted, forcing the AI to generate a fresh, accurate response.
|
|
128
|
+
|
|
129
|
+
### N-to-1 Tag Mapping
|
|
130
|
+
To solve the problem of "brittle exact matching" (where *"How does login work?"* misses a cache for *"How does the login work?"*), agents can assign **tags** to cached answers.
|
|
131
|
+
Agents can use `vault_search_questions("login")` to hit the FTS5 index, discover the exact phrasing of the cached question, and then fetch the answer—bypassing the need for heavy vector databases!
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
agent_vault/__init__.py
|
|
4
|
+
agent_vault/server.py
|
|
5
|
+
agent_vault/core/__init__.py
|
|
6
|
+
agent_vault/core/dsa.py
|
|
7
|
+
agent_vault/core/storage.py
|
|
8
|
+
agent_vault_mcp.egg-info/PKG-INFO
|
|
9
|
+
agent_vault_mcp.egg-info/SOURCES.txt
|
|
10
|
+
agent_vault_mcp.egg-info/dependency_links.txt
|
|
11
|
+
agent_vault_mcp.egg-info/entry_points.txt
|
|
12
|
+
agent_vault_mcp.egg-info/requires.txt
|
|
13
|
+
agent_vault_mcp.egg-info/top_level.txt
|
|
14
|
+
tests/test_prompt_cache.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
mcp<3,>=2
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
agent_vault
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "agent-vault-mcp"
|
|
3
|
+
version = "0.2.0"
|
|
4
|
+
description = "DSA-Optimized AI Knowledge Bank MCP Server"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.10"
|
|
7
|
+
dependencies = [
|
|
8
|
+
"mcp>=2,<3",
|
|
9
|
+
]
|
|
10
|
+
|
|
11
|
+
[project.scripts]
|
|
12
|
+
agent-vault = "agent_vault.server:main"
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import time
|
|
3
|
+
from agent_vault.core.storage import VaultStorage
|
|
4
|
+
|
|
5
|
+
def run_tests():
|
|
6
|
+
db_path = "test_vault.db"
|
|
7
|
+
if os.path.exists(db_path):
|
|
8
|
+
os.remove(db_path)
|
|
9
|
+
|
|
10
|
+
storage = VaultStorage(db_path)
|
|
11
|
+
|
|
12
|
+
# Create a dummy dependency file
|
|
13
|
+
with open("dummy_dep.py", "w") as f:
|
|
14
|
+
f.write("print('hello')\n")
|
|
15
|
+
|
|
16
|
+
# We must cache the file in file_cache first, because search_answer relies on file_cache having the digest
|
|
17
|
+
storage.cache_file("dummy_dep.py", "A dummy file")
|
|
18
|
+
|
|
19
|
+
prompt = "How does the dummy work?"
|
|
20
|
+
response = "It prints hello."
|
|
21
|
+
deps = ["dummy_dep.py"]
|
|
22
|
+
|
|
23
|
+
# Cache the answer
|
|
24
|
+
storage.cache_answer(prompt, response, deps)
|
|
25
|
+
print("Answer cached.")
|
|
26
|
+
|
|
27
|
+
# Search the answer
|
|
28
|
+
res = storage.search_answer(prompt)
|
|
29
|
+
if res == response:
|
|
30
|
+
print("Hit! (Expected)")
|
|
31
|
+
else:
|
|
32
|
+
print("Miss! (Unexpected)")
|
|
33
|
+
|
|
34
|
+
# Modify the dependency file
|
|
35
|
+
with open("dummy_dep.py", "a") as f:
|
|
36
|
+
f.write("print('world')\n")
|
|
37
|
+
|
|
38
|
+
# Search the answer again
|
|
39
|
+
res = storage.search_answer(prompt)
|
|
40
|
+
if res is None:
|
|
41
|
+
print("Miss! (Expected, file changed)")
|
|
42
|
+
else:
|
|
43
|
+
print(f"Hit! (Unexpected, returned {res})")
|
|
44
|
+
|
|
45
|
+
# Clean up
|
|
46
|
+
storage.close()
|
|
47
|
+
if os.path.exists(db_path):
|
|
48
|
+
os.remove(db_path)
|
|
49
|
+
if os.path.exists("dummy_dep.py"):
|
|
50
|
+
os.remove("dummy_dep.py")
|
|
51
|
+
|
|
52
|
+
if __name__ == "__main__":
|
|
53
|
+
run_tests()
|