Compare commits
426
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
600ddba61e | ||
|
|
afc4ae3779 | ||
|
|
35c304bd60 | ||
|
|
33aacf5c6c | ||
|
|
ca180dde45 | ||
|
|
76a1bfa554 | ||
|
|
e90cfa4ac9 | ||
|
|
10785666c7 | ||
|
|
59f9a2bf25 | ||
|
|
30700de670 | ||
|
|
a63253f59f | ||
|
|
afd00c037c | ||
|
|
3d877b05a5 | ||
|
|
902704dfcf | ||
|
|
449a9d70b2 | ||
|
|
e301883952 | ||
|
|
487e2a5e6d | ||
|
|
511ca72361 | ||
|
|
a3b0d2651c | ||
|
|
b79caa9aa8 | ||
|
|
abbd3619c6 | ||
|
|
7126bf8a23 | ||
|
|
ba9d227f4c | ||
|
|
d05b49a24b | ||
|
|
858f0b3a06 | ||
|
|
ce137de643 | ||
|
|
920c56987b | ||
|
|
f5dfe59b90 | ||
|
|
9901aa1e07 | ||
|
|
9181c9a29e | ||
|
|
c8b898bd05 | ||
|
|
41710ba176 | ||
|
|
5d22a8e8fb | ||
|
|
3d6b380515 | ||
|
|
3c1431ad99 | ||
|
|
f75251fc5c | ||
|
|
52c148c203 | ||
|
|
21a22decea | ||
|
|
02ca15de42 | ||
|
|
6ae6663c0d | ||
|
|
ca561aca9e | ||
|
|
b52b483cb5 | ||
|
|
cbc196805d | ||
|
|
69383af896 | ||
|
|
abf24b4e22 | ||
|
|
1cbf9adb67 | ||
|
|
7ddfe9e237 | ||
|
|
2e74a324da | ||
|
|
5437cc0299 | ||
|
|
bca87412f7 | ||
|
|
eb9be90312 | ||
|
|
b8da88c854 | ||
|
|
824db5d5b8 | ||
|
|
eca8526e4a | ||
|
|
3ef8099892 | ||
|
|
8b80959bba | ||
|
|
f890479705 | ||
|
|
576c44d2ff | ||
|
|
f9042e378d | ||
|
|
6a17992e81 | ||
|
|
7d4fd1aa40 | ||
|
|
1f89731406 | ||
|
|
e1e5f36cee | ||
|
|
7ceaa22a66 | ||
|
|
0a16295c1f | ||
|
|
088dfecbc7 | ||
|
|
9e30ae2526 | ||
|
|
13f3052e6e | ||
|
|
654e4c0cfa | ||
|
|
a5e5372192 | ||
|
|
7007ffdb04 | ||
|
|
320d1ce435 | ||
|
|
c571fac7db | ||
|
|
3bedc1cedc | ||
|
|
70d60e96cf | ||
|
|
568e3c3028 | ||
|
|
cbaec36f66 | ||
|
|
d8aada7b0e | ||
|
|
16e1cc4934 | ||
|
|
3ee9437020 | ||
|
|
9671786faf | ||
|
|
33645e08cd | ||
|
|
2f5844b38d | ||
|
|
1267e61edd | ||
|
|
931f2a77ff | ||
|
|
eeff5001af | ||
|
|
f835c731fe | ||
|
|
b6dbd614fc | ||
|
|
32fc9b7477 | ||
|
|
e4f54a6071 | ||
|
|
343b972a95 | ||
|
|
58c02feef0 | ||
|
|
a5f8b58ab5 | ||
|
|
78008a1ad0 | ||
|
|
2128c02e0e | ||
|
|
2eab07834a | ||
|
|
9f41d98172 | ||
|
|
84bab9c5b7 | ||
|
|
d73e552189 | ||
|
|
712a862841 | ||
|
|
f64c5d2097 | ||
|
|
d4bf740618 | ||
|
|
70a7411659 | ||
|
|
dee581396b | ||
|
|
6a1d5fcd30 | ||
|
|
16ed93b9a2 | ||
|
|
581bbf3fc6 | ||
|
|
d00c843262 | ||
|
|
33442f1961 | ||
|
|
a525df4837 | ||
|
|
90a2201655 | ||
|
|
34365c3248 | ||
|
|
06c912df34 | ||
|
|
43dc50dd3f | ||
|
|
7d5d5b2781 | ||
|
|
b79ab2b752 | ||
|
|
cf9918891b | ||
|
|
9372462e13 | ||
|
|
f2fc8f9f26 | ||
|
|
6a80ecbf65 | ||
|
|
870bf4a3d1 | ||
|
|
099f4c925a | ||
|
|
e08faadc17 | ||
|
|
dbd1d1a743 | ||
|
|
c084765950 | ||
|
|
d6ad53986a | ||
|
|
6076354a9c | ||
|
|
9f9c3a1b40 | ||
|
|
8ba862b026 | ||
|
|
fd87de9c15 | ||
|
|
adc85129ba | ||
|
|
2ff805d6e9 | ||
|
|
8125a0d758 | ||
|
|
e1e137b027 | ||
|
|
6b5aa3afe8 | ||
|
|
e28b8c00f6 | ||
|
|
1be5ff33b0 | ||
|
|
aeb0c8b553 | ||
|
|
d0b2ab9ad2 | ||
|
|
93562bfaaf | ||
|
|
9679d8139d | ||
|
|
ab7feb144b | ||
|
|
0c4b79b6d3 | ||
|
|
e9270fd312 | ||
|
|
da21e0727c | ||
|
|
d4b8b3544b | ||
|
|
873223964b | ||
|
|
e5724fcba0 | ||
|
|
848451bd01 | ||
|
|
f82f58fa83 | ||
|
|
2d74007d80 | ||
|
|
05686e1236 | ||
|
|
93300b9104 | ||
|
|
9402572339 | ||
|
|
e9cc771bbd | ||
|
|
2a2b90b0a0 | ||
|
|
2635bbb49e | ||
|
|
2e88bac605 | ||
|
|
2644930561 | ||
|
|
bbd3c5dc04 | ||
|
|
4f9cf15cdd | ||
|
|
2d95f78b09 | ||
|
|
773ef0cb63 | ||
|
|
7b2263ba3b | ||
|
|
d054b88403 | ||
|
|
1c32a7b928 | ||
|
|
d38ecdb9ec | ||
|
|
aad07a141b | ||
|
|
3fc87e767c | ||
|
|
e22ae05f47 | ||
|
|
b57e337fa2 | ||
|
|
c05c491d77 | ||
|
|
fc941d5cae | ||
|
|
576016f5dc | ||
|
|
b3995d1430 | ||
|
|
f519fc4fd0 | ||
|
|
72fd3d59db | ||
|
|
5a61ac50e9 | ||
|
|
1f1716bdb0 | ||
|
|
61a8014f9d | ||
|
|
c5091d29cd | ||
|
|
e82bc56580 | ||
|
|
fa0e63b088 | ||
|
|
27cb7e43e0 | ||
|
|
9e23e83abf | ||
|
|
bdf93f0660 | ||
|
|
b0e8ac0f4d | ||
|
|
f74b577e02 | ||
|
|
3c633e5e16 | ||
|
|
cf0537ba7e | ||
|
|
e5944b63e7 | ||
|
|
37348c859e | ||
|
|
cece2c903c | ||
|
|
d7c73f4342 | ||
|
|
3b9d2db091 | ||
|
|
9790d904e0 | ||
|
|
2463efd0f2 | ||
|
|
6674ee4706 | ||
|
|
57f154454d | ||
|
|
4028dd91f8 | ||
|
|
0e81d1a25e | ||
|
|
8a2388a48f | ||
|
|
48185a4bee | ||
|
|
7e23f8e149 | ||
|
|
f659bb17c4 | ||
|
|
f31f82627c | ||
|
|
e1c6220f0e | ||
|
|
66cbdda3cb | ||
|
|
cf4bd598b4 | ||
|
|
443c94c827 | ||
|
|
26794aab09 | ||
|
|
7863ffeb49 | ||
|
|
9e2890ba81 | ||
|
|
e0e65c44f6 | ||
|
|
6881f63781 | ||
|
|
6cb309f72b | ||
|
|
f9fe6953a3 | ||
|
|
07de798c3b | ||
|
|
cefa75545a | ||
|
|
cd99eef4c5 | ||
|
|
cd4b3e96e2 | ||
|
|
e02e7ad3d4 | ||
|
|
98fee1e380 | ||
|
|
906b740dd7 | ||
|
|
7990381f6a | ||
|
|
045e8910d1 | ||
|
|
81441ee9af | ||
|
|
30a319a6ab | ||
|
|
9cfdd464a9 | ||
|
|
8d1bfbbd2b | ||
|
|
7d6c570a3a | ||
|
|
26a64cc00e | ||
|
|
087545cc1b | ||
|
|
7415ebff7c | ||
|
|
0c97b555ab | ||
|
|
4d117cc274 | ||
|
|
a757765ab2 | ||
|
|
300d089b6a | ||
|
|
1a1fb35cb0 | ||
|
|
914ba7962c | ||
|
|
6f173b10a7 | ||
|
|
ea834bc7dc | ||
|
|
4fd7c5d1f8 | ||
|
|
36783df320 | ||
|
|
ee4510a762 | ||
|
|
f3f2c6b023 | ||
|
|
6f7437be21 | ||
|
|
d7f6723546 | ||
|
|
2c32ffadc9 | ||
|
|
baf5447de2 | ||
|
|
84985ee9bc | ||
|
|
ecaa1ad1e0 | ||
|
|
ea0c616240 | ||
|
|
0b29378eb8 | ||
|
|
c2e801ccb0 | ||
|
|
410c208746 | ||
|
|
c475c6bb56 | ||
|
|
f841bcb92d | ||
|
|
fa82efc886 | ||
|
|
627ec5d524 | ||
|
|
bdb33c58d1 | ||
|
|
1dbbe39ea1 | ||
|
|
cef42d8154 | ||
|
|
f8f62030e3 | ||
|
|
4768bf39ef | ||
|
|
865fb91298 | ||
|
|
1f5dc8bd15 | ||
|
|
d3d2684b11 | ||
|
|
41025c3b7c | ||
|
|
1b5c262a8a | ||
|
|
0096115678 | ||
|
|
b104bad02c | ||
|
|
45ffc7fe90 | ||
|
|
99122055f0 | ||
|
|
75e2679cf1 | ||
|
|
d93dfea8ce | ||
|
|
e5209b18b3 | ||
|
|
a7adfbb0df | ||
|
|
3461398b52 | ||
|
|
a915584e39 | ||
|
|
2c72af5525 | ||
|
|
41bb6d710b | ||
|
|
7af01e35e9 | ||
|
|
a3e458ad43 | ||
|
|
704e41fa27 | ||
|
|
f30ca3deda | ||
|
|
d61517d502 | ||
|
|
df17570d8a | ||
|
|
7a9e99998a | ||
|
|
cc3cdc2f83 | ||
|
|
88630f93d7 | ||
|
|
73460fa4e6 | ||
|
|
6a28373ecf | ||
|
|
66071fec9d | ||
|
|
b8b40e458b | ||
|
|
e65bd361cc | ||
|
|
b8fa0e8cfc | ||
|
|
b739b9e36a | ||
|
|
5b16882a5e | ||
|
|
56489a58d1 | ||
|
|
a8a63818c7 | ||
|
|
d8050387e4 | ||
|
|
38e03e419d | ||
|
|
6488c9bc77 | ||
|
|
d2965e64e6 | ||
|
|
ecf16ea1e6 | ||
|
|
0b17a67c70 | ||
|
|
e7c9a6832d | ||
|
|
2d787c4ffd | ||
|
|
111e8c70a2 | ||
|
|
d441ab814d | ||
|
|
f8285b7b90 | ||
|
|
a209ef1ae2 | ||
|
|
3573e53b1d | ||
|
|
585ac76f39 | ||
|
|
b32767caa8 | ||
|
|
cd4d449f8e | ||
|
|
7a3dbc1958 | ||
|
|
a69bdbb55f | ||
|
|
f50cc25dfb | ||
|
|
6e90df9818 | ||
|
|
dffb87080f | ||
|
|
1cac35728f | ||
|
|
26e6877b53 | ||
|
|
1ac80bda6f | ||
|
|
3c78b717b0 | ||
|
|
9321c59bf1 | ||
|
|
696d99ca1e | ||
|
|
4b584e4d0b | ||
|
|
e5c7e166c5 | ||
|
|
083295dc6f | ||
|
|
731238707d | ||
|
|
62c0992075 | ||
|
|
02b0f7799d | ||
|
|
7c18723fd9 | ||
|
|
417fac61e4 | ||
|
|
105cdf1fbf | ||
|
|
a0cea84d82 | ||
|
|
200bab233e | ||
|
|
c9ff37dcbf | ||
|
|
91397190c0 | ||
|
|
fd88c0efa5 | ||
|
|
ea4df8dbb5 | ||
|
|
b6a4f17cbe | ||
|
|
ffc96bec97 | ||
|
|
8cb8b9128e | ||
|
|
64d96a9c53 | ||
|
|
9dedac1dbd | ||
|
|
413ddbb45d | ||
|
|
246912f596 | ||
|
|
2d31b67d0c | ||
|
|
349c112c61 | ||
|
|
939cb40a73 | ||
|
|
c5700ff5b4 | ||
|
|
6bb83f4600 | ||
|
|
a94a90ea3f | ||
|
|
9e5a066d26 | ||
|
|
5095d5e36f | ||
|
|
22ca6a8d73 | ||
|
|
0ff36548e0 | ||
|
|
d344ef26da | ||
|
|
4fed005662 | ||
|
|
b42b35bf93 | ||
|
|
db70fdbe5e | ||
|
|
c5273f5fd4 | ||
|
|
4285e94406 | ||
|
|
35dfd3aa0c | ||
|
|
0bcbf8491b | ||
|
|
f0f0d554f2 | ||
|
|
0ad6ee3156 | ||
|
|
39bf6820d6 | ||
|
|
8ef9c48a62 | ||
|
|
7fe773c0ee | ||
|
|
58e68f3e4a | ||
|
|
4f533dde94 | ||
|
|
08d2c78ae7 | ||
|
|
8e2e2d5bf2 | ||
|
|
4a55068db7 | ||
|
|
e1f539c612 | ||
|
|
742f212b2f | ||
|
|
f2b0ff7d38 | ||
|
|
8ae3ae13a6 | ||
|
|
546d595c9f | ||
|
|
26944e25bc | ||
|
|
e6333719ee | ||
|
|
d886d3acb9 | ||
|
|
35b2cbb6ed | ||
|
|
f4390bdc2e | ||
|
|
e0f0da5d2d | ||
|
|
a9e6d9f731 | ||
|
|
8ce06e3e7c | ||
|
|
365fa3ce50 | ||
|
|
2eb1019da9 | ||
|
|
2f2db2a6e2 | ||
|
|
caa53ee370 | ||
|
|
5cdc714a38 | ||
|
|
78aa7c537e | ||
|
|
3f31cbf505 | ||
|
|
b7abf8565a | ||
|
|
682cbf38ee | ||
|
|
5e8952c54a | ||
|
|
8364b9c5d5 | ||
|
|
5a486883e8 | ||
|
|
d2c32cb8e4 | ||
|
|
ce691549ba | ||
|
|
72b61214f6 | ||
|
|
103994c25f | ||
|
|
36b5627d2c | ||
|
|
d284de28c7 | ||
|
|
93609f74ab | ||
|
|
9a5f83adb4 | ||
|
|
b4320254b2 | ||
|
|
97f7a365e8 | ||
|
|
20e17f28ad | ||
|
|
80b1badf74 | ||
|
|
ea662d062e | ||
|
|
94cf89b570 | ||
|
|
439424559e | ||
|
|
4c4b3568db | ||
|
|
a706905653 | ||
|
|
fe12be47a0 | ||
|
|
a56cd044e5 | ||
|
|
438ce98b40 | ||
|
|
446c75f3e2 | ||
|
|
276a4ba7e8 | ||
|
|
31f1c53c8f |
@@ -0,0 +1,15 @@
|
||||
{
|
||||
"$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
|
||||
"name": "hindsight",
|
||||
"description": "Official Hindsight integrations for Claude Code",
|
||||
"owner": {
|
||||
"name": "vectorize-io"
|
||||
},
|
||||
"plugins": [
|
||||
{
|
||||
"name": "hindsight-memory",
|
||||
"description": "Automatic long-term memory for Claude Code via Hindsight",
|
||||
"source": "./hindsight-integrations/claude-code"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
---
|
||||
name: code-review
|
||||
description: Review changed code against project standards. Checks for missing tests, dead code, type safety, lint issues, and coding conventions. Run after completing any implementation work.
|
||||
user_invocable: true
|
||||
---
|
||||
|
||||
# Code Review
|
||||
|
||||
Review all changed code against the project's quality standards and coding conventions.
|
||||
|
||||
## Code Standards
|
||||
|
||||
Read and internalize these standards before writing code. The review steps below verify compliance.
|
||||
|
||||
### Python Style
|
||||
- Python 3.11+, type hints required
|
||||
- Async throughout (asyncpg, async FastAPI)
|
||||
- Pydantic models for request/response
|
||||
- Ruff for linting (line-length 120)
|
||||
- No Python files at project root - maintain clean directory structure
|
||||
- **Never use multi-item tuple return values** — not even for internal/private functions. Always use a dataclass or Pydantic model. No exceptions, no "it's just two values" shortcuts. If a function returns more than one value, define a named type for it.
|
||||
|
||||
### Type Safety with Pydantic Models
|
||||
**NEVER use raw `dict` types for structured data** — this applies to all code, including internal helpers and private functions. If the dict has known keys, it must be a dataclass or Pydantic model:
|
||||
- Use Pydantic `BaseModel` for all data structures passed between functions
|
||||
- Use `@dataclass` for lightweight internal data containers when Pydantic validation isn't needed
|
||||
- Add `@field_validator` for type coercion (e.g., ensuring datetimes are timezone-aware)
|
||||
- Avoid `dict.get()` patterns - use typed model attributes instead
|
||||
- Parse external data (JSON, API responses) into Pydantic models at the boundary
|
||||
- This catches type errors at parse time, not deep in business logic
|
||||
- The only acceptable `dict` usage is for truly dynamic/unknown keys (e.g., arbitrary metadata, JSON blobs with no fixed schema)
|
||||
|
||||
```python
|
||||
# BAD - error-prone dict access
|
||||
def process(data: dict) -> str:
|
||||
return data.get("name", "") # No validation, silent failures
|
||||
|
||||
# GOOD - typed and validated
|
||||
class UserData(BaseModel):
|
||||
name: str
|
||||
created_at: datetime
|
||||
|
||||
@field_validator("created_at", mode="before")
|
||||
@classmethod
|
||||
def ensure_tz_aware(cls, v):
|
||||
if isinstance(v, str):
|
||||
v = datetime.fromisoformat(v.replace("Z", "+00:00"))
|
||||
if v.tzinfo is None:
|
||||
return v.replace(tzinfo=timezone.utc)
|
||||
return v
|
||||
|
||||
def process(data: UserData) -> str:
|
||||
return data.name # Type-safe, validated at construction
|
||||
```
|
||||
|
||||
### TypeScript Style
|
||||
- Next.js App Router for control plane
|
||||
- Tailwind CSS with shadcn/ui components
|
||||
|
||||
### Code Comments
|
||||
- **Always comment non-trivial technical decisions** with the reasoning behind the choice. If someone would ask "why is it done this way?", there should be a comment.
|
||||
- **Keep comments up to date with history** — when changing an approach, update the comment to explain what was tried before and why it was changed. Comments serve as a tracker of previous implementations that likely had problems.
|
||||
- Don't comment obvious code — only where the "why" isn't self-evident from the code itself.
|
||||
|
||||
```python
|
||||
# BAD - no context for future readers
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
# GOOD - explains the non-obvious choice
|
||||
# Use return_exceptions=True to avoid cancelling sibling tasks on failure.
|
||||
# Previously we used TaskGroup but it cancelled all tasks when one failed,
|
||||
# causing partial writes that left orphaned entity links (see #412).
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
```
|
||||
|
||||
### Branch Hygiene
|
||||
- **Always start new feature branches from `origin/main`** — rebase to ensure a clean base.
|
||||
- **Only include commits relevant to the PR/branch/feature** — no unrelated changes. If the branch contains commits that don't belong, they must be removed before merging.
|
||||
|
||||
### General Principles
|
||||
- Don't add features, refactor code, or make "improvements" beyond what was asked
|
||||
- Don't add unnecessary error handling for impossible scenarios
|
||||
- Don't create helpers or abstractions for one-time operations
|
||||
- No backwards-compatibility hacks (unused vars, re-exports, "removed" comments)
|
||||
- Three similar lines of code is better than a premature abstraction
|
||||
|
||||
## Review Steps
|
||||
|
||||
### 1. Check branch hygiene
|
||||
|
||||
- Run `git log --oneline main..HEAD` to list all commits on the branch.
|
||||
- Verify every commit is relevant to the feature/PR. Flag any unrelated commits.
|
||||
- Check the branch is based on a recent `origin/main` (no stale base).
|
||||
|
||||
### 2. Identify changed files
|
||||
|
||||
Run `git diff --name-only HEAD` (unstaged) and `git diff --cached --name-only` (staged) to get all changed files. If there are no local changes, diff against the base branch using `git diff main...HEAD --name-only` and `git diff main...HEAD` to review all commits on the current branch.
|
||||
|
||||
### 3. Run linters
|
||||
|
||||
```bash
|
||||
./scripts/hooks/lint.sh
|
||||
```
|
||||
|
||||
Report any failures. Do NOT fix them yourself — just report.
|
||||
|
||||
### 4. Check for dead code
|
||||
|
||||
For each changed Python file, check for:
|
||||
- Unused imports (Ruff should catch these, but verify)
|
||||
- Functions/methods/classes that were added but are never called from anywhere
|
||||
- Variables assigned but never read
|
||||
- Commented-out code blocks that should be removed
|
||||
|
||||
For each changed TypeScript file, check for:
|
||||
- Unused imports
|
||||
- Unused variables or functions
|
||||
- Commented-out code
|
||||
|
||||
### 5. Check type safety (Python)
|
||||
|
||||
For each changed Python file, check for violations:
|
||||
- **No raw `dict` for structured data** — must use Pydantic model or dataclass, even for internal/private functions (only exception: truly dynamic/unknown keys)
|
||||
- **No multi-item tuple returns** — must use dataclass or Pydantic model, even for internal/private functions (no exceptions)
|
||||
- **Missing type hints** on function parameters and return types
|
||||
- **Missing `@field_validator`** for datetime fields that should be timezone-aware
|
||||
|
||||
### 6. Check for missing tests
|
||||
|
||||
For each new or significantly changed function/endpoint/class:
|
||||
- Check if there is a corresponding test addition or update
|
||||
- New API endpoints MUST have integration tests
|
||||
- New utility functions MUST have unit tests
|
||||
- Bug fixes SHOULD have a regression test
|
||||
|
||||
Flag any new logic that lacks test coverage.
|
||||
|
||||
### 7. Check API consistency
|
||||
|
||||
If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
|
||||
- Were the OpenAPI specs regenerated? (`./scripts/generate-openapi.sh`)
|
||||
- Were the client SDKs regenerated? (`./scripts/generate-clients.sh`)
|
||||
- Were the control plane proxy routes updated? (`hindsight-control-plane/src/app/api/`)
|
||||
|
||||
### 8. Check code comments
|
||||
|
||||
For each non-trivial change:
|
||||
- **New non-obvious logic** — is there a comment explaining the reasoning?
|
||||
- **Changed approach** — does the comment include what was done before and why it changed?
|
||||
- **Stale comments** — do existing comments near the changed code still accurately describe the behavior?
|
||||
|
||||
### 9. Check integration completeness
|
||||
|
||||
If any files in `hindsight-integrations/` were added or changed, verify:
|
||||
- **Tests exist** — the integration must have tests that simulate/exercise the external framework (not just pure unit tests of helpers). Check for a `tests/` directory with meaningful test files.
|
||||
- **CI job exists** — check `.github/workflows/test.yml` for a corresponding `test-<name>-integration` job. If missing, flag it.
|
||||
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh`. If missing, flag it.
|
||||
- **Code standards** — the integration code must follow all Python style rules (type hints, no raw dicts, no tuple returns, etc.).
|
||||
|
||||
### 10. Check MCP tool registration completeness
|
||||
|
||||
If any new MCP tools were added or existing tools renamed in `hindsight-api-slim/hindsight_api/mcp_tools.py`:
|
||||
- **`_ALL_TOOLS` set** in `mcp_tools.py` — must include the new tool name
|
||||
- **`tools_to_register` default set** in `register_mcp_tools()` in `mcp_tools.py` — must include the new tool name
|
||||
- **`_SINGLE_BANK_TOOLS` set** in `hindsight-api-slim/hindsight_api/api/mcp.py` — must include the new tool if it is bank-scoped (not a bank-management tool like `list_banks`/`create_bank`)
|
||||
- **`MCP_TOOL_GROUPS`** in `hindsight-control-plane/src/components/bank-config-view.tsx` — must include the new tool in the appropriate group for the UI tool selector
|
||||
- **Tool count assertions** in tests (e.g., `test_mcp_tools.py`) — must be updated to reflect the new count
|
||||
|
||||
### 11. Review against other coding standards
|
||||
|
||||
Check the diff for violations of the standards listed above:
|
||||
- Python files at project root (not allowed)
|
||||
- Missing async patterns (should be async throughout)
|
||||
- Pydantic models for request/response
|
||||
- Line length > 120 chars
|
||||
- New features/code beyond what was asked (over-engineering)
|
||||
- Unnecessary error handling for impossible scenarios
|
||||
- Premature abstractions or speculative helpers
|
||||
- Backwards-compatibility hacks (unused vars, re-exports, "removed" comments)
|
||||
|
||||
### 12. Report findings
|
||||
|
||||
Present a clear summary organized by severity:
|
||||
|
||||
**Must fix** — issues that will break CI or violate hard project rules:
|
||||
- Unrelated commits on the branch
|
||||
- Lint failures
|
||||
- Missing type hints on public functions
|
||||
- Raw dict usage for structured data (including internal code)
|
||||
- Multi-item tuple returns (including internal code)
|
||||
- Missing tests for new endpoints
|
||||
- New integration missing tests, CI job, or release-integration.sh entry
|
||||
|
||||
**Should fix** — issues that hurt code quality:
|
||||
- Dead code / unused imports missed by linter
|
||||
- Missing tests for non-trivial utility functions
|
||||
- Over-engineering beyond the task scope
|
||||
|
||||
**Note** — observations that may or may not need action:
|
||||
- API changes that might need client regeneration
|
||||
- Patterns that deviate from nearby code style
|
||||
|
||||
For each finding, include the file path, line number, and a brief explanation.
|
||||
|
||||
Do NOT auto-fix any issues. Report all findings and let the user decide what to address. If there are no findings, confirm the code looks good.
|
||||
+2
-1
@@ -2,7 +2,7 @@
|
||||
# Copy this file to .env and fill in your values
|
||||
|
||||
# LLM Configuration (Required)
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, volcano
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
@@ -44,6 +44,7 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
|
||||
# Database (Optional - uses embedded pg0 by default)
|
||||
# HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@host:5432/db
|
||||
# HINDSIGHT_API_MIGRATION_DATABASE_URL= # Direct PostgreSQL URL for migrations (bypasses PgBouncer). Falls back to DATABASE_URL.
|
||||
# HINDSIGHT_API_DATABASE_SCHEMA=public # PostgreSQL schema name (default: public)
|
||||
|
||||
# Vector Extension (Optional - uses pgvector by default)
|
||||
|
||||
@@ -34,7 +34,7 @@ jobs:
|
||||
env:
|
||||
UMAMI_URL: https://analytics.hindsight.vectorize.io
|
||||
UMAMI_WEBSITE_ID: ${{ secrets.UMAMI_WEBSITE_ID }}
|
||||
- uses: actions/upload-pages-artifact@v4
|
||||
- uses: actions/upload-pages-artifact@v5
|
||||
with:
|
||||
path: hindsight-docs/build
|
||||
deploy:
|
||||
@@ -44,5 +44,5 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
needs: build
|
||||
steps:
|
||||
- uses: actions/deploy-pages@v4
|
||||
- uses: actions/deploy-pages@v5
|
||||
id: deployment
|
||||
|
||||
@@ -0,0 +1,120 @@
|
||||
name: Release Integration
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- 'integrations/**'
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
id-token: write # for PyPI trusted publishing
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Extract integration info
|
||||
id: info
|
||||
run: |
|
||||
# refs/tags/integrations/litellm/v0.1.0 → integration=litellm, version=0.1.0
|
||||
TAG="${GITHUB_REF#refs/tags/}"
|
||||
INTEGRATION=$(echo "$TAG" | cut -d'/' -f2)
|
||||
VERSION=$(echo "$TAG" | cut -d'/' -f3 | sed 's/^v//')
|
||||
echo "integration=$INTEGRATION" >> $GITHUB_OUTPUT
|
||||
echo "version=$VERSION" >> $GITHUB_OUTPUT
|
||||
echo "tag=$TAG" >> $GITHUB_OUTPUT
|
||||
echo "Integration: $INTEGRATION, Version: $VERSION"
|
||||
|
||||
- name: Detect integration type
|
||||
id: type
|
||||
run: |
|
||||
if [ -f "hindsight-integrations/${{ steps.info.outputs.integration }}/pyproject.toml" ]; then
|
||||
echo "type=python" >> $GITHUB_OUTPUT
|
||||
elif [ -f "hindsight-integrations/${{ steps.info.outputs.integration }}/package.json" ]; then
|
||||
echo "type=typescript" >> $GITHUB_OUTPUT
|
||||
else
|
||||
echo "type=plugin" >> $GITHUB_OUTPUT
|
||||
fi
|
||||
|
||||
# ── Python integrations (litellm, pydantic-ai, crewai) ──────────────────
|
||||
|
||||
- name: Install uv
|
||||
if: steps.type.outputs.type == 'python'
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.type.outputs.type == 'python'
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build Python package
|
||||
if: steps.type.outputs.type == 'python'
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Publish Python package to PyPI
|
||||
if: steps.type.outputs.type == 'python'
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/${{ steps.info.outputs.integration }}/dist
|
||||
skip-existing: true
|
||||
|
||||
# ── TypeScript integrations (ai-sdk, chat, openclaw) ────────────────────
|
||||
|
||||
# ── Plugin integrations (claude-code) — no package to publish ───────────
|
||||
|
||||
- name: Plugin release
|
||||
if: steps.type.outputs.type == 'plugin'
|
||||
run: |
|
||||
echo "Plugin integration ${{ steps.info.outputs.integration }} v${{ steps.info.outputs.version }} — no package to publish."
|
||||
echo "Users install via: claude plugin marketplace add vectorize-io/hindsight --sparse hindsight-integrations"
|
||||
|
||||
# ── TypeScript integrations (ai-sdk, chat, openclaw) ────────────────────
|
||||
|
||||
- name: Set up Node.js
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
|
||||
# Guard: fail fast if the integration's lockfile resolves any dep from a
|
||||
# monorepo workspace (link=true) or a relative file path. The release
|
||||
# runner has no pre-built workspace `dist/` so `npm run build` would
|
||||
# later fail at tsc with "Cannot find module". See:
|
||||
# https://github.com/vectorize-io/hindsight/issues/… (0.6.0 openclaw retry)
|
||||
- name: Check integration lockfile
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
run: ./scripts/check-integration-lockfiles.sh
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
run: npm ci
|
||||
|
||||
- name: Build TypeScript package
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
run: npm run build
|
||||
|
||||
- name: Publish TypeScript package to npm
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
run: |
|
||||
set +e
|
||||
OUTPUT=$(npm publish --access public 2>&1)
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
if echo "$OUTPUT" | grep -q "cannot publish over"; then
|
||||
echo "Package version already published, skipping..."
|
||||
exit 0
|
||||
fi
|
||||
exit $EXIT_CODE
|
||||
fi
|
||||
env:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
|
||||
+20
-193
@@ -46,30 +46,10 @@ jobs:
|
||||
working-directory: ./hindsight-all-slim
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-litellm
|
||||
working-directory: ./hindsight-integrations/litellm
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-embed
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-crewai
|
||||
working-directory: ./hindsight-integrations/crewai
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-pydantic-ai
|
||||
working-directory: ./hindsight-integrations/pydantic-ai
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-hermes
|
||||
working-directory: ./hindsight-integrations/hermes
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-agno
|
||||
working-directory: ./hindsight-integrations/agno
|
||||
run: uv build --out-dir dist
|
||||
|
||||
# Publish in order (client and api-slim first, then api/all wrappers which depend on them)
|
||||
- name: Publish hindsight-client to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
@@ -101,42 +81,12 @@ jobs:
|
||||
packages-dir: ./hindsight-all-slim/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-litellm to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/litellm/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-embed to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-embed/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-crewai to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/crewai/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-pydantic-ai to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/pydantic-ai/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-hermes to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/hermes/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-agno to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/agno/dist
|
||||
skip-existing: true
|
||||
|
||||
# Upload artifacts for GitHub release
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
@@ -148,12 +98,7 @@ jobs:
|
||||
hindsight-api/dist/*
|
||||
hindsight-all/dist/*
|
||||
hindsight-all-slim/dist/*
|
||||
hindsight-integrations/litellm/dist/*
|
||||
hindsight-embed/dist/*
|
||||
hindsight-integrations/crewai/dist/*
|
||||
hindsight-integrations/pydantic-ai/dist/*
|
||||
hindsight-integrations/hermes/dist/*
|
||||
hindsight-integrations/agno/dist/*
|
||||
retention-days: 1
|
||||
|
||||
release-typescript-client:
|
||||
@@ -205,7 +150,7 @@ jobs:
|
||||
path: hindsight-clients/typescript/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-openclaw-integration:
|
||||
release-hindsight-all-npm:
|
||||
runs-on: ubuntu-latest
|
||||
environment: npm
|
||||
|
||||
@@ -217,17 +162,17 @@ jobs:
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: package-lock.json
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm ci
|
||||
run: npm ci --workspace=hindsight-all-npm
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm run build
|
||||
run: npm run build --workspace=hindsight-all-npm
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
working-directory: ./hindsight-all-npm
|
||||
run: |
|
||||
set +e
|
||||
OUTPUT=$(npm publish --access public 2>&1)
|
||||
@@ -244,112 +189,14 @@ jobs:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
|
||||
|
||||
- name: Pack for GitHub release
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
working-directory: ./hindsight-all-npm
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: openclaw-integration
|
||||
path: hindsight-integrations/openclaw/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-ai-sdk-integration:
|
||||
runs-on: ubuntu-latest
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: npm ci
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: npm run build
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: |
|
||||
set +e
|
||||
OUTPUT=$(npm publish --access public 2>&1)
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
if echo "$OUTPUT" | grep -q "cannot publish over"; then
|
||||
echo "Package version already published, skipping..."
|
||||
exit 0
|
||||
fi
|
||||
exit $EXIT_CODE
|
||||
fi
|
||||
env:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
|
||||
|
||||
- name: Pack for GitHub release
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: ai-sdk-integration
|
||||
path: hindsight-integrations/ai-sdk/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-chat-integration:
|
||||
runs-on: ubuntu-latest
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/chat
|
||||
run: npm ci
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/chat
|
||||
run: npm run build
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-integrations/chat
|
||||
run: |
|
||||
set +e
|
||||
OUTPUT=$(npm publish --access public 2>&1)
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
if echo "$OUTPUT" | grep -q "cannot publish over"; then
|
||||
echo "Package version already published, skipping..."
|
||||
exit 0
|
||||
fi
|
||||
exit $EXIT_CODE
|
||||
fi
|
||||
env:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
|
||||
|
||||
- name: Pack for GitHub release
|
||||
working-directory: ./hindsight-integrations/chat
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: chat-integration
|
||||
path: hindsight-integrations/chat/*.tgz
|
||||
name: hindsight-all-npm
|
||||
path: hindsight-all-npm/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-control-plane:
|
||||
@@ -584,7 +431,7 @@ jobs:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Install Helm
|
||||
uses: azure/setup-helm@v4
|
||||
uses: azure/setup-helm@v5
|
||||
with:
|
||||
version: 'latest'
|
||||
|
||||
@@ -609,7 +456,7 @@ jobs:
|
||||
|
||||
create-github-release:
|
||||
runs-on: ubuntu-latest
|
||||
needs: [release-python-packages, release-typescript-client, release-openclaw-integration, release-ai-sdk-integration, release-chat-integration, release-control-plane, release-rust-cli, release-docker-images, release-helm-chart]
|
||||
needs: [release-python-packages, release-typescript-client, release-hindsight-all-npm, release-control-plane, release-rust-cli, release-docker-images, release-helm-chart]
|
||||
permissions:
|
||||
contents: write
|
||||
|
||||
@@ -632,30 +479,18 @@ jobs:
|
||||
name: typescript-client
|
||||
path: ./artifacts/typescript-client
|
||||
|
||||
- name: Download OpenClaw Integration
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: openclaw-integration
|
||||
path: ./artifacts/openclaw-integration
|
||||
|
||||
- name: Download AI SDK Integration
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: ai-sdk-integration
|
||||
path: ./artifacts/ai-sdk-integration
|
||||
|
||||
- name: Download Chat Integration
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: chat-integration
|
||||
path: ./artifacts/chat-integration
|
||||
|
||||
- name: Download Control Plane
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: control-plane
|
||||
path: ./artifacts/control-plane
|
||||
|
||||
- name: Download hindsight-embed npm wrapper
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: hindsight-all-npm
|
||||
path: ./artifacts/hindsight-all-npm
|
||||
|
||||
- name: Download Rust CLI (Linux)
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
@@ -689,19 +524,11 @@ jobs:
|
||||
cp artifacts/python-packages/hindsight-api/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-all/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-all-slim/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-integrations/litellm/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-integrations/pydantic-ai/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-integrations/hermes/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-integrations/agno/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-embed/dist/* release-assets/ || true
|
||||
# TypeScript client
|
||||
cp artifacts/typescript-client/*.tgz release-assets/ || true
|
||||
# OpenClaw Integration
|
||||
cp artifacts/openclaw-integration/*.tgz release-assets/ || true
|
||||
# AI SDK Integration
|
||||
cp artifacts/ai-sdk-integration/*.tgz release-assets/ || true
|
||||
# Chat Integration
|
||||
cp artifacts/chat-integration/*.tgz release-assets/ || true
|
||||
# hindsight-embed npm wrapper
|
||||
cp artifacts/hindsight-all-npm/*.tgz release-assets/ || true
|
||||
# Control Plane
|
||||
cp artifacts/control-plane/*.tgz release-assets/ || true
|
||||
# Rust CLI binaries
|
||||
@@ -713,7 +540,7 @@ jobs:
|
||||
ls -la release-assets/
|
||||
|
||||
- name: Create GitHub Release
|
||||
uses: softprops/action-gh-release@v2
|
||||
uses: softprops/action-gh-release@v3
|
||||
with:
|
||||
files: release-assets/*
|
||||
generate_release_notes: true
|
||||
|
||||
+1008
-12
File diff suppressed because it is too large
Load Diff
+2
-1
@@ -50,7 +50,8 @@ hindsight-dev/benchmarks/perf/results/
|
||||
benchmarks/results/
|
||||
hindsight-cli/target
|
||||
hindsight-clients/rust/target
|
||||
.claude
|
||||
.claude/*
|
||||
!.claude/skills/
|
||||
whats-next.md
|
||||
TASK.md
|
||||
# Changelog is now tracked in hindsight-docs/src/pages/changelog.md
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
{
|
||||
"semi": true,
|
||||
"singleQuote": false,
|
||||
"tabWidth": 2,
|
||||
"trailingComma": "es5",
|
||||
"printWidth": 100
|
||||
}
|
||||
@@ -11,9 +11,15 @@ Hindsight is an agent memory system that provides long-term memory for AI agents
|
||||
|
||||
## Development Commands
|
||||
|
||||
### Local Development (API + UI)
|
||||
```bash
|
||||
# Start both API server and control plane UI
|
||||
./scripts/dev/start.sh
|
||||
```
|
||||
|
||||
### API Server (Python/FastAPI)
|
||||
```bash
|
||||
# Start API server (loads .env automatically)
|
||||
# Start API server only (loads .env automatically)
|
||||
./scripts/dev/start-api.sh
|
||||
|
||||
# Run all tests (parallelized with pytest-xdist)
|
||||
@@ -73,17 +79,16 @@ cd hindsight-control-plane && npm run dev
|
||||
|
||||
### Monorepo Structure
|
||||
- **hindsight-api-slim/**: Core FastAPI server with memory engine (Python, uv)
|
||||
- **hindsight/**: Embedded Python bundle (hindsight-all package)
|
||||
- **hindsight-control-plane/**: Admin UI (Next.js, npm)
|
||||
- **hindsight-cli/**: CLI tool (Rust, cargo, uses progenitor for API client)
|
||||
- **hindsight-clients/**: Generated SDK clients (Python, TypeScript, Rust)
|
||||
- **hindsight-docs/**: Docusaurus documentation site
|
||||
- **hindsight-integrations/**: Framework integrations (LiteLLM, OpenAI)
|
||||
- **hindsight-integrations/**: Framework integrations (LiteLLM, CrewAI, LangGraph, Pydantic AI, AG2, Claude Code, etc.)
|
||||
- **hindsight-dev/**: Development tools and benchmarks
|
||||
|
||||
### Core Engine (hindsight-api-slim/hindsight_api/engine/)
|
||||
- `memory_engine.py`: Main orchestrator (~170KB) for retain/recall/reflect operations
|
||||
- `llm_wrapper.py`: LLM abstraction supporting OpenAI, Anthropic, Gemini, Groq, MiniMax, Ollama, LM Studio
|
||||
- `memory_engine.py`: Main orchestrator for retain/recall/reflect operations
|
||||
- `llm_wrapper.py`: LLM abstraction supporting OpenAI, Anthropic, Gemini, VertexAI, Groq, MiniMax, Ollama, LM Studio, LiteLLM, Claude Code
|
||||
- `embeddings.py`: Embedding generation (local sentence-transformers or TEI)
|
||||
- `cross_encoder.py`: Reranking (local or TEI)
|
||||
- `entity_resolver.py`: Entity extraction and normalization
|
||||
@@ -96,13 +101,13 @@ cd hindsight-control-plane && npm run dev
|
||||
|
||||
**search/**: Multi-strategy retrieval
|
||||
- `retrieval.py`: Main retrieval orchestrator
|
||||
- `graph_retrieval.py`: Entity/relationship graph traversal
|
||||
- `mpfp_retrieval.py`: Multi-Path Fact Propagation retrieval
|
||||
- `graph_retrieval.py`: Graph retrieval abstract base class
|
||||
- `link_expansion_retrieval.py`: Link expansion graph retrieval
|
||||
- `fusion.py`: Reciprocal rank fusion for combining results
|
||||
- `reranking.py`: Cross-encoder reranking
|
||||
|
||||
### API Layer (hindsight-api-slim/hindsight_api/api/)
|
||||
- `http.py`: FastAPI HTTP routers (~80KB) for all REST endpoints
|
||||
- `http.py`: FastAPI HTTP routers for all REST endpoints
|
||||
- `mcp.py`: Model Context Protocol server implementation
|
||||
|
||||
Main operations:
|
||||
@@ -164,11 +169,17 @@ Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links`
|
||||
## Key Conventions
|
||||
|
||||
### Code Quality
|
||||
|
||||
**Before writing code, read `.claude/skills/code-review/SKILL.md`** for the full coding standards (Python style, type safety, TypeScript style, general principles).
|
||||
|
||||
**Always run the lint script after making Python or TypeScript/Node changes:**
|
||||
```bash
|
||||
./scripts/hooks/lint.sh
|
||||
```
|
||||
This runs the same checks as the pre-commit hook (Ruff for Python, ESLint/Prettier for TypeScript).
|
||||
|
||||
**After completing any implementation work, run `/code-review`** to verify your changes against project standards (missing tests, dead code, type safety, etc.). Fix any "must fix" issues before considering the task done.
|
||||
|
||||
**MANDATORY: Run `/code-review` before pushing code or creating a pull request.** Do not push or create a PR until all "must fix" issues are resolved.
|
||||
|
||||
### Memory Banks
|
||||
- Each bank is an isolated memory store (like a "brain" for one user/agent)
|
||||
@@ -200,48 +211,20 @@ When adding or modifying parameters in the dataplane API (hindsight-api), you mu
|
||||
- Update the client type definition in `lib/api.ts`
|
||||
- Update any UI components that need to use the new parameter
|
||||
|
||||
### Python Style
|
||||
- Python 3.11+, type hints required
|
||||
- Async throughout (asyncpg, async FastAPI)
|
||||
- Pydantic models for request/response
|
||||
- Ruff for linting (line-length 120)
|
||||
- No Python files at project root - maintain clean directory structure
|
||||
- **Never use multi-item tuple return values** - prefer dataclass or Pydantic model for structured returns
|
||||
### Adding New Integrations
|
||||
|
||||
### Type Safety with Pydantic Models
|
||||
**NEVER use raw `dict` types for structured data.** Always use Pydantic models:
|
||||
- Use Pydantic `BaseModel` for all data structures passed between functions
|
||||
- Add `@field_validator` for type coercion (e.g., ensuring datetimes are timezone-aware)
|
||||
- Avoid `dict.get()` patterns - use typed model attributes instead
|
||||
- Parse external data (JSON, API responses) into Pydantic models at the boundary
|
||||
- This catches type errors at parse time, not deep in business logic
|
||||
Every new integration in `hindsight-integrations/` must satisfy all of the following before it can be merged:
|
||||
|
||||
```python
|
||||
# BAD - error-prone dict access
|
||||
def process(data: dict) -> str:
|
||||
return data.get("name", "") # No validation, silent failures
|
||||
1. **Tests are required** — tests must simulate or exercise the external system (mock the framework's interfaces and verify the integration actually calls Hindsight correctly). Pure unit tests of helper functions are not sufficient.
|
||||
2. **CI job** — add a test job in `.github/workflows/test.yml` following the existing pattern (e.g., `test-crewai-integration`). The job must build, install deps, and run `uv run pytest tests -v`. Also add the integration to `detect-changes` outputs so it only runs when its files change.
|
||||
3. **Release process** — add the integration name to the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh` so it can be released via the standard release workflow.
|
||||
4. **Follow project code standards** — Python style, type safety, no raw dicts for structured data, no multi-item tuple returns (see `.claude/skills/code-review/SKILL.md`).
|
||||
|
||||
# GOOD - typed and validated
|
||||
class UserData(BaseModel):
|
||||
name: str
|
||||
created_at: datetime
|
||||
If any of these are missing, the integration is incomplete and must not be pushed or merged.
|
||||
|
||||
@field_validator("created_at", mode="before")
|
||||
@classmethod
|
||||
def ensure_tz_aware(cls, v):
|
||||
if isinstance(v, str):
|
||||
v = datetime.fromisoformat(v.replace("Z", "+00:00"))
|
||||
if v.tzinfo is None:
|
||||
return v.replace(tzinfo=timezone.utc)
|
||||
return v
|
||||
### Changelogs
|
||||
|
||||
def process(data: UserData) -> str:
|
||||
return data.name # Type-safe, validated at construction
|
||||
```
|
||||
|
||||
### TypeScript Style
|
||||
- Next.js App Router for control plane
|
||||
- Tailwind CSS with shadcn/ui components
|
||||
Never add "Unreleased" entries to changelogs (e.g. `hindsight-docs/src/pages/changelog/**`). Changelog entries are written by the release script (`./scripts/release-integration.sh`) when a version is actually cut. If a bug fix or feature needs documenting before release, describe it in the PR/commit — the release tooling will surface it in the published changelog section.
|
||||
|
||||
### Adding New API Configuration Flags
|
||||
|
||||
@@ -255,17 +238,17 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
|
||||
- Add `ENV_*` constant for the environment variable name (e.g., `ENV_MY_SETTING = "HINDSIGHT_API_MY_SETTING"`)
|
||||
- Add `DEFAULT_*` constant for the default value
|
||||
- Add field to `HindsightConfig` dataclass with type annotation
|
||||
- **Mark as hierarchical or static** by adding to `_HIERARCHICAL_FIELDS` set (hierarchical) or leaving it out (static)
|
||||
- **Mark as configurable** by adding to `_CONFIGURABLE_FIELDS` set if the field should be overridable per-tenant/bank via API
|
||||
- Add initialization in `from_env()` method
|
||||
|
||||
```python
|
||||
# Hierarchical field (can be overridden per-bank)
|
||||
_HIERARCHICAL_FIELDS = {
|
||||
# Configurable field (can be overridden per-tenant/bank via API)
|
||||
_CONFIGURABLE_FIELDS = {
|
||||
...,
|
||||
"my_setting", # Add here for hierarchical
|
||||
"my_setting", # Add here for configurable
|
||||
}
|
||||
|
||||
# Static field - just don't add to _HIERARCHICAL_FIELDS
|
||||
# Static field - just don't add to _CONFIGURABLE_FIELDS
|
||||
```
|
||||
|
||||
2. **main.py** (`hindsight-api-slim/hindsight_api/main.py`):
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
|
||||
[](https://join.slack.com/t/hindsight-space/shared_invite/zt-3nhbm4w29-LeSJ5Ixi6j8PdiYOCPlOgg)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
[](https://gitcgr.com/vectorize-io/hindsight)
|
||||

|
||||

|
||||
<br/>
|
||||
|
||||
@@ -1,6 +1,28 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
# =============================================================================
|
||||
# Embedded pg0 data integrity check (#675)
|
||||
#
|
||||
# When using embedded pg0, check if the data directory has existing PostgreSQL
|
||||
# data before starting. If the directory exists but appears empty/corrupt
|
||||
# (e.g., missing PG_VERSION file), log a warning. This helps diagnose data
|
||||
# loss scenarios where a container restart caused the data directory to be
|
||||
# wiped despite a volume mount being present.
|
||||
# =============================================================================
|
||||
PG0_DATA_DIR="${HOME}/.pg0"
|
||||
if [ -d "$PG0_DATA_DIR" ]; then
|
||||
# Look for actual PostgreSQL data directories (pg0 creates subdirs per instance)
|
||||
if compgen -G "$PG0_DATA_DIR"/*/PG_VERSION > /dev/null 2>&1; then
|
||||
echo "✅ Existing pg0 data directory detected at $PG0_DATA_DIR"
|
||||
elif [ "$(ls -A "$PG0_DATA_DIR" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $PG0_DATA_DIR but no PG_VERSION found."
|
||||
echo " This may indicate data corruption or an incomplete previous shutdown."
|
||||
echo " If you see all migrations running from scratch after this, your data may have been lost."
|
||||
echo " See: https://github.com/vectorize-io/hindsight/issues/675"
|
||||
fi
|
||||
fi
|
||||
|
||||
# Service flags (default to true if not set)
|
||||
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
|
||||
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
|
||||
@@ -71,6 +93,63 @@ if [ "${HINDSIGHT_WAIT_FOR_DEPS:-false}" = "true" ]; then
|
||||
done
|
||||
fi
|
||||
|
||||
# =============================================================================
|
||||
# Graceful shutdown handler (#675)
|
||||
#
|
||||
# Docker sends SIGTERM on `docker stop`/`docker restart`. Without a trap, child
|
||||
# processes (hindsight-api + pg0, control-plane) are killed abruptly. For the
|
||||
# embedded pg0 database this can cause data loss when the data directory is on
|
||||
# a Docker volume that gets remounted after restart.
|
||||
#
|
||||
# The trap forwards SIGTERM to all tracked child PIDs so that:
|
||||
# - hindsight-api receives the signal and can run its shutdown hooks
|
||||
# - pg0 gets a clean PostgreSQL shutdown (checkpoint + WAL flush)
|
||||
# - The control-plane Node.js process exits cleanly
|
||||
# =============================================================================
|
||||
# Guard against concurrent cleanup (e.g., child crash + SIGTERM arriving together)
|
||||
SHUTTING_DOWN=false
|
||||
|
||||
cleanup() {
|
||||
if $SHUTTING_DOWN; then return; fi
|
||||
SHUTTING_DOWN=true
|
||||
|
||||
echo ""
|
||||
echo "🛑 Received shutdown signal, stopping services gracefully..."
|
||||
for pid in "${PIDS[@]}"; do
|
||||
if kill -0 "$pid" 2>/dev/null; then
|
||||
kill -TERM "$pid" 2>/dev/null
|
||||
fi
|
||||
done
|
||||
# Give processes time to shut down cleanly (pg0 needs to flush WAL).
|
||||
# NOTE: Docker's default stop_grace_period is 10s. If you use the default,
|
||||
# either set stop_grace_period: 30s in your compose file / docker stop -t 30,
|
||||
# or Docker will SIGKILL the container before this timeout expires.
|
||||
local timeout=30
|
||||
for ((i=1; i<=timeout; i++)); do
|
||||
local all_stopped=true
|
||||
for pid in "${PIDS[@]}"; do
|
||||
if kill -0 "$pid" 2>/dev/null; then
|
||||
all_stopped=false
|
||||
break
|
||||
fi
|
||||
done
|
||||
if $all_stopped; then
|
||||
echo "✅ All services stopped cleanly"
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
# Force kill if still running after timeout
|
||||
echo "⚠️ Timeout reached, forcing shutdown..."
|
||||
for pid in "${PIDS[@]}"; do
|
||||
if kill -0 "$pid" 2>/dev/null; then
|
||||
kill -9 "$pid" 2>/dev/null
|
||||
fi
|
||||
done
|
||||
exit 1
|
||||
}
|
||||
trap cleanup SIGTERM SIGINT
|
||||
|
||||
# Track PIDs for wait
|
||||
PIDS=()
|
||||
|
||||
@@ -138,8 +217,21 @@ if [ ${#PIDS[@]} -eq 0 ]; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Wait for any process to exit
|
||||
wait -n
|
||||
|
||||
# Exit with status of first exited process
|
||||
exit $?
|
||||
# Wait for any process to exit (use wait -n with trap-safe loop)
|
||||
while true; do
|
||||
# wait -n returns when any child exits; it also returns on signal delivery
|
||||
# (the trap handler will run and exit, so this loop is just for robustness).
|
||||
# `&& true` prevents `set -e` from killing the script when wait -n returns
|
||||
# non-zero (child exited with error or no backgrounded children remain).
|
||||
wait -n && true
|
||||
# Check if any tracked PID has exited
|
||||
for pid in "${PIDS[@]}"; do
|
||||
if ! kill -0 "$pid" 2>/dev/null; then
|
||||
wait "$pid" 2>/dev/null
|
||||
exit_code=$?
|
||||
echo "⚠️ Service (PID $pid) exited with code $exit_code"
|
||||
# Trigger cleanup for remaining services
|
||||
cleanup
|
||||
fi
|
||||
done
|
||||
done
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.4.19
|
||||
appVersion: "0.4.19"
|
||||
version: 0.5.4
|
||||
appVersion: "0.5.4"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -95,6 +95,27 @@ spec:
|
||||
{{- toYaml .Values.api.readinessProbe | nindent 10 }}
|
||||
resources:
|
||||
{{- toYaml .Values.api.resources | nindent 10 }}
|
||||
{{- if or .Values.api.persistence.modelCache.enabled .Values.api.extraVolumeMounts }}
|
||||
volumeMounts:
|
||||
{{- if .Values.api.persistence.modelCache.enabled }}
|
||||
- name: model-cache
|
||||
mountPath: /home/hindsight/.cache
|
||||
{{- end }}
|
||||
{{- with .Values.api.extraVolumeMounts }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- if or .Values.api.persistence.modelCache.enabled .Values.api.extraVolumes }}
|
||||
volumes:
|
||||
{{- if .Values.api.persistence.modelCache.enabled }}
|
||||
- name: model-cache
|
||||
persistentVolumeClaim:
|
||||
claimName: {{ include "hindsight.fullname" . }}-api-model-cache
|
||||
{{- end }}
|
||||
{{- with .Values.api.extraVolumes }}
|
||||
{{- toYaml . | nindent 6 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
{{- if and .Values.api.enabled .Values.api.persistence.modelCache.enabled }}
|
||||
apiVersion: v1
|
||||
kind: PersistentVolumeClaim
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-api-model-cache
|
||||
labels:
|
||||
{{- include "hindsight.api.labels" . | nindent 4 }}
|
||||
{{- with .Values.api.persistence.modelCache.annotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
accessModes:
|
||||
{{- toYaml .Values.api.persistence.modelCache.accessModes | nindent 4 }}
|
||||
{{- if .Values.api.persistence.modelCache.storageClass }}
|
||||
storageClassName: {{ .Values.api.persistence.modelCache.storageClass }}
|
||||
{{- end }}
|
||||
resources:
|
||||
requests:
|
||||
storage: {{ .Values.api.persistence.modelCache.size }}
|
||||
{{- end }}
|
||||
@@ -95,6 +95,16 @@ spec:
|
||||
{{- toYaml .Values.worker.readinessProbe | nindent 10 }}
|
||||
resources:
|
||||
{{- toYaml .Values.worker.resources | nindent 10 }}
|
||||
{{- if or .Values.worker.persistence.modelCache.enabled .Values.worker.extraVolumeMounts }}
|
||||
volumeMounts:
|
||||
{{- if .Values.worker.persistence.modelCache.enabled }}
|
||||
- name: model-cache
|
||||
mountPath: /home/hindsight/.cache
|
||||
{{- end }}
|
||||
{{- with .Values.worker.extraVolumeMounts }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
@@ -107,4 +117,26 @@ spec:
|
||||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.worker.extraVolumes }}
|
||||
volumes:
|
||||
{{- toYaml . | nindent 6 }}
|
||||
{{- end }}
|
||||
{{- if .Values.worker.persistence.modelCache.enabled }}
|
||||
volumeClaimTemplates:
|
||||
- metadata:
|
||||
name: model-cache
|
||||
{{- with .Values.worker.persistence.modelCache.annotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
accessModes:
|
||||
{{- toYaml .Values.worker.persistence.modelCache.accessModes | nindent 8 }}
|
||||
{{- if .Values.worker.persistence.modelCache.storageClass }}
|
||||
storageClassName: {{ .Values.worker.persistence.modelCache.storageClass }}
|
||||
{{- end }}
|
||||
resources:
|
||||
requests:
|
||||
storage: {{ .Values.worker.persistence.modelCache.size }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
@@ -67,6 +67,33 @@ api:
|
||||
# Pod affinity/anti-affinity (overrides global affinity for this component)
|
||||
# affinity: {}
|
||||
|
||||
# Persistent volume for local model cache (reranker, embeddings)
|
||||
# Models are downloaded to /home/hindsight/.cache on first use.
|
||||
# Without persistence, models are re-downloaded on every pod restart.
|
||||
persistence:
|
||||
modelCache:
|
||||
enabled: false
|
||||
size: 5Gi
|
||||
storageClass: ""
|
||||
accessModes:
|
||||
- ReadWriteOnce
|
||||
annotations: {}
|
||||
|
||||
# Extra volume mounts for the api container
|
||||
# e.g.
|
||||
# extraVolumeMounts:
|
||||
# - name: my-volume
|
||||
# mountPath: /mnt/my-volume
|
||||
extraVolumeMounts: []
|
||||
|
||||
# Extra volumes for the api pod
|
||||
# e.g.
|
||||
# extraVolumes:
|
||||
# - name: my-volume
|
||||
# configMap:
|
||||
# name: my-configmap
|
||||
extraVolumes: []
|
||||
|
||||
# Environment variables
|
||||
env:
|
||||
#HINDSIGHT_API_LLM_PROVIDER: "groq"
|
||||
@@ -140,6 +167,32 @@ worker:
|
||||
# Pod affinity/anti-affinity (overrides global affinity for this component)
|
||||
# affinity: {}
|
||||
|
||||
# Persistent volume for local model cache (reranker, embeddings)
|
||||
# Uses volumeClaimTemplates since worker is a StatefulSet.
|
||||
persistence:
|
||||
modelCache:
|
||||
enabled: false
|
||||
size: 5Gi
|
||||
storageClass: ""
|
||||
accessModes:
|
||||
- ReadWriteOnce
|
||||
annotations: {}
|
||||
|
||||
# Extra volume mounts for the worker container
|
||||
# e.g.
|
||||
# extraVolumeMounts:
|
||||
# - name: my-volume
|
||||
# mountPath: /mnt/my-volume
|
||||
extraVolumeMounts: []
|
||||
|
||||
# Extra volumes for the worker pod
|
||||
# e.g.
|
||||
# extraVolumes:
|
||||
# - name: my-volume
|
||||
# configMap:
|
||||
# name: my-configmap
|
||||
extraVolumes: []
|
||||
|
||||
# Secret environment variables (inherited from api.secrets if not specified)
|
||||
secrets: {}
|
||||
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
node_modules
|
||||
dist
|
||||
*.tgz
|
||||
.DS_Store
|
||||
@@ -0,0 +1,80 @@
|
||||
# @vectorize-io/hindsight-all
|
||||
|
||||
Node.js equivalent of the Python [`hindsight-all`](https://pypi.org/project/hindsight-all/) package — programmatic lifecycle manager for a local Hindsight daemon. Use this when you want to embed Hindsight in a Node application without hand-rolling subprocess management.
|
||||
|
||||
This package deliberately does **not** ship an HTTP client. Once the daemon is running, talk to it with [`@vectorize-io/hindsight-client`](https://www.npmjs.com/package/@vectorize-io/hindsight-client) against `server.getBaseUrl()`. The two packages compose — one owns the daemon process, the other owns the HTTP API surface.
|
||||
|
||||
## Requirements
|
||||
|
||||
- **Node.js >= 22** — uses global `fetch` and `AbortSignal.timeout`.
|
||||
- **`uv` / `uvx`** on `PATH` — used to download and run the underlying `hindsight-embed` daemon on first use. Install via <https://docs.astral.sh/uv/>.
|
||||
|
||||
## Install
|
||||
|
||||
```bash
|
||||
npm install @vectorize-io/hindsight-all @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
## Example
|
||||
|
||||
```ts
|
||||
import { HindsightServer, consoleLogger } from '@vectorize-io/hindsight-all';
|
||||
import { HindsightClient } from '@vectorize-io/hindsight-client';
|
||||
|
||||
const server = new HindsightServer({
|
||||
profile: 'my-app',
|
||||
port: 9077,
|
||||
env: {
|
||||
HINDSIGHT_API_LLM_PROVIDER: 'anthropic',
|
||||
HINDSIGHT_API_LLM_API_KEY: process.env.ANTHROPIC_API_KEY,
|
||||
HINDSIGHT_API_LLM_MODEL: 'claude-sonnet-4-20250514',
|
||||
HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT: '0',
|
||||
},
|
||||
logger: consoleLogger,
|
||||
});
|
||||
|
||||
await server.start();
|
||||
|
||||
const client = new HindsightClient({ baseUrl: server.getBaseUrl() });
|
||||
|
||||
await client.retain('user-123', 'User prefers dark mode and concise answers.', {
|
||||
documentId: 'pref-2026-04-01',
|
||||
});
|
||||
|
||||
const recall = await client.recall('user-123', 'what are the user preferences?');
|
||||
console.log(recall.results);
|
||||
|
||||
await server.stop();
|
||||
```
|
||||
|
||||
For a remote Hindsight API, skip `HindsightServer` entirely and just point `HindsightClient` at the remote URL.
|
||||
|
||||
## Open config — forward-compatible with new daemon flags
|
||||
|
||||
`HindsightServerOptions` is designed so every new environment variable or CLI flag in the underlying Hindsight daemon can be used without waiting for a wrapper release:
|
||||
|
||||
- **`env`** accepts an arbitrary `Record<string, string>`. Every entry is exported into the daemon process and written into the profile config via `--env KEY=VALUE`.
|
||||
- **`extraProfileCreateArgs`** / **`extraDaemonStartArgs`** append raw args to the respective commands.
|
||||
|
||||
## Development against a local checkout
|
||||
|
||||
If you're hacking on the Python `hindsight-embed` package in the same monorepo, point the server at the local path — it'll use `uv run --directory <path>` instead of `uvx`:
|
||||
|
||||
```ts
|
||||
new HindsightServer({
|
||||
embedPackagePath: '/path/to/hindsight-embed',
|
||||
// ...
|
||||
});
|
||||
```
|
||||
|
||||
## API surface
|
||||
|
||||
- `HindsightServer` — daemon lifecycle (`start`, `stop`, `checkHealth`, `getBaseUrl`, `getProfile`).
|
||||
- `Logger` interface plus `silentLogger` (default) and `consoleLogger` helpers.
|
||||
- `getEmbedCommand(opts)` — low-level helper that returns the `[cmd, ...args]` tuple used to invoke the underlying Python CLI.
|
||||
|
||||
For memory operations (retain, recall, reflect, bank management, stats) use [`@vectorize-io/hindsight-client`](https://www.npmjs.com/package/@vectorize-io/hindsight-client).
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1,57 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.5.4",
|
||||
"description": "Node.js programmatic lifecycle manager for Hindsight — embeds a local hindsight daemon in a Node application. Pair with @vectorize-io/hindsight-client for memory operations.",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
"type": "module",
|
||||
"exports": {
|
||||
".": {
|
||||
"types": "./dist/index.d.ts",
|
||||
"import": "./dist/index.js"
|
||||
}
|
||||
},
|
||||
"keywords": [
|
||||
"hindsight",
|
||||
"hindsight-all",
|
||||
"memory",
|
||||
"ai",
|
||||
"agent",
|
||||
"long-term-memory",
|
||||
"llm",
|
||||
"embedded-server"
|
||||
],
|
||||
"author": "Vectorize <[email protected]>",
|
||||
"license": "MIT",
|
||||
"repository": {
|
||||
"type": "git",
|
||||
"url": "https://github.com/vectorize-io/hindsight.git",
|
||||
"directory": "hindsight-all-npm"
|
||||
},
|
||||
"files": [
|
||||
"dist",
|
||||
"README.md"
|
||||
],
|
||||
"scripts": {
|
||||
"build": "tsup",
|
||||
"dev": "tsup --watch",
|
||||
"clean": "rm -rf dist",
|
||||
"test": "vitest run src",
|
||||
"test:watch": "vitest src",
|
||||
"prepublishOnly": "npm run clean && npm run build"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/node": "^22.0.0",
|
||||
"tsup": "^8.5.1",
|
||||
"typescript": "^5.7.0",
|
||||
"vitest": "^4.1.2"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=22"
|
||||
},
|
||||
"overrides": {
|
||||
"rollup": "^4.59.0",
|
||||
"picomatch": ">=2.3.2 <3.0.0 || >=4.0.4",
|
||||
"vite": ">=8.0.5"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { getEmbedCommand } from './command.js';
|
||||
|
||||
describe('getEmbedCommand', () => {
|
||||
it('defaults to uvx hindsight-embed@latest', () => {
|
||||
expect(getEmbedCommand()).toEqual(['uvx', 'hindsight-embed@latest']);
|
||||
});
|
||||
|
||||
it('honours an explicit version', () => {
|
||||
expect(getEmbedCommand({ embedVersion: '0.5.0' })).toEqual(['uvx', '[email protected]']);
|
||||
});
|
||||
|
||||
it('treats an empty version as latest', () => {
|
||||
expect(getEmbedCommand({ embedVersion: '' })).toEqual(['uvx', 'hindsight-embed@latest']);
|
||||
});
|
||||
|
||||
it('uses uv run --directory when a local path is given', () => {
|
||||
expect(getEmbedCommand({ embedPackagePath: '/abs/path' })).toEqual([
|
||||
'uv',
|
||||
'run',
|
||||
'--directory',
|
||||
'/abs/path',
|
||||
'hindsight-embed',
|
||||
]);
|
||||
});
|
||||
|
||||
it('local path takes precedence over version', () => {
|
||||
expect(
|
||||
getEmbedCommand({ embedPackagePath: '/abs/path', embedVersion: '0.5.0' }),
|
||||
).toEqual(['uv', 'run', '--directory', '/abs/path', 'hindsight-embed']);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,25 @@
|
||||
/**
|
||||
* Resolve the command that invokes the `hindsight-embed` Python CLI.
|
||||
*
|
||||
* - If `embedPackagePath` is set, runs the package from a local checkout via
|
||||
* `uv run --directory <path> hindsight-embed`. Used for in-repo development.
|
||||
* - Otherwise runs it via `uvx hindsight-embed@<version>` so no global install
|
||||
* is required.
|
||||
*
|
||||
* Returns the argv as `[command, ...baseArgs]` suitable for `spawn()` /
|
||||
* `execFile()` (never shell-interpolated).
|
||||
*/
|
||||
export interface EmbedCommandOptions {
|
||||
/** Version spec passed to uvx (e.g. "latest", "0.5.0"). Default: "latest". */
|
||||
embedVersion?: string;
|
||||
/** Local checkout path. When set, overrides `embedVersion` and uses `uv run`. */
|
||||
embedPackagePath?: string;
|
||||
}
|
||||
|
||||
export function getEmbedCommand(opts: EmbedCommandOptions = {}): string[] {
|
||||
if (opts.embedPackagePath) {
|
||||
return ['uv', 'run', '--directory', opts.embedPackagePath, 'hindsight-embed'];
|
||||
}
|
||||
const version = opts.embedVersion && opts.embedVersion.length > 0 ? opts.embedVersion : 'latest';
|
||||
return ['uvx', `hindsight-embed@${version}`];
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
export { HindsightServer } from './server.js';
|
||||
export { getEmbedCommand } from './command.js';
|
||||
export { silentLogger, consoleLogger } from './logger.js';
|
||||
|
||||
export type { Logger } from './logger.js';
|
||||
export type { EmbedCommandOptions } from './command.js';
|
||||
export type { HindsightServerOptions } from './types.js';
|
||||
@@ -0,0 +1,29 @@
|
||||
/**
|
||||
* Pluggable logger interface.
|
||||
*
|
||||
* This package does not own any logging infrastructure — consumers inject
|
||||
* whatever they want (console, pino, openclaw's logger, a no-op). The default
|
||||
* is silent so embedding this package never adds noise to an unrelated app.
|
||||
*/
|
||||
export interface Logger {
|
||||
debug(msg: string): void;
|
||||
info(msg: string): void;
|
||||
warn(msg: string): void;
|
||||
error(msg: string): void;
|
||||
}
|
||||
|
||||
/** Logger that drops every call. Used when no logger is passed. */
|
||||
export const silentLogger: Logger = {
|
||||
debug: () => {},
|
||||
info: () => {},
|
||||
warn: () => {},
|
||||
error: () => {},
|
||||
};
|
||||
|
||||
/** Logger that writes to the standard console. Handy for CLIs and tests. */
|
||||
export const consoleLogger: Logger = {
|
||||
debug: (msg) => console.debug(msg),
|
||||
info: (msg) => console.log(msg),
|
||||
warn: (msg) => console.warn(msg),
|
||||
error: (msg) => console.error(msg),
|
||||
};
|
||||
@@ -0,0 +1,35 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { HindsightServer } from './server.js';
|
||||
|
||||
describe('HindsightServer construction', () => {
|
||||
it('defaults base URL to http://127.0.0.1:8888', () => {
|
||||
const server = new HindsightServer();
|
||||
expect(server.getBaseUrl()).toBe('http://127.0.0.1:8888');
|
||||
expect(server.getProfile()).toBe('default');
|
||||
});
|
||||
|
||||
it('honours custom profile, port, and host', () => {
|
||||
const server = new HindsightServer({ profile: 'app', port: 9077, host: '0.0.0.0' });
|
||||
expect(server.getProfile()).toBe('app');
|
||||
expect(server.getBaseUrl()).toBe('http://0.0.0.0:9077');
|
||||
});
|
||||
|
||||
it('accepts open env pass-through without complaining about unknown keys', () => {
|
||||
const server = new HindsightServer({
|
||||
env: {
|
||||
HINDSIGHT_API_LLM_PROVIDER: 'openai',
|
||||
HINDSIGHT_API_LLM_MODEL: 'gpt-4o-mini',
|
||||
// A field that does not exist today — should still be accepted
|
||||
HINDSIGHT_FUTURE_FLAG: 'enabled',
|
||||
},
|
||||
});
|
||||
expect(server).toBeInstanceOf(HindsightServer);
|
||||
});
|
||||
|
||||
it('exposes checkHealth that returns false when no daemon is running', async () => {
|
||||
// Random high port that nothing is listening on.
|
||||
const server = new HindsightServer({ port: 1, readyTimeoutMs: 100 });
|
||||
const healthy = await server.checkHealth();
|
||||
expect(healthy).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,322 @@
|
||||
import { spawn } from 'child_process';
|
||||
import { getEmbedCommand } from './command.js';
|
||||
import { silentLogger } from './logger.js';
|
||||
import type { Logger } from './logger.js';
|
||||
import type { HindsightServerOptions } from './types.js';
|
||||
|
||||
const DEFAULT_PORT = 8888;
|
||||
const DEFAULT_HOST = '127.0.0.1';
|
||||
const DEFAULT_PROFILE = 'default';
|
||||
const DEFAULT_READY_TIMEOUT_MS = 30_000;
|
||||
const DEFAULT_READY_POLL_INTERVAL_MS = 1_000;
|
||||
|
||||
/**
|
||||
* Manages the lifecycle of a local Hindsight daemon from a Node.js process.
|
||||
*
|
||||
* On {@link start}, this class:
|
||||
* 1. Resolves the `hindsight-embed` command (via `uvx` or a local `uv run`).
|
||||
* 2. Runs `profile create <name> --merge --port <port> [--env K=V ...]`
|
||||
* with every entry in {@link HindsightServerOptions.env} forwarded as
|
||||
* an `--env` flag.
|
||||
* 3. Runs `daemon --profile <name> start` and waits for the start command
|
||||
* to exit.
|
||||
* 4. Polls `http://host:port/health` until it returns `200` or the
|
||||
* `readyTimeoutMs` budget is exhausted.
|
||||
*
|
||||
* On {@link stop}, it runs `daemon --profile <name> stop` and returns once
|
||||
* the command exits (or after a short grace period).
|
||||
*
|
||||
* This is the Node.js equivalent of the Python `hindsight-all` package's
|
||||
* `HindsightServer`: a thin programmatic lifecycle wrapper around the
|
||||
* Hindsight daemon. It does NOT ship an HTTP client — once `start()`
|
||||
* resolves, use `@vectorize-io/hindsight-client` against `getBaseUrl()` for
|
||||
* retain / recall / reflect.
|
||||
*
|
||||
* The class is deliberately transparent about the daemon: new CLI flags or
|
||||
* environment variables never require a code change here — callers can pass
|
||||
* them via `env`, `extraProfileCreateArgs`, or `extraDaemonStartArgs`.
|
||||
*/
|
||||
export class HindsightServer {
|
||||
private readonly profile: string;
|
||||
private readonly port: number;
|
||||
private readonly host: string;
|
||||
private readonly baseUrl: string;
|
||||
private readonly embedVersion: string | undefined;
|
||||
private readonly embedPackagePath: string | undefined;
|
||||
private readonly userEnv: Record<string, string | undefined>;
|
||||
private readonly extraProfileCreateArgs: string[];
|
||||
private readonly extraDaemonStartArgs: string[];
|
||||
private readonly platformCpuWorkaround: boolean;
|
||||
private readonly readyTimeoutMs: number;
|
||||
private readonly readyPollIntervalMs: number;
|
||||
private readonly logger: Logger;
|
||||
|
||||
constructor(opts: HindsightServerOptions = {}) {
|
||||
this.profile = opts.profile ?? DEFAULT_PROFILE;
|
||||
this.port = opts.port ?? DEFAULT_PORT;
|
||||
this.host = opts.host ?? DEFAULT_HOST;
|
||||
this.baseUrl = `http://${this.host}:${this.port}`;
|
||||
this.embedVersion = opts.embedVersion;
|
||||
this.embedPackagePath = opts.embedPackagePath;
|
||||
this.userEnv = opts.env ?? {};
|
||||
this.extraProfileCreateArgs = opts.extraProfileCreateArgs ?? [];
|
||||
this.extraDaemonStartArgs = opts.extraDaemonStartArgs ?? [];
|
||||
this.platformCpuWorkaround = opts.platformCpuWorkaround ?? (process.platform === 'darwin');
|
||||
this.readyTimeoutMs = opts.readyTimeoutMs ?? DEFAULT_READY_TIMEOUT_MS;
|
||||
this.readyPollIntervalMs = opts.readyPollIntervalMs ?? DEFAULT_READY_POLL_INTERVAL_MS;
|
||||
this.logger = opts.logger ?? silentLogger;
|
||||
}
|
||||
|
||||
/** The base URL the daemon listens on (`http://host:port`). */
|
||||
getBaseUrl(): string {
|
||||
return this.baseUrl;
|
||||
}
|
||||
|
||||
/** The profile name this server operates on. */
|
||||
getProfile(): string {
|
||||
return this.profile;
|
||||
}
|
||||
|
||||
/**
|
||||
* Ensure the daemon is configured and running. Idempotent — the underlying
|
||||
* `profile create --merge` and `daemon start` commands tolerate re-runs.
|
||||
*/
|
||||
async start(): Promise<void> {
|
||||
this.logger.info(`[hindsight] starting daemon for profile "${this.profile}"`);
|
||||
|
||||
const env = this.buildEnv();
|
||||
await this.configureProfile(env);
|
||||
await this.startDaemon(env);
|
||||
await this.waitForReady();
|
||||
|
||||
this.logger.info(`[hindsight] daemon ready at ${this.baseUrl}`);
|
||||
}
|
||||
|
||||
/** Stop the daemon. Never throws — logs and resolves even on failure. */
|
||||
async stop(): Promise<void> {
|
||||
this.logger.info(`[hindsight] stopping daemon for profile "${this.profile}"`);
|
||||
|
||||
const [cmd, ...baseArgs] = getEmbedCommand({
|
||||
embedVersion: this.embedVersion,
|
||||
embedPackagePath: this.embedPackagePath,
|
||||
});
|
||||
const args = [...baseArgs, 'daemon', '--profile', this.profile, 'stop'];
|
||||
|
||||
const child = spawn(cmd, args, { stdio: 'pipe' });
|
||||
this.pipeOutput(child, 'daemon.stop');
|
||||
|
||||
await new Promise<void>((resolve) => {
|
||||
const timeout = setTimeout(() => {
|
||||
this.logger.warn(`[hindsight] daemon stop timed out after 5s`);
|
||||
resolve();
|
||||
}, 5_000);
|
||||
child.on('exit', () => {
|
||||
clearTimeout(timeout);
|
||||
this.logger.info(`[hindsight] daemon stopped`);
|
||||
resolve();
|
||||
});
|
||||
child.on('error', (err) => {
|
||||
clearTimeout(timeout);
|
||||
this.logger.warn(`[hindsight] error stopping daemon: ${err.message}`);
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** Probe `/health` once with a short timeout. */
|
||||
async checkHealth(): Promise<boolean> {
|
||||
try {
|
||||
const res = await fetch(`${this.baseUrl}/health`, {
|
||||
signal: AbortSignal.timeout(2_000),
|
||||
});
|
||||
return res.ok;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------
|
||||
// Internal
|
||||
// -------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Merge the process env, the caller-supplied `env`, and (on macOS) the
|
||||
* embeddings CPU workaround. Caller-supplied values always win over the
|
||||
* workaround; undefined values are dropped.
|
||||
*/
|
||||
private buildEnv(): NodeJS.ProcessEnv {
|
||||
const merged: NodeJS.ProcessEnv = { ...process.env };
|
||||
|
||||
if (this.platformCpuWorkaround && process.platform === 'darwin') {
|
||||
merged['HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU'] = '1';
|
||||
merged['HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU'] = '1';
|
||||
}
|
||||
|
||||
for (const [key, value] of Object.entries(this.userEnv)) {
|
||||
if (value !== undefined) {
|
||||
merged[key] = value;
|
||||
}
|
||||
}
|
||||
|
||||
return merged;
|
||||
}
|
||||
|
||||
/**
|
||||
* Run `profile create <name> --merge --port <port> [--env K=V ...]`.
|
||||
* Every entry in the merged env that was passed via {@link userEnv} (or
|
||||
* auto-applied by the CPU workaround) is forwarded as `--env`.
|
||||
*/
|
||||
private async configureProfile(env: NodeJS.ProcessEnv): Promise<void> {
|
||||
this.logger.info(`[hindsight] configuring profile "${this.profile}"`);
|
||||
|
||||
const [cmd, ...baseArgs] = getEmbedCommand({
|
||||
embedVersion: this.embedVersion,
|
||||
embedPackagePath: this.embedPackagePath,
|
||||
});
|
||||
const createArgs = [
|
||||
...baseArgs,
|
||||
'profile',
|
||||
'create',
|
||||
this.profile,
|
||||
'--merge',
|
||||
'--port',
|
||||
String(this.port),
|
||||
];
|
||||
|
||||
// Forward every env var that the caller intended for the daemon as --env.
|
||||
// We only forward keys the caller explicitly set (userEnv) plus the CPU
|
||||
// workaround values — not the entire process.env, to avoid leaking random
|
||||
// host state into profile config.
|
||||
const envForProfile = this.collectProfileEnv(env);
|
||||
for (const [key, value] of Object.entries(envForProfile)) {
|
||||
createArgs.push('--env', `${key}=${value}`);
|
||||
}
|
||||
|
||||
createArgs.push(...this.extraProfileCreateArgs);
|
||||
|
||||
await this.runCommand(cmd, createArgs, env, 'profile.create');
|
||||
}
|
||||
|
||||
/** Collect only the env vars that should be written into the profile file. */
|
||||
private collectProfileEnv(env: NodeJS.ProcessEnv): Record<string, string> {
|
||||
const out: Record<string, string> = {};
|
||||
|
||||
// 1. User-supplied env — always forwarded.
|
||||
for (const [key, value] of Object.entries(this.userEnv)) {
|
||||
if (value !== undefined) {
|
||||
out[key] = value;
|
||||
}
|
||||
}
|
||||
|
||||
// 2. CPU workaround — only if auto-applied and not already overridden.
|
||||
if (this.platformCpuWorkaround && process.platform === 'darwin') {
|
||||
const cpuKeys = [
|
||||
'HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU',
|
||||
'HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU',
|
||||
];
|
||||
for (const key of cpuKeys) {
|
||||
if (!(key in out) && env[key] !== undefined) {
|
||||
out[key] = env[key] as string;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
private async startDaemon(env: NodeJS.ProcessEnv): Promise<void> {
|
||||
const [cmd, ...baseArgs] = getEmbedCommand({
|
||||
embedVersion: this.embedVersion,
|
||||
embedPackagePath: this.embedPackagePath,
|
||||
});
|
||||
const args = [
|
||||
...baseArgs,
|
||||
'daemon',
|
||||
'--profile',
|
||||
this.profile,
|
||||
'start',
|
||||
...this.extraDaemonStartArgs,
|
||||
];
|
||||
|
||||
await this.runCommand(cmd, args, env, 'daemon.start');
|
||||
}
|
||||
|
||||
/**
|
||||
* Spawn `cmd` with `args`, pipe its output through the logger, and resolve
|
||||
* once it exits with code 0. Rejects on non-zero exit or spawn error.
|
||||
*/
|
||||
private async runCommand(
|
||||
cmd: string,
|
||||
args: string[],
|
||||
env: NodeJS.ProcessEnv,
|
||||
label: string,
|
||||
): Promise<void> {
|
||||
const child = spawn(cmd, args, { stdio: 'pipe', env });
|
||||
let output = '';
|
||||
child.stdout?.on('data', (data: Buffer) => {
|
||||
const text = data.toString();
|
||||
output += text;
|
||||
for (const line of text.trimEnd().split('\n')) {
|
||||
if (line) this.logger.info(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
child.stderr?.on('data', (data: Buffer) => {
|
||||
const text = data.toString();
|
||||
output += text;
|
||||
for (const line of text.trimEnd().split('\n')) {
|
||||
if (line) this.logger.warn(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
child.on('exit', (code) => {
|
||||
if (code === 0) {
|
||||
resolve();
|
||||
} else {
|
||||
reject(new Error(`${label} failed with code ${code}: ${output.trim()}`));
|
||||
}
|
||||
});
|
||||
child.on('error', (err) => {
|
||||
reject(new Error(`${label} failed to spawn: ${err.message}`, { cause: err }));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** Stream a spawned child's stdout/stderr through the logger without blocking. */
|
||||
private pipeOutput(child: ReturnType<typeof spawn>, label: string): void {
|
||||
child.stdout?.on('data', (data: Buffer) => {
|
||||
for (const line of data.toString().trimEnd().split('\n')) {
|
||||
if (line) this.logger.info(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
child.stderr?.on('data', (data: Buffer) => {
|
||||
for (const line of data.toString().trimEnd().split('\n')) {
|
||||
if (line) this.logger.warn(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/** Poll `/health` until it succeeds or `readyTimeoutMs` elapses. */
|
||||
private async waitForReady(): Promise<void> {
|
||||
const deadline = Date.now() + this.readyTimeoutMs;
|
||||
let attempt = 0;
|
||||
while (Date.now() < deadline) {
|
||||
attempt++;
|
||||
try {
|
||||
const res = await fetch(`${this.baseUrl}/health`, {
|
||||
signal: AbortSignal.timeout(this.readyPollIntervalMs),
|
||||
});
|
||||
if (res.ok) {
|
||||
this.logger.debug(`[hindsight] health check passed (attempt ${attempt})`);
|
||||
return;
|
||||
}
|
||||
} catch {
|
||||
// expected while the daemon is still booting
|
||||
}
|
||||
await new Promise((resolve) => setTimeout(resolve, this.readyPollIntervalMs));
|
||||
}
|
||||
throw new Error(
|
||||
`Hindsight daemon did not become ready within ${this.readyTimeoutMs}ms at ${this.baseUrl}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
import type { Logger } from './logger.js';
|
||||
|
||||
/**
|
||||
* Options for {@link HindsightServer}.
|
||||
*
|
||||
* The server is intentionally thin and pass-through: anything configurable
|
||||
* on the daemon side (env vars or CLI flags) can be set here without needing
|
||||
* a new dedicated option. Use {@link env} for `HINDSIGHT_*` / `OPENAI_API_KEY` /
|
||||
* custom provider settings, and the two `extra*` arrays to append raw CLI
|
||||
* args to `profile create` or `daemon start`.
|
||||
*
|
||||
* For talking to the daemon after `start()`, use `@vectorize-io/hindsight-client`
|
||||
* against `server.getBaseUrl()`. This package does not ship its own HTTP
|
||||
* client.
|
||||
*/
|
||||
export interface HindsightServerOptions {
|
||||
/** Profile name used for `--profile <name>` on every sub-command. Default: `"default"`. */
|
||||
profile?: string;
|
||||
/** TCP port the daemon listens on. Default: `8888`. */
|
||||
port?: number;
|
||||
/** Hostname the daemon binds to (for health checks). Default: `127.0.0.1`. */
|
||||
host?: string;
|
||||
/** Version of the underlying `hindsight-embed` PyPI package to run via `uvx`. Default: `"latest"`. */
|
||||
embedVersion?: string;
|
||||
/** Local path to a `hindsight-embed` checkout — takes precedence over `embedVersion`. */
|
||||
embedPackagePath?: string;
|
||||
/**
|
||||
* Environment variables passed to the daemon process AND written into the
|
||||
* profile via repeated `--env KEY=VALUE` flags. This is the preferred way
|
||||
* to surface any `HINDSIGHT_API_*` / `HINDSIGHT_EMBED_*` setting — adding a
|
||||
* new daemon env var never requires a wrapper update.
|
||||
*
|
||||
* Values of `undefined` are dropped (so you can spread conditionally).
|
||||
*/
|
||||
env?: Record<string, string | undefined>;
|
||||
/** Extra args appended verbatim to `hindsight-embed profile create <name> --merge ...`. */
|
||||
extraProfileCreateArgs?: string[];
|
||||
/** Extra args appended verbatim to `hindsight-embed daemon --profile <name> start ...`. */
|
||||
extraDaemonStartArgs?: string[];
|
||||
/**
|
||||
* On macOS, automatically set
|
||||
* `HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU=1` and
|
||||
* `HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=1` to avoid Metal/MPS crashes in
|
||||
* daemon mode. Default: `true` on `darwin`, ignored elsewhere. Any value set
|
||||
* explicitly in {@link env} wins over the auto-applied value.
|
||||
*/
|
||||
platformCpuWorkaround?: boolean;
|
||||
/** Max time (ms) to wait for `/health` to return 200. Default: `30_000`. */
|
||||
readyTimeoutMs?: number;
|
||||
/** Polling interval (ms) while waiting for `/health`. Default: `1_000`. */
|
||||
readyPollIntervalMs?: number;
|
||||
/** Optional pluggable logger. Default: silent. */
|
||||
logger?: Logger;
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"target": "ES2022",
|
||||
"module": "ES2022",
|
||||
"lib": ["ES2022"],
|
||||
"moduleResolution": "node",
|
||||
"declaration": true,
|
||||
"outDir": "./dist",
|
||||
"rootDir": "./src",
|
||||
"strict": true,
|
||||
"esModuleInterop": true,
|
||||
"skipLibCheck": true,
|
||||
"forceConsistentCasingInFileNames": true,
|
||||
"resolveJsonModule": true
|
||||
},
|
||||
"include": ["src/**/*"],
|
||||
"exclude": ["node_modules", "dist", "src/**/*.test.ts"]
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
import { defineConfig } from 'tsup';
|
||||
|
||||
export default defineConfig({
|
||||
entry: ['src/index.ts'],
|
||||
format: ['esm'],
|
||||
dts: true,
|
||||
outDir: 'dist',
|
||||
clean: true,
|
||||
sourcemap: true,
|
||||
bundle: true,
|
||||
});
|
||||
@@ -0,0 +1,8 @@
|
||||
import { defineConfig } from 'vitest/config';
|
||||
|
||||
export default defineConfig({
|
||||
test: {
|
||||
include: ['src/**/*.test.ts'],
|
||||
environment: 'node',
|
||||
},
|
||||
});
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.4.19"
|
||||
version = "0.5.4"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
|
||||
@@ -34,7 +34,6 @@ Using context manager:
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from typing import Optional
|
||||
|
||||
@@ -72,8 +71,11 @@ class HindsightEmbedded:
|
||||
llm_model: Model name to use
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override (default: profile-specific pg0)
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (default: 300)
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (default: 0, disabled)
|
||||
log_level: Daemon log level (default: "info")
|
||||
ui: Whether to start the control plane web UI alongside the daemon (default: False)
|
||||
ui_port: Port for the UI. Defaults to daemon_port + 10000.
|
||||
ui_hostname: Hostname to bind the UI to. Defaults to "0.0.0.0".
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -84,8 +86,11 @@ class HindsightEmbedded:
|
||||
llm_model: str = "openai/gpt-oss-120b",
|
||||
llm_base_url: Optional[str] = None,
|
||||
database_url: Optional[str] = None,
|
||||
idle_timeout: int = 300,
|
||||
idle_timeout: int = 0,
|
||||
log_level: str = "info",
|
||||
ui: bool = False,
|
||||
ui_port: Optional[int] = None,
|
||||
ui_hostname: str = "0.0.0.0",
|
||||
):
|
||||
"""
|
||||
Initialize the embedded client (daemon starts on first use).
|
||||
@@ -97,8 +102,11 @@ class HindsightEmbedded:
|
||||
llm_model: Model name to use
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override
|
||||
idle_timeout: Seconds before daemon auto-exits when idle
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (0 = disabled)
|
||||
log_level: Daemon log level
|
||||
ui: Whether to start the control plane web UI alongside the daemon
|
||||
ui_port: Port for the UI (defaults to daemon_port + 10000)
|
||||
ui_hostname: Hostname to bind the UI to (defaults to "0.0.0.0")
|
||||
"""
|
||||
self.profile = profile
|
||||
|
||||
@@ -117,6 +125,10 @@ class HindsightEmbedded:
|
||||
if database_url:
|
||||
self.config["HINDSIGHT_EMBED_API_DATABASE_URL"] = database_url
|
||||
|
||||
self._ui = ui
|
||||
self._ui_port = ui_port
|
||||
self._ui_hostname = ui_hostname
|
||||
|
||||
self._client: Optional[Hindsight] = None
|
||||
self._lock = threading.Lock()
|
||||
self._started = False
|
||||
@@ -130,23 +142,50 @@ class HindsightEmbedded:
|
||||
self._memories_api: Optional[MemoriesAPI] = None
|
||||
|
||||
def _ensure_started(self):
|
||||
"""Ensure daemon is running (thread-safe)."""
|
||||
"""Ensure daemon is running (thread-safe), restarting if crashed."""
|
||||
if self._started and self._client is not None:
|
||||
return
|
||||
if self._manager.is_running(self.profile):
|
||||
return
|
||||
# Daemon crashed — reset state and fall through to restart
|
||||
logger.warning(
|
||||
"Daemon for profile '%s' is no longer responsive, restarting...",
|
||||
self.profile,
|
||||
)
|
||||
try:
|
||||
self._client.close()
|
||||
except Exception:
|
||||
logger.debug("Error closing stale client", exc_info=True)
|
||||
self._client = None
|
||||
self._started = False
|
||||
|
||||
with self._lock:
|
||||
# Double-check after acquiring lock
|
||||
if self._started and self._client is not None:
|
||||
return
|
||||
if self._manager.is_running(self.profile):
|
||||
return
|
||||
logger.warning(
|
||||
"Daemon for profile '%s' is no longer responsive (lock path), restarting...",
|
||||
self.profile,
|
||||
)
|
||||
try:
|
||||
self._client.close()
|
||||
except Exception:
|
||||
logger.debug("Error closing stale client", exc_info=True)
|
||||
self._client = None
|
||||
self._started = False
|
||||
|
||||
if self._closed:
|
||||
raise RuntimeError("Cannot use HindsightEmbedded after it has been closed")
|
||||
raise RuntimeError(
|
||||
"Cannot use HindsightEmbedded after it has been closed"
|
||||
)
|
||||
|
||||
# Use embed manager interface for daemon management
|
||||
logger.info(f"Ensuring daemon is running for profile '{self.profile}'...")
|
||||
success = self._manager.ensure_running(self.config, self.profile)
|
||||
if not success:
|
||||
raise RuntimeError(f"Failed to start daemon for profile '{self.profile}'")
|
||||
raise RuntimeError(
|
||||
f"Failed to start daemon for profile '{self.profile}'"
|
||||
)
|
||||
|
||||
# Get daemon URL and create client
|
||||
daemon_url = self._manager.get_url(self.profile)
|
||||
@@ -154,6 +193,15 @@ class HindsightEmbedded:
|
||||
self._started = True
|
||||
logger.info(f"Connected to daemon at {daemon_url}")
|
||||
|
||||
# Start UI if requested
|
||||
if self._ui:
|
||||
logger.info(f"Starting UI for profile '{self.profile}'...")
|
||||
ui_started = self._manager.start_ui(
|
||||
self.profile, self._ui_port, self._ui_hostname
|
||||
)
|
||||
if not ui_started:
|
||||
logger.warning(f"Failed to start UI for profile '{self.profile}'")
|
||||
|
||||
def _cleanup(self, stop_daemon_on_close: bool = False):
|
||||
"""
|
||||
Cleanup client resources (idempotent).
|
||||
@@ -165,20 +213,47 @@ class HindsightEmbedded:
|
||||
if self._closed:
|
||||
return
|
||||
|
||||
with self._lock:
|
||||
acquired = self._lock.acquire(timeout=5.0)
|
||||
if not acquired:
|
||||
# Lock is held by another thread (e.g. _ensure_started).
|
||||
# Mark closed to prevent new operations but skip shared-state
|
||||
# teardown — the daemon's idle timeout handles the rest.
|
||||
logger.warning(
|
||||
"Cleanup lock acquisition timed out for profile '%s'; "
|
||||
"marking closed, daemon will idle-stop on its own",
|
||||
self.profile,
|
||||
)
|
||||
self._closed = True
|
||||
return
|
||||
|
||||
try:
|
||||
if self._closed:
|
||||
return
|
||||
|
||||
if self._client is not None:
|
||||
self._client.close()
|
||||
try:
|
||||
self._client.close()
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"Error closing client for profile '%s'",
|
||||
self.profile,
|
||||
exc_info=True,
|
||||
)
|
||||
self._client = None
|
||||
|
||||
# Stop UI if it was started
|
||||
if self._ui and self._started:
|
||||
logger.info(f"Stopping UI for profile '{self.profile}'...")
|
||||
self._manager.stop_ui(self.profile, self._ui_port)
|
||||
|
||||
# Optionally stop daemon (daemon has idle timeout, so not required)
|
||||
if stop_daemon_on_close and self._started:
|
||||
logger.info(f"Stopping daemon for profile '{self.profile}'...")
|
||||
self._manager.stop(self.profile)
|
||||
|
||||
self._closed = True
|
||||
finally:
|
||||
self._lock.release()
|
||||
|
||||
def close(self, stop_daemon: bool = False):
|
||||
"""
|
||||
@@ -201,23 +276,10 @@ class HindsightEmbedded:
|
||||
This allows HindsightEmbedded to expose all HindsightClient methods
|
||||
without manually wrapping each one.
|
||||
"""
|
||||
# Ensure server is started before proxying
|
||||
# Ensure server is started (and restart if crashed) before proxying
|
||||
self._ensure_started()
|
||||
|
||||
# Get the attribute from the underlying client
|
||||
attr = getattr(self._client, name)
|
||||
|
||||
# If it's a callable, wrap it to ensure server is started
|
||||
# (shouldn't be needed since _ensure_started already called, but defensive)
|
||||
if callable(attr):
|
||||
|
||||
def wrapper(*args, **kwargs):
|
||||
self._ensure_started()
|
||||
return attr(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return attr
|
||||
return getattr(self._client, name)
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry - ensures server is started."""
|
||||
@@ -342,11 +404,8 @@ class HindsightEmbedded:
|
||||
"""
|
||||
Get the underlying Hindsight client for direct access.
|
||||
|
||||
WARNING: Using this property directly means daemon restarts won't be
|
||||
handled automatically. Prefer using the API namespaces (banks, mental_models,
|
||||
directives, memories) or direct method calls on HindsightEmbedded instead.
|
||||
|
||||
Ensures daemon is started before returning the client.
|
||||
Ensures daemon is started (and restarts it if it has crashed) before
|
||||
returning the client.
|
||||
|
||||
Returns:
|
||||
Hindsight: The underlying client instance
|
||||
@@ -357,9 +416,8 @@ class HindsightEmbedded:
|
||||
|
||||
embedded = HindsightEmbedded(profile="myapp", ...)
|
||||
|
||||
# Direct access (not recommended - daemon crashes won't be handled)
|
||||
client = embedded.client
|
||||
banks = client.list_banks() # If daemon crashes, this will fail
|
||||
banks = client.list_banks()
|
||||
```
|
||||
"""
|
||||
self._ensure_started()
|
||||
@@ -373,5 +431,15 @@ class HindsightEmbedded:
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Check if the client is initialized."""
|
||||
return self._started and not self._closed and self._client is not None
|
||||
"""Check if the client is initialized and the daemon is responsive."""
|
||||
return (
|
||||
self._started
|
||||
and not self._closed
|
||||
and self._client is not None
|
||||
and self._manager.is_running(self.profile)
|
||||
)
|
||||
|
||||
@property
|
||||
def ui_url(self) -> str:
|
||||
"""Get the UI URL for this profile."""
|
||||
return self._manager.get_ui_url(self.profile)
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.4.19"
|
||||
version = "0.5.4"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -20,6 +20,9 @@ hindsight-client = { workspace = true }
|
||||
hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]>=0.4.17",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
"pytest-asyncio>=0.21.0",
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
"""
|
||||
Unit test for _cleanup lock timeout behavior.
|
||||
|
||||
Verifies that _cleanup completes even when the lock is held by another thread,
|
||||
instead of hanging indefinitely (fixes #952).
|
||||
"""
|
||||
|
||||
import threading
|
||||
import time
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_cleanup_completes_when_lock_held():
|
||||
"""
|
||||
_cleanup should complete (best-effort) even when self._lock is held
|
||||
by another thread, e.g. during a long _ensure_started call.
|
||||
"""
|
||||
with patch.dict("sys.modules", {
|
||||
"hindsight_client": MagicMock(),
|
||||
"hindsight_embed": MagicMock(),
|
||||
"hindsight.api_namespaces": MagicMock(),
|
||||
}):
|
||||
from hindsight.embedded import HindsightEmbedded
|
||||
|
||||
client = HindsightEmbedded.__new__(HindsightEmbedded)
|
||||
client.profile = "test"
|
||||
client._lock = threading.Lock()
|
||||
client._closed = False
|
||||
client._client = None
|
||||
client._started = False
|
||||
client._ui = False
|
||||
|
||||
# Simulate another thread holding the lock
|
||||
client._lock.acquire()
|
||||
|
||||
cleanup_done = threading.Event()
|
||||
|
||||
def run_cleanup():
|
||||
client._cleanup()
|
||||
cleanup_done.set()
|
||||
|
||||
t = threading.Thread(target=run_cleanup)
|
||||
t.start()
|
||||
|
||||
# Cleanup should complete within the timeout (5s) + margin
|
||||
assert cleanup_done.wait(timeout=8.0), (
|
||||
"_cleanup hung instead of timing out on lock acquisition"
|
||||
)
|
||||
|
||||
# Release the lock from the simulating thread
|
||||
client._lock.release()
|
||||
t.join(timeout=1.0)
|
||||
|
||||
assert client._closed, "Client should be marked as closed after cleanup"
|
||||
@@ -15,6 +15,8 @@ import os
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
import urllib.request
|
||||
import json
|
||||
|
||||
from hindsight import HindsightEmbedded
|
||||
|
||||
@@ -23,12 +25,20 @@ from hindsight import HindsightEmbedded
|
||||
def llm_config():
|
||||
"""Get LLM configuration from environment (session-scoped)."""
|
||||
# Try both naming conventions
|
||||
provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER") or os.getenv("HINDSIGHT_LLM_PROVIDER", "groq")
|
||||
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY") or os.getenv("HINDSIGHT_LLM_API_KEY", "")
|
||||
model = os.getenv("HINDSIGHT_API_LLM_MODEL") or os.getenv("HINDSIGHT_LLM_MODEL", "openai/gpt-oss-120b")
|
||||
provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER") or os.getenv(
|
||||
"HINDSIGHT_LLM_PROVIDER", "groq"
|
||||
)
|
||||
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY") or os.getenv(
|
||||
"HINDSIGHT_LLM_API_KEY", ""
|
||||
)
|
||||
model = os.getenv("HINDSIGHT_API_LLM_MODEL") or os.getenv(
|
||||
"HINDSIGHT_LLM_MODEL", "openai/gpt-oss-120b"
|
||||
)
|
||||
|
||||
if not api_key:
|
||||
pytest.skip("LLM API key not configured. Set HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_LLM_API_KEY.")
|
||||
pytest.skip(
|
||||
"LLM API key not configured. Set HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_LLM_API_KEY."
|
||||
)
|
||||
|
||||
return {
|
||||
"llm_provider": provider,
|
||||
@@ -78,7 +88,9 @@ def test_embedded_context_manager(llm_config):
|
||||
|
||||
# Recall memory
|
||||
recall_results = client.recall(bank_id=bank_id, query="context")
|
||||
assert isinstance(recall_results.results, list), "Recall should return results list"
|
||||
assert isinstance(recall_results.results, list), (
|
||||
"Recall should return results list"
|
||||
)
|
||||
|
||||
# Server should be stopped after context exit
|
||||
# Note: We can't check client.is_running here as client is out of scope
|
||||
@@ -105,7 +117,9 @@ def test_embedded_complete_workflow(llm_config):
|
||||
# Step 1: Create a memory bank
|
||||
print(f"\n1. Creating memory bank: {bank_id}")
|
||||
bank_response = client.create_bank(
|
||||
bank_id=bank_id, name="Test Assistant", mission="Help with programming tasks"
|
||||
bank_id=bank_id,
|
||||
name="Test Assistant",
|
||||
mission="Help with programming tasks",
|
||||
)
|
||||
assert bank_response.bank_id == bank_id
|
||||
|
||||
@@ -126,7 +140,9 @@ def test_embedded_complete_workflow(llm_config):
|
||||
items=[
|
||||
{"content": "User works with pandas and numpy."},
|
||||
{"content": "User likes matplotlib for visualization."},
|
||||
{"content": "User is interested in machine learning with scikit-learn."},
|
||||
{
|
||||
"content": "User is interested in machine learning with scikit-learn."
|
||||
},
|
||||
],
|
||||
)
|
||||
assert batch_response.success
|
||||
@@ -134,7 +150,9 @@ def test_embedded_complete_workflow(llm_config):
|
||||
|
||||
# Step 4: Recall memories
|
||||
print("\n4. Recalling memories...")
|
||||
recall_response = client.recall(bank_id=bank_id, query="What tools does the user prefer?", max_tokens=2000)
|
||||
recall_response = client.recall(
|
||||
bank_id=bank_id, query="What tools does the user prefer?", max_tokens=2000
|
||||
)
|
||||
assert isinstance(recall_response.results, list)
|
||||
assert len(recall_response.results) > 0
|
||||
print(f" Found {len(recall_response.results)} relevant memories")
|
||||
@@ -152,7 +170,9 @@ def test_embedded_complete_workflow(llm_config):
|
||||
|
||||
# Verify answer mentions relevant tools
|
||||
answer_lower = reflect_response.text.lower()
|
||||
assert any(term in answer_lower for term in ["python", "pandas", "numpy", "data"])
|
||||
assert any(
|
||||
term in answer_lower for term in ["python", "pandas", "numpy", "data"]
|
||||
)
|
||||
|
||||
# Step 6: List memories
|
||||
print("\n6. Listing memories...")
|
||||
@@ -215,7 +235,9 @@ def test_embedded_method_proxying(llm_config):
|
||||
assert bank.bank_id == bank_id
|
||||
|
||||
# Test mission setting
|
||||
mission_response = client.set_mission(bank_id=bank_id, mission="Test mission for proxying")
|
||||
mission_response = client.set_mission(
|
||||
bank_id=bank_id, mission="Test mission for proxying"
|
||||
)
|
||||
assert mission_response.bank_id == bank_id
|
||||
|
||||
# Test retain
|
||||
@@ -264,7 +286,9 @@ def test_embedded_multiple_banks(llm_config):
|
||||
|
||||
# Create second bank and store data
|
||||
client.create_bank(bank_id=bank2_id, name="Bank 2")
|
||||
client.retain(bank_id=bank2_id, content="Bob uses JavaScript for web development")
|
||||
client.retain(
|
||||
bank_id=bank2_id, content="Bob uses JavaScript for web development"
|
||||
)
|
||||
|
||||
# Recall from both banks
|
||||
results1 = client.recall(bank_id=bank1_id, query="programming language")
|
||||
@@ -275,9 +299,9 @@ def test_embedded_multiple_banks(llm_config):
|
||||
|
||||
# Verify banks are isolated (each should only see their own content)
|
||||
# This is a basic check - content isolation is tested more thoroughly in other tests
|
||||
assert results1.results[0].text != results2.results[0].text or len(results1.results) != len(
|
||||
results2.results
|
||||
)
|
||||
assert results1.results[0].text != results2.results[0].text or len(
|
||||
results1.results
|
||||
) != len(results2.results)
|
||||
|
||||
finally:
|
||||
client.close()
|
||||
@@ -296,10 +320,14 @@ def test_embedded_profile_isolation(llm_config):
|
||||
|
||||
try:
|
||||
# Store data in profile1
|
||||
client1.retain(bank_id=bank_id, content="User likes TypeScript for frontend development")
|
||||
client1.retain(
|
||||
bank_id=bank_id, content="User likes TypeScript for frontend development"
|
||||
)
|
||||
|
||||
# Store different data in profile2
|
||||
client2.retain(bank_id=bank_id, content="User prefers Rust for systems programming")
|
||||
client2.retain(
|
||||
bank_id=bank_id, content="User prefers Rust for systems programming"
|
||||
)
|
||||
|
||||
# Each profile should only see its own data
|
||||
results1 = client1.recall(bank_id=bank_id, query="programming preference")
|
||||
@@ -334,5 +362,81 @@ def test_embedded_error_after_close(llm_config):
|
||||
assert not client.is_running
|
||||
|
||||
# Trying to use it after close should raise an error
|
||||
with pytest.raises(RuntimeError, match="Cannot use HindsightEmbedded after it has been closed"):
|
||||
with pytest.raises(
|
||||
RuntimeError, match="Cannot use HindsightEmbedded after it has been closed"
|
||||
):
|
||||
client.retain(bank_id=bank_id, content="This should fail")
|
||||
|
||||
|
||||
def test_embedded_ui_flag(llm_config):
|
||||
"""
|
||||
Test that ui=True starts the control plane UI alongside the daemon,
|
||||
and that the UI's health endpoint reports a connected dataplane.
|
||||
"""
|
||||
profile = f"test_ui_{uuid.uuid4().hex[:8]}"
|
||||
bank_id = f"bank_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
client = HindsightEmbedded(profile=profile, log_level="info", ui=True, **llm_config)
|
||||
|
||||
try:
|
||||
# First use triggers daemon + UI startup
|
||||
result = client.retain(bank_id=bank_id, content="UI integration test content")
|
||||
assert result.success, "Retain should succeed"
|
||||
assert client.is_running, "Daemon should be running"
|
||||
|
||||
# Verify UI is reachable and reports connected dataplane
|
||||
ui_url = client.ui_url
|
||||
assert ui_url, "ui_url should be set"
|
||||
|
||||
health_url = f"{ui_url}/api/health"
|
||||
with urllib.request.urlopen(health_url, timeout=10) as resp:
|
||||
health = json.loads(resp.read().decode())
|
||||
|
||||
assert health["status"] == "ok", (
|
||||
f"UI health status should be 'ok', got: {health['status']}"
|
||||
)
|
||||
assert health["dataplane"]["status"] == "connected", (
|
||||
f"Dataplane should be connected, got: {health['dataplane']}"
|
||||
)
|
||||
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def test_embedded_daemon_crash_recovery(llm_config):
|
||||
"""
|
||||
Test that HindsightEmbedded recovers when the daemon crashes.
|
||||
|
||||
Simulates a crash by stopping the daemon, then verifies
|
||||
that the next operation transparently restarts it.
|
||||
"""
|
||||
profile = f"test_crash_{uuid.uuid4().hex[:8]}"
|
||||
bank_id = f"bank_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
client = HindsightEmbedded(profile=profile, log_level="info", **llm_config)
|
||||
|
||||
try:
|
||||
# Start daemon and store a memory
|
||||
result = client.retain(bank_id=bank_id, content="Before crash")
|
||||
assert result.success, "Initial retain should succeed"
|
||||
assert client.is_running, "Daemon should be running"
|
||||
|
||||
original_url = client.url
|
||||
|
||||
# Simulate daemon crash by stopping it
|
||||
client._manager.stop(client.profile)
|
||||
assert not client._manager.is_running(client.profile), (
|
||||
"Daemon should be stopped after simulated crash"
|
||||
)
|
||||
|
||||
# Next operation should transparently restart the daemon
|
||||
result2 = client.retain(bank_id=bank_id, content="After crash recovery")
|
||||
assert result2.success, "Retain after crash recovery should succeed"
|
||||
assert client.is_running, "Daemon should be running again after recovery"
|
||||
|
||||
# Verify recall still works
|
||||
recall_result = client.recall(bank_id=bank_id, query="crash")
|
||||
assert isinstance(recall_result.results, list), "Recall should return results"
|
||||
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
@@ -46,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.4.19"
|
||||
__version__ = "0.5.4"
|
||||
|
||||
@@ -249,7 +249,7 @@ async def _run_migration(
|
||||
schemas = list(dict.fromkeys(schemas))
|
||||
|
||||
for schema in schemas:
|
||||
run_migrations(resolved_url, schema=schema)
|
||||
run_migrations(resolved_url, schema=schema, migration_database_url=config.migration_database_url)
|
||||
|
||||
if embedding_dimension is not None:
|
||||
for schema in schemas:
|
||||
@@ -375,6 +375,140 @@ def decommission_worker(
|
||||
typer.echo(f"No tasks found for worker '{worker_id}'")
|
||||
|
||||
|
||||
async def _decommission_all_workers(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
|
||||
"""Release all processing tasks from all workers, setting them back to pending status."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
|
||||
conn = await asyncpg.connect(resolved_url)
|
||||
try:
|
||||
table = _fq_table("async_operations", schema)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
|
||||
WHERE status = 'processing'
|
||||
RETURNING operation_id, worker_id, operation_type
|
||||
""",
|
||||
)
|
||||
return [dict(r) for r in rows]
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
|
||||
@app.command(name="decommission-workers")
|
||||
def decommission_workers(
|
||||
schema: str = typer.Option("public", "--schema", "-s", help="Database schema"),
|
||||
yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt"),
|
||||
):
|
||||
"""Release all processing tasks from all workers (sets status back to pending).
|
||||
|
||||
Use this command to recover from situations where one or more workers have crashed
|
||||
or been removed without graceful shutdown. All tasks currently in 'processing' status
|
||||
will be released back to the queue regardless of which worker owns them.
|
||||
"""
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
if not config.database_url:
|
||||
typer.echo("Error: Database URL not configured.", err=True)
|
||||
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
if not yes:
|
||||
typer.confirm(
|
||||
"This will release ALL processing tasks from ALL workers back to pending. Continue?",
|
||||
abort=True,
|
||||
)
|
||||
|
||||
typer.echo(f"Decommissioning all workers (schema: {schema})...")
|
||||
|
||||
released = asyncio.run(_decommission_all_workers(config.database_url, schema))
|
||||
|
||||
if released:
|
||||
# Group by worker_id for summary
|
||||
by_worker: dict[str, int] = {}
|
||||
for row in released:
|
||||
wid = row["worker_id"] or "unknown"
|
||||
by_worker[wid] = by_worker.get(wid, 0) + 1
|
||||
|
||||
typer.echo(f"Released {len(released)} task(s):")
|
||||
for wid, count in by_worker.items():
|
||||
typer.echo(f" {wid}: {count} task(s)")
|
||||
else:
|
||||
typer.echo("No processing tasks found")
|
||||
|
||||
|
||||
async def _worker_status(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
|
||||
"""Get all processing tasks grouped by worker with their last update time."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
|
||||
conn = await asyncpg.connect(resolved_url)
|
||||
try:
|
||||
table = _fq_table("async_operations", schema)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT worker_id, operation_id, operation_type, bank_id,
|
||||
claimed_at, updated_at,
|
||||
now() - claimed_at AS running_for,
|
||||
now() - updated_at AS last_update_ago
|
||||
FROM {table}
|
||||
WHERE status = 'processing'
|
||||
ORDER BY worker_id, claimed_at
|
||||
""",
|
||||
)
|
||||
return [dict(r) for r in rows]
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
|
||||
@app.command(name="worker-status")
|
||||
def worker_status(
|
||||
schema: str = typer.Option("public", "--schema", "-s", help="Database schema"),
|
||||
):
|
||||
"""Show all currently processing tasks grouped by worker.
|
||||
|
||||
Displays each worker's active tasks with operation type, bank, how long
|
||||
the task has been running, and when it was last updated. Useful for
|
||||
identifying dead workers with orphaned tasks.
|
||||
"""
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
if not config.database_url:
|
||||
typer.echo("Error: Database URL not configured.", err=True)
|
||||
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
rows = asyncio.run(_worker_status(config.database_url, schema))
|
||||
|
||||
if not rows:
|
||||
typer.echo("No processing tasks found")
|
||||
return
|
||||
|
||||
# Group by worker_id
|
||||
by_worker: dict[str, list[dict[str, Any]]] = {}
|
||||
for row in rows:
|
||||
wid = row["worker_id"] or "unknown"
|
||||
by_worker.setdefault(wid, []).append(row)
|
||||
|
||||
typer.echo(f"Processing tasks across {len(by_worker)} worker(s):\n")
|
||||
for wid, tasks in by_worker.items():
|
||||
typer.echo(f"Worker: {wid} ({len(tasks)} task(s))")
|
||||
for task in tasks:
|
||||
op_id = str(task["operation_id"])[:8]
|
||||
running_for = task["running_for"]
|
||||
last_update = task["last_update_ago"]
|
||||
typer.echo(
|
||||
f" {op_id} {task['operation_type']:<20s} bank={task['bank_id']}"
|
||||
f" running={running_for} last_update={last_update} ago"
|
||||
)
|
||||
typer.echo("")
|
||||
|
||||
|
||||
def main():
|
||||
app()
|
||||
|
||||
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
"""Recreate entities trigram index on LOWER(canonical_name) for case-insensitive matching
|
||||
|
||||
The previous GIN trigram index on canonical_name was case-sensitive, causing
|
||||
"Alice" and "alice" to have different trigram sets. This recreates it on
|
||||
LOWER(canonical_name) so the % operator matches case-insensitively.
|
||||
|
||||
Revision ID: 2eee35aa3cfc
|
||||
Revises: d6e7f8a9b0c1
|
||||
Create Date: 2026-03-31
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "2eee35aa3cfc"
|
||||
down_revision: str | Sequence[str] | None = "d6e7f8a9b0c1"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# Drop the old case-sensitive trigram index
|
||||
op.execute("DROP INDEX IF EXISTS entities_canonical_name_trgm_idx")
|
||||
# Create case-insensitive trigram index on LOWER(canonical_name)
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS entities_canonical_name_lower_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (LOWER(canonical_name) gin_trgm_ops)"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP INDEX IF EXISTS entities_canonical_name_lower_trgm_idx")
|
||||
schema = _get_schema_prefix()
|
||||
# Restore original case-sensitive index
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
"""Merge divergent migration heads for v0.5.3
|
||||
|
||||
v0.5.3 shipped with two migration heads that were never unified:
|
||||
|
||||
* ``c4x5y6z7a8b9`` — delta-refresh chain
|
||||
(``add_last_refreshed_source_query`` ->
|
||||
``add_structured_content_to_mental_models`` ->
|
||||
``backsweep_orphan_observations_v2``)
|
||||
|
||||
* ``h3i4j5k6l7m8`` — per-bank vector indexes / audit log chain
|
||||
(the ``merge_heads_and_add_unit_entities_index`` subtree)
|
||||
|
||||
Both fork from ``z1u2v3w4x5y6``. Upgrades from v0.5.2 still succeed — the
|
||||
walker applies the three c4x5 revisions and leaves the database stamped at
|
||||
both heads — but the result is a split DAG: ``alembic upgrade head``
|
||||
(singular) is ambiguous, and any future migration has to pick one head as
|
||||
its parent, orphaning the other.
|
||||
|
||||
This revision linearises the DAG into a single head. It has no schema
|
||||
effect.
|
||||
|
||||
Revision ID: 8c6fa6f7230b
|
||||
Revises: c4x5y6z7a8b9, h3i4j5k6l7m8
|
||||
Create Date: 2026-04-18
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
revision: str = "8c6fa6f7230b"
|
||||
down_revision: str | Sequence[str] | None = ("c4x5y6z7a8b9", "h3i4j5k6l7m8")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
pass
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
"""Add last_refreshed_source_query column to mental_models
|
||||
|
||||
Revision ID: a2v3w4x5y6z7
|
||||
Revises: z1u2v3w4x5y6
|
||||
Create Date: 2026-04-15
|
||||
|
||||
Tracks the source_query that was used during the most recent refresh.
|
||||
Used by delta-mode refresh to detect when the query has changed: if it has,
|
||||
delta mode falls back to a full regeneration because the surgical-edit
|
||||
assumption (same topic, new facts) no longer holds.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a2v3w4x5y6z7"
|
||||
down_revision: str | Sequence[str] | None = "z1u2v3w4x5y6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS last_refreshed_source_query TEXT
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS last_refreshed_source_query")
|
||||
+142
@@ -0,0 +1,142 @@
|
||||
"""Fix per-bank vector indexes to match configured extension
|
||||
|
||||
Revision ID: a4b5c6d7e8f9
|
||||
Revises: 2eee35aa3cfc
|
||||
Create Date: 2026-04-01
|
||||
|
||||
Migration d5e6f7a8b9c0 hardcoded HNSW when creating per-bank partial vector
|
||||
indexes, ignoring HINDSIGHT_API_VECTOR_EXTENSION. Banks that existed when that
|
||||
migration ran got HNSW indexes even when pgvectorscale (DiskANN) or vchord
|
||||
was configured.
|
||||
|
||||
This migration detects the mismatch and recreates the affected indexes with
|
||||
the correct type. Skipped entirely when the configured extension is pgvector
|
||||
(the default), since those indexes are already correct.
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
revision: str = "a4b5c6d7e8f9"
|
||||
down_revision: str | Sequence[str] | None = "2eee35aa3cfc"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_FACT_TYPES: dict[str, str] = {
|
||||
"world": "worl",
|
||||
"experience": "expr",
|
||||
"observation": "obsv",
|
||||
}
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _target_index_type() -> str | None:
|
||||
"""Return the target index type, or None if pgvector (no fix needed)."""
|
||||
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
|
||||
if ext == "pgvectorscale":
|
||||
return "diskann"
|
||||
elif ext == "vchord":
|
||||
return "vchordrq"
|
||||
return None
|
||||
|
||||
|
||||
def _vector_index_using_clause() -> str:
|
||||
"""Return the USING clause based on the configured vector extension."""
|
||||
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
elif ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
else:
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
target = _target_index_type()
|
||||
if target is None:
|
||||
# pgvector — indexes are already HNSW, nothing to fix
|
||||
return
|
||||
|
||||
bind = op.get_bind()
|
||||
schema_name = context.config.get_main_option("target_schema")
|
||||
schema = _get_schema_prefix()
|
||||
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
|
||||
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
|
||||
using_clause = _vector_index_using_clause()
|
||||
pg_schema = schema_name or "public"
|
||||
|
||||
rows = bind.execute(text(f"SELECT bank_id, internal_id FROM {banks_ref}")).fetchall() # noqa: S608
|
||||
for row in rows:
|
||||
bank_id = row[0]
|
||||
internal_id = str(row[1]).replace("-", "")[:16]
|
||||
escaped_bank_id = bank_id.replace("'", "''")
|
||||
for ft, ft_short in _FACT_TYPES.items():
|
||||
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
|
||||
|
||||
# Check if this index exists and what type it is
|
||||
idx_info = bind.execute(
|
||||
text("SELECT indexdef FROM pg_indexes WHERE schemaname = :schema AND indexname = :idx"),
|
||||
{"schema": pg_schema, "idx": idx_name},
|
||||
).fetchone()
|
||||
|
||||
if idx_info is None:
|
||||
# Index doesn't exist — create it with the correct type
|
||||
bind.execute(
|
||||
text(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx_name} "
|
||||
f"ON {table_ref} {using_clause} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
indexdef = idx_info[0].lower()
|
||||
if target in indexdef:
|
||||
# Already the correct type
|
||||
continue
|
||||
|
||||
# Wrong type — drop and recreate
|
||||
bind.execute(text(f"DROP INDEX IF EXISTS {schema}{idx_name}"))
|
||||
bind.execute(
|
||||
text(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx_name} "
|
||||
f"ON {table_ref} {using_clause} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# Downgrade recreates indexes as HNSW (the original hardcoded behavior)
|
||||
target = _target_index_type()
|
||||
if target is None:
|
||||
return
|
||||
|
||||
bind = op.get_bind()
|
||||
schema_name = context.config.get_main_option("target_schema")
|
||||
schema = _get_schema_prefix()
|
||||
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
|
||||
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
|
||||
|
||||
rows = bind.execute(text(f"SELECT bank_id, internal_id FROM {banks_ref}")).fetchall() # noqa: S608
|
||||
for row in rows:
|
||||
bank_id = row[0]
|
||||
internal_id = str(row[1]).replace("-", "")[:16]
|
||||
escaped_bank_id = bank_id.replace("'", "''")
|
||||
for ft, ft_short in _FACT_TYPES.items():
|
||||
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
|
||||
bind.execute(text(f"DROP INDEX IF EXISTS {schema}{idx_name}"))
|
||||
bind.execute(
|
||||
text(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx_name} "
|
||||
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
|
||||
)
|
||||
)
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
"""add content_hash to chunks table for delta retain
|
||||
|
||||
Revision ID: b3c4d5e6f7a8
|
||||
Revises: a3b4c5d6e7f8
|
||||
Create Date: 2026-03-25
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "b3c4d5e6f7a8"
|
||||
down_revision: str | Sequence[str] | None = "a3b4c5d6e7f8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# Add content_hash column to chunks table for delta comparison
|
||||
op.execute(f"ALTER TABLE {schema}chunks ADD COLUMN IF NOT EXISTS content_hash TEXT")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}chunks DROP COLUMN IF EXISTS content_hash")
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
"""Add structured_content JSONB column to mental_models
|
||||
|
||||
Revision ID: b3w4x5y6z7a8
|
||||
Revises: a2v3w4x5y6z7
|
||||
Create Date: 2026-04-16
|
||||
|
||||
Stores the structured representation of a mental model document (sections,
|
||||
blocks). The plain ``content`` column remains the rendered markdown shown to
|
||||
users. ``structured_content`` is the source of truth for delta-mode refreshes:
|
||||
each refresh applies a list of typed operations to the structured doc, then
|
||||
re-renders to markdown — so unchanged sections come through byte-identical
|
||||
without an LLM round-trip.
|
||||
|
||||
Nullable: existing markdown-only mental models continue to work in full mode;
|
||||
the column is populated lazily the first time a model is refreshed in delta
|
||||
mode.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "b3w4x5y6z7a8"
|
||||
down_revision: str | Sequence[str] | None = "a2v3w4x5y6z7"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS structured_content JSONB
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS structured_content")
|
||||
+15
-2
@@ -11,6 +11,7 @@ block; see migrations.py for how this is handled safely.
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "c1a2b3d4e5f6"
|
||||
@@ -25,9 +26,21 @@ def _get_schema_prefix() -> str:
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# pg_trgm ships with every standard PostgreSQL installation as a contrib module.
|
||||
# pg_trgm ships with most PostgreSQL installations as a contrib module.
|
||||
# It enables fast similarity lookups via GIN indexes, used for entity name matching.
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm")
|
||||
# On managed services (e.g. Azure Flexible Server), the extension may not be
|
||||
# available or may require manual enablement. We gracefully skip the index
|
||||
# creation if the extension cannot be loaded — the entity resolver will
|
||||
# auto-detect and fall back to the "full" lookup strategy at runtime. See #626.
|
||||
conn = op.get_bind()
|
||||
try:
|
||||
conn.execute(sa.text("CREATE EXTENSION IF NOT EXISTS pg_trgm"))
|
||||
except Exception:
|
||||
# Extension not available (managed Postgres, insufficient privileges, etc.)
|
||||
# Roll back the failed statement and skip index creation.
|
||||
conn.execute(sa.text("ROLLBACK"))
|
||||
conn.execute(sa.text("BEGIN"))
|
||||
return
|
||||
|
||||
schema = _get_schema_prefix()
|
||||
# GIN index on canonical_name enables sub-millisecond trigram similarity queries
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Add audit_log table for feature usage tracking.
|
||||
|
||||
Merge migration that combines the two existing heads (a3b4c5d6e7f8 + c8e5f2a3b4d1).
|
||||
|
||||
Stores raw request/response as JSONB for expandability without future migrations.
|
||||
The metadata JSONB column allows adding arbitrary fields in the future.
|
||||
|
||||
Revision ID: c2d3e4f5g6h7
|
||||
Revises: a3b4c5d6e7f8, c8e5f2a3b4d1
|
||||
Create Date: 2026-03-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "c2d3e4f5g6h7"
|
||||
down_revision: str | Sequence[str] | None = ("a3b4c5d6e7f8", "c8e5f2a3b4d1")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}audit_log (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
action TEXT NOT NULL,
|
||||
transport TEXT NOT NULL,
|
||||
bank_id TEXT,
|
||||
started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
ended_at TIMESTAMPTZ,
|
||||
request JSONB,
|
||||
response JSONB,
|
||||
metadata JSONB DEFAULT '{{}}'::jsonb
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_audit_log_action_started ON {schema}audit_log (action, started_at DESC)"
|
||||
)
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_audit_log_bank_started ON {schema}audit_log (bank_id, started_at DESC)")
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_audit_log_started ON {schema}audit_log (started_at DESC)")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_started")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_bank_started")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_action_started")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}audit_log")
|
||||
+66
@@ -0,0 +1,66 @@
|
||||
"""backsweep_orphan_observations_v2
|
||||
|
||||
Re-run of Pass 2 from migration ``g7h8i9j0k1l2_backsweep_orphan_observations``
|
||||
to sweep observations that became orphaned between then and now.
|
||||
|
||||
Why we need it again:
|
||||
``fact_storage.handle_document_tracking`` (the retain/upsert path) deleted
|
||||
the existing document via the FK cascade — which removes the source
|
||||
``memory_units`` — but never invalidated the observations derived from
|
||||
them. Only the explicit ``MemoryEngine.delete_document`` API called
|
||||
``_delete_stale_observations_for_memories``. Every document re-ingest
|
||||
therefore left orphan observations whose ``source_memory_ids`` arrays
|
||||
pointed at IDs that no longer existed in ``memory_units``.
|
||||
|
||||
``handle_document_tracking`` now calls the same cleanup helper before the
|
||||
cascade, so no new orphans will accumulate going forward. This migration
|
||||
cleans up the historical residue.
|
||||
|
||||
Identical to Pass 2 of g7h8i9j0k1l2. Pass 1 (memory_units whose bank is
|
||||
gone) is intentionally not re-run; that scenario has no fresh source.
|
||||
|
||||
Revision ID: c4x5y6z7a8b9
|
||||
Revises: b3w4x5y6z7a8
|
||||
Create Date: 2026-04-16
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "c4x5y6z7a8b9"
|
||||
down_revision: str | Sequence[str] | None = "b3w4x5y6z7a8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
mu = f"{schema}memory_units"
|
||||
|
||||
# Delete observations whose every source_memory_id refers to a now-deleted
|
||||
# memory_unit (or the array is empty). Observations with at least one
|
||||
# surviving source are left alone — the consolidation engine will refresh
|
||||
# their text on the next pass.
|
||||
op.execute(
|
||||
f"""
|
||||
DELETE FROM {mu} orphan
|
||||
WHERE orphan.fact_type = 'observation'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {mu} src
|
||||
WHERE src.id = ANY(orphan.source_memory_ids)
|
||||
AND src.bank_id = orphan.bank_id
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# Deleted rows cannot be restored.
|
||||
pass
|
||||
+48
@@ -0,0 +1,48 @@
|
||||
"""Add bank_id column to memory_links for direct filtering
|
||||
|
||||
The stats endpoint JOINs memory_links to memory_units just to filter by
|
||||
bank_id. With millions of links this takes 18+ seconds. Adding bank_id
|
||||
directly to memory_links lets Postgres push the filter down before the JOIN.
|
||||
|
||||
Revision ID: c5d6e7f8a9b0
|
||||
Revises: b3c4d5e6f7a8
|
||||
Create Date: 2026-03-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "c5d6e7f8a9b0"
|
||||
down_revision: str | Sequence[str] | None = "b3c4d5e6f7a8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# 1. Add nullable column
|
||||
op.execute(f"ALTER TABLE {schema}memory_links ADD COLUMN IF NOT EXISTS bank_id TEXT")
|
||||
|
||||
# 2. Backfill from memory_units
|
||||
op.execute(f"""
|
||||
UPDATE {schema}memory_links ml
|
||||
SET bank_id = mu.bank_id
|
||||
FROM {schema}memory_units mu
|
||||
WHERE ml.from_unit_id = mu.id
|
||||
AND ml.bank_id IS NULL
|
||||
""")
|
||||
|
||||
# 3. Set NOT NULL
|
||||
op.execute(f"ALTER TABLE {schema}memory_links ALTER COLUMN bank_id SET NOT NULL")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}memory_links DROP COLUMN IF EXISTS bank_id")
|
||||
+27
-19
@@ -1,4 +1,4 @@
|
||||
"""Add internal_id to banks and per-(bank, fact_type) partial HNSW indexes
|
||||
"""Add internal_id to banks and per-(bank, fact_type) partial vector indexes
|
||||
|
||||
Revision ID: d5e6f7a8b9c0
|
||||
Revises: a3b4c5d6e7f8
|
||||
@@ -6,25 +6,20 @@ Create Date: 2026-03-11
|
||||
|
||||
This migration:
|
||||
1. Adds internal_id UUID column to banks (stable identifier for index naming)
|
||||
2. Drops the global HNSW index (competes with per-bank partial indexes)
|
||||
3. Creates per-(bank_id, fact_type) partial HNSW indexes for all existing banks
|
||||
(new banks get indexes created at bank-creation time via bank_utils.create_bank_hnsw_indexes)
|
||||
2. Drops the global vector index (competes with per-bank partial indexes)
|
||||
3. Creates per-(bank_id, fact_type) partial vector indexes for all existing banks
|
||||
using the configured vector extension (HNSW for pgvector, DiskANN for
|
||||
pgvectorscale, vchordrq for vchord).
|
||||
(new banks get indexes created at bank-creation time via bank_utils.create_bank_vector_indexes)
|
||||
|
||||
Why per-(bank, fact_type) indexes:
|
||||
- fact_type-only partial indexes are never chosen by the planner when bank_id is in the WHERE
|
||||
clause, because the idx_memory_units_bank_id B-tree index always wins at planning time.
|
||||
- Per-(bank, fact_type) partial indexes have both predicates matching → planner selects them.
|
||||
- The global HNSW index competes for larger partitions (world, observation) and must be dropped.
|
||||
|
||||
For large deployments, create indexes CONCURRENTLY before running this migration:
|
||||
SELECT internal_id, bank_id FROM banks;
|
||||
-- for each bank and each fact_type in (world, experience, observation):
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_mu_emb_{ft}_{uid16}
|
||||
ON memory_units USING hnsw (embedding vector_cosine_ops)
|
||||
WHERE fact_type = '{ft}' AND bank_id = '{bank_id}';
|
||||
DROP INDEX CONCURRENTLY IF EXISTS idx_memory_units_embedding;
|
||||
- The global vector index competes for larger partitions (world, observation) and must be dropped.
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
@@ -35,7 +30,7 @@ down_revision: str | Sequence[str] | None = "c3d4e5f6g7h8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_HNSW_FACT_TYPES: dict[str, str] = {
|
||||
_FACT_TYPES: dict[str, str] = {
|
||||
"world": "worl",
|
||||
"experience": "expr",
|
||||
"observation": "obsv",
|
||||
@@ -47,6 +42,17 @@ def _get_schema_prefix() -> str:
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _vector_index_using_clause() -> str:
|
||||
"""Return the USING clause based on the configured vector extension."""
|
||||
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
elif ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
else:
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
@@ -56,33 +62,35 @@ def upgrade() -> None:
|
||||
)
|
||||
op.execute(f"ALTER TABLE {schema}banks ADD CONSTRAINT banks_internal_id_unique UNIQUE (internal_id)")
|
||||
|
||||
# 2. Drop any fact_type-only partial HNSW indexes that may exist from prior migrations
|
||||
# 2. Drop any fact_type-only partial indexes that may exist from prior migrations
|
||||
# (bank_id B-tree always wins over them when bank_id is in the WHERE clause)
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_world")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_observation")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_experience")
|
||||
|
||||
# 4. Drop global HNSW index (competes with per-bank partial indexes)
|
||||
# 4. Drop global vector index (competes with per-bank partial indexes)
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_embedding")
|
||||
|
||||
# 5. Create per-(bank, fact_type) partial HNSW indexes for all existing banks
|
||||
# 5. Create per-(bank, fact_type) partial vector indexes for all existing banks
|
||||
# using the configured extension (HNSW / DiskANN / vchordrq)
|
||||
bind = op.get_bind()
|
||||
schema_name = context.config.get_main_option("target_schema")
|
||||
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
|
||||
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
|
||||
using_clause = _vector_index_using_clause()
|
||||
|
||||
rows = bind.execute(text(f"SELECT bank_id, internal_id FROM {banks_ref}")).fetchall() # noqa: S608
|
||||
for row in rows:
|
||||
bank_id = row[0]
|
||||
internal_id = str(row[1]).replace("-", "")[:16]
|
||||
escaped_bank_id = bank_id.replace("'", "''")
|
||||
for ft, ft_short in _HNSW_FACT_TYPES.items():
|
||||
for ft, ft_short in _FACT_TYPES.items():
|
||||
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
|
||||
# Index name is schema-unqualified (indexes live in the schema of their table)
|
||||
bind.execute(
|
||||
text(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx_name} "
|
||||
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
|
||||
f"ON {table_ref} {using_clause} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
|
||||
)
|
||||
)
|
||||
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
"""Drop unused metadata column from documents table
|
||||
|
||||
Revision ID: d6e7f8a9b0c1
|
||||
Revises: c2d3e4f5g6h7, c5d6e7f8a9b0
|
||||
Create Date: 2026-03-30
|
||||
|
||||
The metadata column on documents was always stored as an empty dict {}.
|
||||
Actual document metadata is stored inside retain_params.metadata.
|
||||
|
||||
This migration was originally shipped in v0.4.22, then its file was deleted
|
||||
in v0.5.0 (and its revision ID accidentally reused by 2eee35aa3cfc).
|
||||
Restoring the file so that databases stamped at this revision can upgrade
|
||||
cleanly to v0.5.x+.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "d6e7f8a9b0c1"
|
||||
down_revision: str | Sequence[str] | None = ("c2d3e4f5g6h7", "c5d6e7f8a9b0")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}documents DROP COLUMN IF EXISTS metadata")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}documents ADD COLUMN IF NOT EXISTS metadata jsonb DEFAULT '{{}}'")
|
||||
+2
-2
@@ -5,7 +5,7 @@ Revises: e0a1b2c3d4e5
|
||||
Create Date: 2025-01-12
|
||||
|
||||
Add composite index on memory_links (from_unit_id, link_type, weight DESC)
|
||||
to optimize MPFP graph traversal queries that need top-k edges per type.
|
||||
to optimize graph traversal queries that need top-k edges per type.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
@@ -26,7 +26,7 @@ def _get_schema_prefix() -> str:
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Add composite index for efficient MPFP edge loading."""
|
||||
"""Add composite index for efficient graph retrieval edge loading."""
|
||||
schema = _get_schema_prefix()
|
||||
# Create composite index for efficient top-k per (from_node, link_type) queries
|
||||
# This enables LATERAL joins to use index-only scans with early termination
|
||||
|
||||
+22
-3
@@ -24,9 +24,28 @@ def upgrade() -> None:
|
||||
the memory_units rows survived with chunk_id = NULL, leaving ghost records.
|
||||
Switching to CASCADE ensures they are removed together with their chunk.
|
||||
"""
|
||||
op.drop_constraint("memory_units_chunk_fkey", "memory_units", type_="foreignkey")
|
||||
op.create_foreign_key(
|
||||
"memory_units_chunk_fkey", "memory_units", "chunks", ["chunk_id"], ["chunk_id"], ondelete="CASCADE"
|
||||
from alembic import context
|
||||
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
schema_prefix = f'"{schema}".' if schema else ""
|
||||
# Use raw SQL with IF EXISTS so this is safe on schemas where the FK was
|
||||
# already dropped or never existed under this name.
|
||||
op.execute(f"ALTER TABLE {schema_prefix}memory_units DROP CONSTRAINT IF EXISTS memory_units_chunk_fkey")
|
||||
# Use a DO block so the ADD is also idempotent: if the FK already exists (e.g.
|
||||
# the schema was provisioned after the base migration already added it) the
|
||||
# duplicate_object exception is swallowed rather than failing the migration.
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$ BEGIN
|
||||
ALTER TABLE {schema_prefix}memory_units
|
||||
ADD CONSTRAINT memory_units_chunk_fkey
|
||||
FOREIGN KEY (chunk_id)
|
||||
REFERENCES {schema_prefix}chunks (chunk_id)
|
||||
ON DELETE CASCADE;
|
||||
EXCEPTION
|
||||
WHEN duplicate_object THEN NULL;
|
||||
END $$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
|
||||
+83
@@ -0,0 +1,83 @@
|
||||
"""remove_opinion_fact_type
|
||||
|
||||
Revision ID: g2h3i4j5k6l7
|
||||
Revises: f1a2b3c4d5e6
|
||||
Create Date: 2026-04-02
|
||||
|
||||
Remove the deprecated 'opinion' fact type: drop opinion-specific indexes,
|
||||
update CHECK constraints, delete any remaining opinion rows, and drop the
|
||||
confidence_score column (was only used for opinions, always NULL otherwise).
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "g2h3i4j5k6l7"
|
||||
down_revision: str | Sequence[str] | None = "f1a2b3c4d5e6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# 1. Delete any remaining opinion rows
|
||||
op.execute(f"DELETE FROM {schema}memory_units WHERE fact_type = 'opinion'")
|
||||
|
||||
# 2. Drop opinion-specific indexes
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_opinion_confidence")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_opinion_date")
|
||||
|
||||
# 3. Drop confidence_score constraints and column (only used for opinions, always NULL otherwise)
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS confidence_score_fact_type_check")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS memory_units_confidence_score_check")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS confidence_score")
|
||||
|
||||
# 4. Replace fact_type CHECK constraint
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS memory_units_fact_type_check")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT memory_units_fact_type_check "
|
||||
f"CHECK (fact_type IN ('world', 'experience', 'observation'))"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Restore confidence_score column
|
||||
op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS confidence_score float")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT memory_units_confidence_score_check "
|
||||
f"CHECK (confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0))"
|
||||
)
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT confidence_score_fact_type_check "
|
||||
f"CHECK ((fact_type = 'opinion' AND confidence_score IS NOT NULL) OR "
|
||||
f"(fact_type = 'observation') OR "
|
||||
f"(fact_type NOT IN ('opinion', 'observation') AND confidence_score IS NULL))"
|
||||
)
|
||||
|
||||
# Restore original fact_type CHECK constraint (with opinion)
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP CONSTRAINT IF EXISTS memory_units_fact_type_check")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}memory_units ADD CONSTRAINT memory_units_fact_type_check "
|
||||
f"CHECK (fact_type IN ('world', 'experience', 'opinion', 'observation'))"
|
||||
)
|
||||
|
||||
# Recreate opinion indexes
|
||||
op.execute(
|
||||
f"CREATE INDEX idx_memory_units_opinion_confidence ON {schema}memory_units "
|
||||
f"(bank_id, confidence_score DESC) WHERE fact_type = 'opinion'"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX idx_memory_units_opinion_date ON {schema}memory_units "
|
||||
f"(bank_id, event_date DESC) WHERE fact_type = 'opinion'"
|
||||
)
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
"""Merge 3 migration heads and add unit_entities composite index
|
||||
|
||||
Revision ID: h3i4j5k6l7m8
|
||||
Revises: a4b5c6d7e8f9, g2h3i4j5k6l7
|
||||
Create Date: 2026-04-07
|
||||
|
||||
Merges three unmerged migration heads into one, and adds a composite index
|
||||
(entity_id, unit_id) on unit_entities for index-only scans in the LATERAL
|
||||
entity expansion query.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "h3i4j5k6l7m8"
|
||||
down_revision: str | Sequence[str] | None = ("a4b5c6d7e8f9", "g2h3i4j5k6l7")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# Composite index enables index-only scans for entity_id -> unit_id lookups
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_unit_entities_entity_unit ON {schema}unit_entities (entity_id, unit_id)"
|
||||
)
|
||||
# Drop the now-redundant single-column index
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_unit_entities_entity")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_unit_entities_entity_unit")
|
||||
# Restore the single-column index
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_unit_entities_entity ON {schema}unit_entities (entity_id)")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -12,44 +12,9 @@ from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.memory_engine import _current_schema
|
||||
from hindsight_api.extensions import MCPExtension, load_extension
|
||||
from hindsight_api.extensions.tenant import AuthenticationError
|
||||
from hindsight_api.mcp_tools import MCPToolsConfig, register_mcp_tools
|
||||
from hindsight_api.mcp_tools import _ALL_TOOLS, MCPToolsConfig, register_mcp_tools
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
# All tools available in the system (explicit list — no wildcards)
|
||||
_ALL_TOOLS: frozenset[str] = frozenset(
|
||||
{
|
||||
"retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_banks",
|
||||
"create_bank",
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"create_mental_model",
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"delete_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
"list_operations",
|
||||
"get_operation",
|
||||
"cancel_operation",
|
||||
"list_tags",
|
||||
"get_bank",
|
||||
"get_bank_stats",
|
||||
"update_bank",
|
||||
"delete_bank",
|
||||
"clear_memories",
|
||||
}
|
||||
)
|
||||
|
||||
# Configure logging from HINDSIGHT_API_LOG_LEVEL environment variable
|
||||
_log_level_str = os.environ.get("HINDSIGHT_API_LOG_LEVEL", "info").lower()
|
||||
_log_level_map = {
|
||||
@@ -83,6 +48,9 @@ _current_api_key: ContextVar[str | None] = ContextVar("current_api_key", default
|
||||
_current_tenant_id: ContextVar[str | None] = ContextVar("current_tenant_id", default=None)
|
||||
_current_api_key_id: ContextVar[str | None] = ContextVar("current_api_key_id", default=None)
|
||||
|
||||
# Context variable for MCP pre-authentication flag (set when MCP_AUTH_TOKEN validates)
|
||||
_current_mcp_authenticated: ContextVar[bool] = ContextVar("current_mcp_authenticated", default=False)
|
||||
|
||||
|
||||
def get_current_bank_id() -> str | None:
|
||||
"""Get the current bank_id from context."""
|
||||
@@ -104,6 +72,11 @@ def get_current_api_key_id() -> str | None:
|
||||
return _current_api_key_id.get()
|
||||
|
||||
|
||||
def get_current_mcp_authenticated() -> bool:
|
||||
"""Get whether the request was pre-authenticated by MCP transport auth."""
|
||||
return _current_mcp_authenticated.get()
|
||||
|
||||
|
||||
def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"""
|
||||
Create and configure the Hindsight MCP server.
|
||||
@@ -124,6 +97,7 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
_SINGLE_BANK_TOOLS: frozenset[str] = frozenset(
|
||||
{
|
||||
"retain",
|
||||
"sync_retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_mental_models",
|
||||
@@ -164,6 +138,7 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
api_key_resolver=get_current_api_key, # Propagate API key for tenant auth
|
||||
tenant_id_resolver=get_current_tenant_id, # Propagate tenant_id for usage metering
|
||||
api_key_id_resolver=get_current_api_key_id, # Propagate api_key_id for usage metering
|
||||
mcp_authenticated_resolver=get_current_mcp_authenticated, # Propagate MCP pre-auth flag
|
||||
include_bank_id_param=multi_bank,
|
||||
tools=base_tools,
|
||||
)
|
||||
@@ -182,24 +157,65 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
return mcp
|
||||
|
||||
|
||||
def _get_mcp_tools(mcp: FastMCP) -> dict:
|
||||
"""Get tool name→object mapping, compatible with FastMCP 2.x and 3.x."""
|
||||
# FastMCP 2.x: _tool_manager._tools
|
||||
if hasattr(mcp, "_tool_manager"):
|
||||
return mcp._tool_manager._tools # type: ignore[union-attr]
|
||||
# FastMCP 3.x: _local_provider._components with "tool:" prefix
|
||||
if hasattr(mcp, "_local_provider"):
|
||||
return {
|
||||
k.split(":")[1].split("@")[0]: v
|
||||
for k, v in mcp._local_provider._components.items() # type: ignore[union-attr]
|
||||
if k.startswith("tool:")
|
||||
}
|
||||
msg = "Cannot locate tools on FastMCP instance"
|
||||
raise AttributeError(msg)
|
||||
|
||||
|
||||
def _make_tools_tolerant(mcp: FastMCP) -> None:
|
||||
"""Wrap all tool run methods to strip unknown arguments before validation.
|
||||
"""Wrap all tool run methods to strip unknown arguments and coerce string-encoded JSON.
|
||||
|
||||
LLMs frequently add extra fields like "explanation" or "reasoning" to tool calls.
|
||||
FastMCP's Pydantic TypeAdapter rejects these with "Unexpected keyword argument".
|
||||
This wraps each tool's run() to filter arguments to only known parameters.
|
||||
|
||||
LLMs also frequently serialize list/dict arguments as JSON strings instead of native
|
||||
types (e.g., tags='["a","b"]' instead of tags=["a","b"]). This auto-coerces them.
|
||||
|
||||
This wraps each tool's run() to apply both fixes before validation.
|
||||
"""
|
||||
try:
|
||||
for name, tool in mcp._tool_manager._tools.items():
|
||||
tools = _get_mcp_tools(mcp)
|
||||
for name, tool in tools.items():
|
||||
if hasattr(tool, "parameters") and tool.parameters:
|
||||
allowed = set(tool.parameters.get("properties", {}).keys())
|
||||
properties = tool.parameters.get("properties", {})
|
||||
allowed = set(properties.keys())
|
||||
|
||||
# Build sets of parameter names that expect array or object types.
|
||||
# Handles both direct types {"type": "array"} and anyOf/oneOf unions
|
||||
# like {"anyOf": [{"type": "array", ...}, {"type": "null"}]}.
|
||||
array_params: set[str] = set()
|
||||
object_params: set[str] = set()
|
||||
for param_name, param_schema in properties.items():
|
||||
_collect_coercible_types(param_schema, param_name, array_params, object_params)
|
||||
|
||||
original_run = tool.run
|
||||
|
||||
async def _tolerant_run(arguments, _allowed=allowed, _orig=original_run):
|
||||
async def _tolerant_run(
|
||||
arguments,
|
||||
_allowed=allowed,
|
||||
_orig=original_run,
|
||||
_array_params=array_params,
|
||||
_object_params=object_params,
|
||||
):
|
||||
extra_keys = set(arguments.keys()) - _allowed
|
||||
if extra_keys:
|
||||
logger.debug(f"Stripping unknown arguments from tool call: {extra_keys}")
|
||||
arguments = {k: v for k, v in arguments.items() if k in _allowed}
|
||||
|
||||
# Coerce string-encoded JSON for list/dict parameters
|
||||
arguments = _coerce_string_json(arguments, _array_params, _object_params)
|
||||
|
||||
return await _orig(arguments)
|
||||
|
||||
# FunctionTool is a Pydantic model with extra='forbid', so use
|
||||
@@ -209,6 +225,59 @@ def _make_tools_tolerant(mcp: FastMCP) -> None:
|
||||
logger.warning(f"Could not make tools tolerant of extra arguments: {e}")
|
||||
|
||||
|
||||
def _collect_coercible_types(schema: dict, param_name: str, array_params: set[str], object_params: set[str]) -> None:
|
||||
"""Check a JSON Schema property and add param_name to array_params/object_params if applicable."""
|
||||
# Direct type
|
||||
schema_type = schema.get("type")
|
||||
if schema_type == "array":
|
||||
array_params.add(param_name)
|
||||
return
|
||||
if schema_type == "object":
|
||||
object_params.add(param_name)
|
||||
return
|
||||
|
||||
# anyOf / oneOf unions (e.g., list[str] | None → {"anyOf": [{"type": "array"}, {"type": "null"}]})
|
||||
for variant in schema.get("anyOf", []) + schema.get("oneOf", []):
|
||||
variant_type = variant.get("type")
|
||||
if variant_type == "array":
|
||||
array_params.add(param_name)
|
||||
return
|
||||
if variant_type == "object":
|
||||
object_params.add(param_name)
|
||||
return
|
||||
|
||||
|
||||
def _coerce_string_json(arguments: dict, array_params: set[str], object_params: set[str]) -> dict:
|
||||
"""Auto-coerce string-encoded JSON arrays/objects to native types.
|
||||
|
||||
LLM agents frequently serialize list and dict tool arguments as JSON strings.
|
||||
This is backward-compatible: native arrays/objects pass through unchanged.
|
||||
"""
|
||||
for param_name in array_params:
|
||||
val = arguments.get(param_name)
|
||||
if isinstance(val, str):
|
||||
try:
|
||||
parsed = json.loads(val)
|
||||
if isinstance(parsed, list):
|
||||
arguments = {**arguments, param_name: parsed}
|
||||
logger.debug(f"Coerced string to list for parameter '{param_name}'")
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
for param_name in object_params:
|
||||
val = arguments.get(param_name)
|
||||
if isinstance(val, str):
|
||||
try:
|
||||
parsed = json.loads(val)
|
||||
if isinstance(parsed, dict):
|
||||
arguments = {**arguments, param_name: parsed}
|
||||
logger.debug(f"Coerced string to dict for parameter '{param_name}'")
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
return arguments
|
||||
|
||||
|
||||
class MCPMiddleware:
|
||||
"""ASGI middleware that intercepts MCP requests and routes to appropriate MCP server.
|
||||
|
||||
@@ -272,10 +341,12 @@ class MCPMiddleware:
|
||||
self.single_bank_server = single_bank_server
|
||||
else:
|
||||
# Create servers internally (for direct construction / tests)
|
||||
global_config = _get_raw_config()
|
||||
stateless = global_config.mcp_stateless
|
||||
self.multi_bank_server = create_mcp_server(memory, multi_bank=True)
|
||||
self.multi_bank_app = self.multi_bank_server.http_app(path="/", stateless_http=True)
|
||||
self.multi_bank_app = self.multi_bank_server.http_app(path="/", stateless_http=stateless)
|
||||
self.single_bank_server = create_mcp_server(memory, multi_bank=False)
|
||||
self.single_bank_app = self.single_bank_server.http_app(path="/", stateless_http=True)
|
||||
self.single_bank_app = self.single_bank_server.http_app(path="/", stateless_http=stateless)
|
||||
|
||||
def _get_header(self, scope: dict, name: str) -> str | None:
|
||||
"""Extract a header value from ASGI scope."""
|
||||
@@ -298,6 +369,17 @@ class MCPMiddleware:
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
# Handle GET-before-POST gracefully (Claude Code v2.1.84+ sends GET probe before POST initialize).
|
||||
# Without a valid Mcp-Session-Id, GET has no meaningful response — return 200 OK so
|
||||
# the client proceeds to POST initialize instead of marking the server as failed.
|
||||
method = scope.get("method", "")
|
||||
if method == "GET":
|
||||
session_id = self._get_header(scope, "Mcp-Session-Id")
|
||||
if not session_id:
|
||||
logger.debug("MCP GET without session ID (client probe) — returning 200 OK")
|
||||
await self._send_ok(send)
|
||||
return
|
||||
|
||||
# Strip prefix from path
|
||||
path = path[len(self.prefix) :] or "/"
|
||||
|
||||
@@ -312,6 +394,7 @@ class MCPMiddleware:
|
||||
tenant_context = None
|
||||
auth_tenant_id: str | None = None
|
||||
auth_api_key_id: str | None = None
|
||||
mcp_pre_authenticated = False
|
||||
if MCP_AUTH_TOKEN:
|
||||
# Legacy authentication mode - validate against static token
|
||||
if not auth_token:
|
||||
@@ -320,8 +403,9 @@ class MCPMiddleware:
|
||||
if auth_token != MCP_AUTH_TOKEN:
|
||||
await self._send_error(send, 401, "Invalid authentication token")
|
||||
return
|
||||
# Legacy mode doesn't use tenant schemas
|
||||
# Legacy mode: mark as pre-authenticated so tenant extension won't re-validate
|
||||
tenant_context = None
|
||||
mcp_pre_authenticated = True
|
||||
else:
|
||||
# Use TenantExtension.authenticate_mcp() for auth
|
||||
try:
|
||||
@@ -368,13 +452,15 @@ class MCPMiddleware:
|
||||
# - Header/env bank_id → multi-bank app (bank_id param, all tools)
|
||||
target_app = self.single_bank_app if bank_id_from_path else self.multi_bank_app
|
||||
|
||||
# Set bank_id, api_key, tenant_id, and api_key_id context
|
||||
# Set bank_id, api_key, tenant_id, api_key_id, and mcp_authenticated context
|
||||
bank_id_token = _current_bank_id.set(bank_id)
|
||||
# Store the auth token for tenant extension to validate
|
||||
api_key_token = _current_api_key.set(auth_token) if auth_token else None
|
||||
# Store tenant_id and api_key_id from authentication for usage metering
|
||||
tenant_id_token = _current_tenant_id.set(auth_tenant_id) if auth_tenant_id else None
|
||||
api_key_id_token = _current_api_key_id.set(auth_api_key_id) if auth_api_key_id else None
|
||||
# Store MCP pre-authentication flag to skip tenant re-validation
|
||||
mcp_auth_token = _current_mcp_authenticated.set(mcp_pre_authenticated)
|
||||
try:
|
||||
new_scope = scope.copy()
|
||||
new_scope["path"] = new_path
|
||||
@@ -419,9 +505,26 @@ class MCPMiddleware:
|
||||
_current_tenant_id.reset(tenant_id_token)
|
||||
if api_key_id_token is not None:
|
||||
_current_api_key_id.reset(api_key_id_token)
|
||||
_current_mcp_authenticated.reset(mcp_auth_token)
|
||||
if schema_token is not None:
|
||||
_current_schema.reset(schema_token)
|
||||
|
||||
async def _send_ok(self, send):
|
||||
"""Send a 200 OK response with empty body (used for GET probes without session)."""
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 200,
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
}
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": b"{}",
|
||||
}
|
||||
)
|
||||
|
||||
async def _send_error(self, send, status: int, message: str, extra_headers: dict[str, str] | None = None):
|
||||
"""Send an error response."""
|
||||
body = json.dumps({"error": message}).encode()
|
||||
@@ -452,10 +555,14 @@ def create_mcp_servers(memory: MemoryEngine):
|
||||
Returns:
|
||||
Tuple of (multi_bank_server, single_bank_server, multi_bank_app, single_bank_app)
|
||||
"""
|
||||
global_config = _get_raw_config()
|
||||
stateless = global_config.mcp_stateless
|
||||
|
||||
multi_bank_server = create_mcp_server(memory, multi_bank=True)
|
||||
multi_bank_app = multi_bank_server.http_app(path="/", stateless_http=True)
|
||||
multi_bank_app = multi_bank_server.http_app(path="/", stateless_http=stateless)
|
||||
|
||||
single_bank_server = create_mcp_server(memory, multi_bank=False)
|
||||
single_bank_app = single_bank_server.http_app(path="/", stateless_http=True)
|
||||
single_bank_app = single_bank_server.http_app(path="/", stateless_http=stateless)
|
||||
|
||||
logger.info(f"MCP servers created (stateless_http={stateless})")
|
||||
return multi_bank_server, single_bank_server, multi_bank_app, single_bank_app
|
||||
|
||||
@@ -118,6 +118,7 @@ def normalize_config_dict(config: dict[str, Any]) -> dict[str, Any]:
|
||||
|
||||
# Environment variable names
|
||||
ENV_DATABASE_URL = "HINDSIGHT_API_DATABASE_URL"
|
||||
ENV_MIGRATION_DATABASE_URL = "HINDSIGHT_API_MIGRATION_DATABASE_URL"
|
||||
ENV_DATABASE_SCHEMA = "HINDSIGHT_API_DATABASE_SCHEMA"
|
||||
ENV_LLM_PROVIDER = "HINDSIGHT_API_LLM_PROVIDER"
|
||||
ENV_LLM_API_KEY = "HINDSIGHT_API_LLM_API_KEY"
|
||||
@@ -130,10 +131,12 @@ ENV_LLM_MAX_BACKOFF = "HINDSIGHT_API_LLM_MAX_BACKOFF"
|
||||
ENV_LLM_TIMEOUT = "HINDSIGHT_API_LLM_TIMEOUT"
|
||||
ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER"
|
||||
ENV_LLM_OPENAI_SERVICE_TIER = "HINDSIGHT_API_LLM_OPENAI_SERVICE_TIER"
|
||||
ENV_LLM_EXTRA_BODY = "HINDSIGHT_API_LLM_EXTRA_BODY"
|
||||
|
||||
# Defaults for service tiers
|
||||
DEFAULT_LLM_GROQ_SERVICE_TIER = "auto" # "on_demand", "flex", or "auto"
|
||||
DEFAULT_LLM_OPENAI_SERVICE_TIER = None # None (default) or "flex" (50% cheaper)
|
||||
DEFAULT_LLM_EXTRA_BODY = None # None = no extra body params; JSON dict merged into OpenAI extra_body
|
||||
|
||||
# Per-operation LLM configuration (optional, falls back to global LLM config)
|
||||
ENV_RETAIN_LLM_PROVIDER = "HINDSIGHT_API_RETAIN_LLM_PROVIDER"
|
||||
@@ -174,6 +177,15 @@ ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL"
|
||||
ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY"
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"
|
||||
ENV_EMBEDDINGS_OPENAI_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL"
|
||||
ENV_EMBEDDINGS_OPENAI_BATCH_SIZE = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"
|
||||
|
||||
# Gemini/Vertex AI embeddings configuration
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY"
|
||||
ENV_EMBEDDINGS_GEMINI_MODEL = "HINDSIGHT_API_EMBEDDINGS_GEMINI_MODEL"
|
||||
ENV_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY"
|
||||
ENV_EMBEDDINGS_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_PROJECT_ID"
|
||||
ENV_EMBEDDINGS_VERTEXAI_REGION = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_REGION"
|
||||
ENV_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY"
|
||||
|
||||
# Cohere configuration (separate for embeddings and reranker)
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY = "HINDSIGHT_API_EMBEDDINGS_COHERE_API_KEY"
|
||||
@@ -183,6 +195,13 @@ ENV_RERANKER_COHERE_API_KEY = "HINDSIGHT_API_RERANKER_COHERE_API_KEY"
|
||||
ENV_RERANKER_COHERE_MODEL = "HINDSIGHT_API_RERANKER_COHERE_MODEL"
|
||||
ENV_RERANKER_COHERE_BASE_URL = "HINDSIGHT_API_RERANKER_COHERE_BASE_URL"
|
||||
|
||||
# OpenRouter configuration (embeddings and reranker)
|
||||
ENV_OPENROUTER_API_KEY = "HINDSIGHT_API_OPENROUTER_API_KEY"
|
||||
ENV_EMBEDDINGS_OPENROUTER_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENROUTER_API_KEY"
|
||||
ENV_EMBEDDINGS_OPENROUTER_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENROUTER_MODEL"
|
||||
ENV_RERANKER_OPENROUTER_API_KEY = "HINDSIGHT_API_RERANKER_OPENROUTER_API_KEY"
|
||||
ENV_RERANKER_OPENROUTER_MODEL = "HINDSIGHT_API_RERANKER_OPENROUTER_MODEL"
|
||||
|
||||
# Deprecated: Legacy shared Cohere API key (for backward compatibility)
|
||||
ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY"
|
||||
|
||||
@@ -199,6 +218,8 @@ ENV_RERANKER_LITELLM_MAX_TOKENS_PER_DOC = "HINDSIGHT_API_RERANKER_LITELLM_MAX_TO
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_BASE = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_BASE"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT"
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY"
|
||||
ENV_RERANKER_LITELLM_SDK_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL"
|
||||
ENV_RERANKER_LITELLM_SDK_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE"
|
||||
@@ -218,6 +239,7 @@ ENV_RERANKER_LOCAL_BATCH_SIZE = "HINDSIGHT_API_RERANKER_LOCAL_BATCH_SIZE"
|
||||
ENV_RERANKER_TEI_URL = "HINDSIGHT_API_RERANKER_TEI_URL"
|
||||
ENV_RERANKER_TEI_BATCH_SIZE = "HINDSIGHT_API_RERANKER_TEI_BATCH_SIZE"
|
||||
ENV_RERANKER_TEI_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_TEI_MAX_CONCURRENT"
|
||||
ENV_RERANKER_TEI_HTTP_TIMEOUT = "HINDSIGHT_API_RERANKER_TEI_HTTP_TIMEOUT"
|
||||
ENV_RERANKER_MAX_CANDIDATES = "HINDSIGHT_API_RERANKER_MAX_CANDIDATES"
|
||||
ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL"
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
|
||||
@@ -225,6 +247,17 @@ ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
|
||||
# ZeroEntropy configuration (reranker only)
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY = "HINDSIGHT_API_RERANKER_ZEROENTROPY_API_KEY"
|
||||
ENV_RERANKER_ZEROENTROPY_MODEL = "HINDSIGHT_API_RERANKER_ZEROENTROPY_MODEL"
|
||||
ENV_RERANKER_ZEROENTROPY_BASE_URL = "HINDSIGHT_API_RERANKER_ZEROENTROPY_BASE_URL"
|
||||
|
||||
# SiliconFlow configuration (reranker only; Cohere-compatible /rerank endpoint)
|
||||
ENV_RERANKER_SILICONFLOW_API_KEY = "HINDSIGHT_API_RERANKER_SILICONFLOW_API_KEY"
|
||||
ENV_RERANKER_SILICONFLOW_MODEL = "HINDSIGHT_API_RERANKER_SILICONFLOW_MODEL"
|
||||
ENV_RERANKER_SILICONFLOW_BASE_URL = "HINDSIGHT_API_RERANKER_SILICONFLOW_BASE_URL"
|
||||
|
||||
# Google Discovery Engine reranker configuration
|
||||
ENV_RERANKER_GOOGLE_MODEL = "HINDSIGHT_API_RERANKER_GOOGLE_MODEL"
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID = "HINDSIGHT_API_RERANKER_GOOGLE_PROJECT_ID"
|
||||
ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY"
|
||||
|
||||
ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION"
|
||||
ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION"
|
||||
@@ -234,16 +267,20 @@ ENV_PORT = "HINDSIGHT_API_PORT"
|
||||
ENV_BASE_PATH = "HINDSIGHT_API_BASE_PATH"
|
||||
ENV_LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL"
|
||||
ENV_LOG_FORMAT = "HINDSIGHT_API_LOG_FORMAT"
|
||||
ENV_LOG_JSON_FIELDS = "HINDSIGHT_API_LOG_JSON_FIELDS"
|
||||
ENV_WORKERS = "HINDSIGHT_API_WORKERS"
|
||||
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
|
||||
ENV_MCP_ENABLED_TOOLS = "HINDSIGHT_API_MCP_ENABLED_TOOLS"
|
||||
ENV_MCP_STATELESS = "HINDSIGHT_API_MCP_STATELESS"
|
||||
ENV_ENABLE_BANK_CONFIG_API = "HINDSIGHT_API_ENABLE_BANK_CONFIG_API"
|
||||
ENV_DEFAULT_BANK_TEMPLATE = "HINDSIGHT_API_DEFAULT_BANK_TEMPLATE"
|
||||
ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
|
||||
ENV_MPFP_TOP_K_NEIGHBORS = "HINDSIGHT_API_MPFP_TOP_K_NEIGHBORS"
|
||||
ENV_RECALL_MAX_CONCURRENT = "HINDSIGHT_API_RECALL_MAX_CONCURRENT"
|
||||
ENV_RECALL_CONNECTION_BUDGET = "HINDSIGHT_API_RECALL_CONNECTION_BUDGET"
|
||||
ENV_RECALL_MAX_QUERY_TOKENS = "HINDSIGHT_API_RECALL_MAX_QUERY_TOKENS"
|
||||
ENV_MENTAL_MODEL_REFRESH_CONCURRENCY = "HINDSIGHT_API_MENTAL_MODEL_REFRESH_CONCURRENCY"
|
||||
ENV_LINK_EXPANSION_PER_ENTITY_LIMIT = "HINDSIGHT_API_LINK_EXPANSION_PER_ENTITY_LIMIT"
|
||||
ENV_LINK_EXPANSION_TIMEOUT = "HINDSIGHT_API_LINK_EXPANSION_TIMEOUT"
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
ENV_OTEL_TRACES_ENABLED = "HINDSIGHT_API_OTEL_TRACES_ENABLED"
|
||||
@@ -251,6 +288,7 @@ ENV_OTEL_EXPORTER_OTLP_ENDPOINT = "HINDSIGHT_API_OTEL_EXPORTER_OTLP_ENDPOINT"
|
||||
ENV_OTEL_EXPORTER_OTLP_HEADERS = "HINDSIGHT_API_OTEL_EXPORTER_OTLP_HEADERS"
|
||||
ENV_OTEL_SERVICE_NAME = "HINDSIGHT_API_OTEL_SERVICE_NAME"
|
||||
ENV_OTEL_DEPLOYMENT_ENVIRONMENT = "HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT"
|
||||
ENV_METRICS_INCLUDE_BANK_ID = "HINDSIGHT_API_METRICS_INCLUDE_BANK_ID"
|
||||
|
||||
# Vertex AI configuration
|
||||
ENV_LLM_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"
|
||||
@@ -272,6 +310,7 @@ ENV_RETAIN_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS"
|
||||
ENV_RETAIN_ENTITY_LOOKUP = "HINDSIGHT_API_RETAIN_ENTITY_LOOKUP"
|
||||
ENV_RETAIN_BATCH_ENABLED = "HINDSIGHT_API_RETAIN_BATCH_ENABLED"
|
||||
ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_RETAIN_BATCH_POLL_INTERVAL_SECONDS"
|
||||
ENV_RETAIN_CHUNK_BATCH_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_BATCH_SIZE"
|
||||
|
||||
# File storage configuration
|
||||
ENV_FILE_STORAGE_TYPE = "HINDSIGHT_API_FILE_STORAGE_TYPE"
|
||||
@@ -297,13 +336,16 @@ ENV_FILE_DELETE_AFTER_RETAIN = "HINDSIGHT_API_FILE_DELETE_AFTER_RETAIN"
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS"
|
||||
ENV_CONSOLIDATION_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_BATCH_SIZE"
|
||||
ENV_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = "HINDSIGHT_API_CONSOLIDATION_MAX_MEMORIES_PER_ROUND"
|
||||
ENV_CONSOLIDATION_LLM_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_LLM_BATCH_SIZE"
|
||||
ENV_CONSOLIDATION_MAX_TOKENS = "HINDSIGHT_API_CONSOLIDATION_MAX_TOKENS"
|
||||
ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = "HINDSIGHT_API_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS"
|
||||
ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION = (
|
||||
"HINDSIGHT_API_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION"
|
||||
)
|
||||
ENV_CONSOLIDATION_MAX_ATTEMPTS = "HINDSIGHT_API_CONSOLIDATION_MAX_ATTEMPTS"
|
||||
ENV_OBSERVATIONS_MISSION = "HINDSIGHT_API_OBSERVATIONS_MISSION"
|
||||
ENV_MAX_OBSERVATIONS_PER_SCOPE = "HINDSIGHT_API_MAX_OBSERVATIONS_PER_SCOPE"
|
||||
ENV_ENABLE_OBSERVATION_HISTORY = "HINDSIGHT_API_ENABLE_OBSERVATION_HISTORY"
|
||||
ENV_ENABLE_MENTAL_MODEL_HISTORY = "HINDSIGHT_API_ENABLE_MENTAL_MODEL_HISTORY"
|
||||
|
||||
@@ -313,6 +355,14 @@ ENV_WEBHOOK_SECRET = "HINDSIGHT_API_WEBHOOK_SECRET"
|
||||
ENV_WEBHOOK_EVENT_TYPES = "HINDSIGHT_API_WEBHOOK_EVENT_TYPES"
|
||||
ENV_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS"
|
||||
|
||||
# Built-in llama.cpp configuration (for provider=llamacpp)
|
||||
ENV_LLAMACPP_MODEL_PATH = "HINDSIGHT_API_LLAMACPP_MODEL_PATH"
|
||||
ENV_LLAMACPP_GPU_LAYERS = "HINDSIGHT_API_LLAMACPP_GPU_LAYERS"
|
||||
ENV_LLAMACPP_CONTEXT_SIZE = "HINDSIGHT_API_LLAMACPP_CONTEXT_SIZE"
|
||||
ENV_LLAMACPP_CHAT_FORMAT = "HINDSIGHT_API_LLAMACPP_CHAT_FORMAT"
|
||||
ENV_LLAMACPP_NO_GRAMMAR = "HINDSIGHT_API_LLAMACPP_NO_GRAMMAR"
|
||||
ENV_LLAMACPP_EXTRA_ARGS = "HINDSIGHT_API_LLAMACPP_EXTRA_ARGS"
|
||||
|
||||
# Optimization flags
|
||||
ENV_SKIP_LLM_VERIFICATION = "HINDSIGHT_API_SKIP_LLM_VERIFICATION"
|
||||
ENV_LAZY_RERANKER = "HINDSIGHT_API_LAZY_RERANKER"
|
||||
@@ -325,6 +375,7 @@ ENV_DB_POOL_MIN_SIZE = "HINDSIGHT_API_DB_POOL_MIN_SIZE"
|
||||
ENV_DB_POOL_MAX_SIZE = "HINDSIGHT_API_DB_POOL_MAX_SIZE"
|
||||
ENV_DB_COMMAND_TIMEOUT = "HINDSIGHT_API_DB_COMMAND_TIMEOUT"
|
||||
ENV_DB_ACQUIRE_TIMEOUT = "HINDSIGHT_API_DB_ACQUIRE_TIMEOUT"
|
||||
ENV_DB_STATEMENT_TIMEOUT = "HINDSIGHT_API_DB_STATEMENT_TIMEOUT"
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED"
|
||||
@@ -334,11 +385,33 @@ ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES"
|
||||
ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT"
|
||||
ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS"
|
||||
ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS"
|
||||
ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT"
|
||||
|
||||
# Reflect agent settings
|
||||
ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS"
|
||||
ENV_REFLECT_MAX_CONTEXT_TOKENS = "HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS"
|
||||
ENV_REFLECT_WALL_TIMEOUT = "HINDSIGHT_API_REFLECT_WALL_TIMEOUT"
|
||||
ENV_REFLECT_MISSION = "HINDSIGHT_API_REFLECT_MISSION"
|
||||
ENV_REFLECT_SOURCE_FACTS_MAX_TOKENS = "HINDSIGHT_API_REFLECT_SOURCE_FACTS_MAX_TOKENS"
|
||||
ENV_RECALL_INCLUDE_CHUNKS = "HINDSIGHT_API_RECALL_INCLUDE_CHUNKS"
|
||||
ENV_RECALL_MAX_TOKENS = "HINDSIGHT_API_RECALL_MAX_TOKENS"
|
||||
ENV_RECALL_CHUNKS_MAX_TOKENS = "HINDSIGHT_API_RECALL_CHUNKS_MAX_TOKENS"
|
||||
|
||||
# Recall budget mapping (budget enum -> thinking_budget integer)
|
||||
ENV_RECALL_BUDGET_FUNCTION = "HINDSIGHT_API_RECALL_BUDGET_FUNCTION"
|
||||
ENV_RECALL_BUDGET_FIXED_LOW = "HINDSIGHT_API_RECALL_BUDGET_FIXED_LOW"
|
||||
ENV_RECALL_BUDGET_FIXED_MID = "HINDSIGHT_API_RECALL_BUDGET_FIXED_MID"
|
||||
ENV_RECALL_BUDGET_FIXED_HIGH = "HINDSIGHT_API_RECALL_BUDGET_FIXED_HIGH"
|
||||
ENV_RECALL_BUDGET_ADAPTIVE_LOW = "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_LOW"
|
||||
ENV_RECALL_BUDGET_ADAPTIVE_MID = "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_MID"
|
||||
ENV_RECALL_BUDGET_ADAPTIVE_HIGH = "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_HIGH"
|
||||
ENV_RECALL_BUDGET_MIN = "HINDSIGHT_API_RECALL_BUDGET_MIN"
|
||||
ENV_RECALL_BUDGET_MAX = "HINDSIGHT_API_RECALL_BUDGET_MAX"
|
||||
|
||||
# Audit log settings
|
||||
ENV_AUDIT_LOG_ENABLED = "HINDSIGHT_API_AUDIT_LOG_ENABLED"
|
||||
ENV_AUDIT_LOG_ACTIONS = "HINDSIGHT_API_AUDIT_LOG_ACTIONS"
|
||||
ENV_AUDIT_LOG_RETENTION_DAYS = "HINDSIGHT_API_AUDIT_LOG_RETENTION_DAYS"
|
||||
|
||||
# Disposition settings
|
||||
ENV_DISPOSITION_SKEPTICISM = "HINDSIGHT_API_DISPOSITION_SKEPTICISM"
|
||||
@@ -353,20 +426,33 @@ DEFAULT_LLM_PROVIDER = "openai"
|
||||
# Provider-specific default models
|
||||
PROVIDER_DEFAULT_MODELS = {
|
||||
"openai": "gpt-4o-mini",
|
||||
"anthropic": "claude-haiku-4-5-20251001",
|
||||
"anthropic": "claude-haiku-4-5",
|
||||
"gemini": "gemini-2.5-flash",
|
||||
"groq": "openai/gpt-oss-120b",
|
||||
"minimax": "MiniMax-M2.7",
|
||||
"ollama": "gemma3:12b",
|
||||
"llamacpp": "gemma-4-e2b-it",
|
||||
"lmstudio": "local-model",
|
||||
"vertexai": "google/gemini-2.5-flash-lite",
|
||||
"openai-codex": "gpt-5.2-codex",
|
||||
"claude-code": "claude-sonnet-4-5-20250929",
|
||||
"mock": "mock-model",
|
||||
"none": "none",
|
||||
"litellm": "gpt-4o-mini",
|
||||
"bedrock": "us.amazon.nova-2-lite-v1:0",
|
||||
"volcano": "doubao-pro-32k",
|
||||
"openrouter": "qwen/qwen3.5-9b",
|
||||
}
|
||||
DEFAULT_LLM_MODEL = "gpt-4o-mini" # Fallback if provider not in table
|
||||
# Built-in llama.cpp defaults
|
||||
DEFAULT_LLAMACPP_GPU_LAYERS = -1 # -1 = offload all layers to GPU (Metal/CUDA)
|
||||
DEFAULT_LLAMACPP_CONTEXT_SIZE = 8192
|
||||
DEFAULT_LLAMACPP_CHAT_FORMAT = None # None = auto-detect from GGUF metadata
|
||||
DEFAULT_LLAMACPP_NO_GRAMMAR = False # True = disable JSON grammar enforcement (faster but less reliable)
|
||||
DEFAULT_LLAMACPP_EXTRA_ARGS = None # Space-separated extra CLI args for llama.cpp server
|
||||
|
||||
DEFAULT_LLM_MAX_CONCURRENT = 32
|
||||
DEFAULT_LLM_MAX_RETRIES = 10 # Max retry attempts for LLM API calls
|
||||
DEFAULT_LLM_MAX_RETRIES = 3 # Max retry attempts for LLM API calls
|
||||
DEFAULT_LLM_INITIAL_BACKOFF = 1.0 # Initial backoff in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_MAX_BACKOFF = 60.0 # Max backoff cap in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_TIMEOUT = 120.0 # seconds
|
||||
@@ -384,6 +470,9 @@ DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5"
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU = False # Force CPU mode for local embeddings (avoids MPS/XPC issues on macOS)
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE = False # Security: disabled by default, required for some models
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL = "text-embedding-3-small"
|
||||
DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE = 100
|
||||
DEFAULT_EMBEDDINGS_GEMINI_MODEL = "gemini-embedding-001"
|
||||
DEFAULT_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY = 768
|
||||
DEFAULT_EMBEDDING_DIMENSION = 384
|
||||
|
||||
DEFAULT_RERANKER_PROVIDER = "local"
|
||||
@@ -398,6 +487,7 @@ DEFAULT_RERANKER_LOCAL_BUCKET_BATCHING = False # Length-sorted bucket batching:
|
||||
DEFAULT_RERANKER_LOCAL_BATCH_SIZE = 32 # Batch size for local reranker predict() calls
|
||||
DEFAULT_RERANKER_TEI_BATCH_SIZE = 128
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT = 8
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT = 30.0 # HTTP timeout for TEI reranker requests (seconds)
|
||||
DEFAULT_RERANKER_MAX_CANDIDATES = 300
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL = "ms-marco-MiniLM-L-12-v2" # Best balance of speed and quality
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory
|
||||
@@ -405,8 +495,17 @@ DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL = "embed-english-v3.0"
|
||||
DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
|
||||
|
||||
# OpenRouter defaults
|
||||
DEFAULT_EMBEDDINGS_OPENROUTER_MODEL = "perplexity/pplx-embed-v1-0.6b"
|
||||
DEFAULT_RERANKER_OPENROUTER_MODEL = "cohere/rerank-v3.5"
|
||||
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL = "zerank-2"
|
||||
|
||||
DEFAULT_RERANKER_SILICONFLOW_MODEL = "BAAI/bge-reranker-v2-m3"
|
||||
DEFAULT_RERANKER_SILICONFLOW_BASE_URL = "https://api.siliconflow.cn/v1"
|
||||
|
||||
DEFAULT_RERANKER_GOOGLE_MODEL = "semantic-ranker-default-004"
|
||||
|
||||
# Vector extension (pgvector, vchord, or pgvectorscale)
|
||||
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord", "pgvectorscale"
|
||||
|
||||
@@ -421,6 +520,7 @@ DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC: int | None = None
|
||||
|
||||
# LiteLLM SDK defaults
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL = "cohere/embed-english-v3.0"
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT = "float"
|
||||
DEFAULT_RERANKER_LITELLM_SDK_MODEL = "cohere/rerank-english-v3.0"
|
||||
|
||||
DEFAULT_HOST = "0.0.0.0"
|
||||
@@ -431,13 +531,16 @@ DEFAULT_LOG_FORMAT = "text" # Options: "text", "json"
|
||||
DEFAULT_WORKERS = 1
|
||||
DEFAULT_MCP_ENABLED = True
|
||||
DEFAULT_MCP_ENABLED_TOOLS: list[str] | None = None # None = all tools enabled
|
||||
DEFAULT_MCP_STATELESS = False # False = stateful (supports SSE/GET); True = stateless (POST-only)
|
||||
DEFAULT_ENABLE_BANK_CONFIG_API = True
|
||||
DEFAULT_GRAPH_RETRIEVER = "link_expansion" # Options: "link_expansion", "mpfp", "bfs"
|
||||
DEFAULT_MPFP_TOP_K_NEIGHBORS = 20 # Fan-out limit per node in MPFP graph traversal
|
||||
DEFAULT_DEFAULT_BANK_TEMPLATE: dict | None = None # BankTemplateManifest dict applied to newly-created banks
|
||||
DEFAULT_GRAPH_RETRIEVER = "link_expansion"
|
||||
DEFAULT_RECALL_MAX_CONCURRENT = 32 # Max concurrent recall operations per worker
|
||||
DEFAULT_RECALL_CONNECTION_BUDGET = 4 # Max concurrent DB connections per recall operation
|
||||
DEFAULT_RECALL_MAX_QUERY_TOKENS = 500 # Maximum tokens allowed in recall query
|
||||
DEFAULT_MENTAL_MODEL_REFRESH_CONCURRENCY = 8 # Max concurrent mental model refreshes
|
||||
DEFAULT_LINK_EXPANSION_PER_ENTITY_LIMIT = 200 # Max target units per entity in graph expansion
|
||||
DEFAULT_LINK_EXPANSION_TIMEOUT = 10.0 # Timeout (seconds) for entity expansion query
|
||||
|
||||
# Retain settings
|
||||
DEFAULT_RETAIN_MAX_COMPLETION_TOKENS = 64000 # Max tokens for fact extraction LLM call
|
||||
@@ -449,6 +552,9 @@ DEFAULT_RETAIN_MISSION = None # Declarative spec of what to retain (injected in
|
||||
DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS = None # Custom extraction guidelines (only used when mode="custom")
|
||||
DEFAULT_RETAIN_DEFAULT_STRATEGY = None # Default strategy name (None = no strategy override)
|
||||
DEFAULT_RETAIN_STRATEGIES: dict | None = None # Named retain strategies (dict of name → config overrides)
|
||||
DEFAULT_RETAIN_CHUNK_BATCH_SIZE = (
|
||||
100 # Max chunks per streaming batch. Each chunk produces ~17 facts, so 100 chunks = ~1700 facts/batch.
|
||||
)
|
||||
DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch for async retain auto-splitting
|
||||
DEFAULT_RETAIN_ENTITY_LOOKUP = "trigram" # "full" or "trigram"
|
||||
DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True)
|
||||
@@ -467,7 +573,11 @@ DEFAULT_FILE_DELETE_AFTER_RETAIN = True # Delete file bytes after retain (saves
|
||||
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
|
||||
DEFAULT_ENABLE_OBSERVATION_HISTORY = True # Observation history tracking enabled by default
|
||||
DEFAULT_ENABLE_MENTAL_MODEL_HISTORY = True # Mental model history tracking enabled by default
|
||||
DEFAULT_CONSOLIDATION_MAX_ATTEMPTS = 3 # Outer retry attempts for consolidation LLM batch calls
|
||||
DEFAULT_CONSOLIDATION_BATCH_SIZE = 50 # Memories to load per batch (internal memory optimization)
|
||||
DEFAULT_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = (
|
||||
100 # Max memories per consolidation round (0 = unlimited). Limits how long one bank holds a worker slot.
|
||||
)
|
||||
DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE = 8 # Facts per LLM call (1 = no batching; >1 = batch mode)
|
||||
DEFAULT_CONSOLIDATION_MAX_TOKENS = 512 # Max tokens for recall when finding related observations
|
||||
DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = (
|
||||
@@ -477,6 +587,7 @@ DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION = (
|
||||
256 # Max tokens of source facts per observation in consolidation prompt (-1 = unlimited)
|
||||
)
|
||||
DEFAULT_OBSERVATIONS_MISSION = None # Declarative spec of what observations are for this bank
|
||||
DEFAULT_MAX_OBSERVATIONS_PER_SCOPE = -1 # Max observations per tag scope (-1 = unlimited)
|
||||
|
||||
# Database migrations
|
||||
DEFAULT_RUN_MIGRATIONS_ON_STARTUP = True
|
||||
@@ -486,6 +597,7 @@ DEFAULT_DB_POOL_MIN_SIZE = 5
|
||||
DEFAULT_DB_POOL_MAX_SIZE = 100
|
||||
DEFAULT_DB_COMMAND_TIMEOUT = 60 # seconds
|
||||
DEFAULT_DB_ACQUIRE_TIMEOUT = 30 # seconds
|
||||
DEFAULT_DB_STATEMENT_TIMEOUT = 600 # seconds (Postgres statement_timeout applied on every pool connection; 0 disables)
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode)
|
||||
@@ -495,10 +607,32 @@ DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed
|
||||
DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health
|
||||
DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker
|
||||
DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks per worker
|
||||
DEFAULT_RETAIN_MAX_CONCURRENT = 4 # Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention.
|
||||
|
||||
# Reflect agent settings
|
||||
DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response
|
||||
DEFAULT_REFLECT_MAX_CONTEXT_TOKENS = 100_000 # Max accumulated context tokens before forcing final prompt
|
||||
DEFAULT_REFLECT_WALL_TIMEOUT = 300 # Wall-clock timeout in seconds for the entire reflect operation (5 minutes)
|
||||
DEFAULT_REFLECT_SOURCE_FACTS_MAX_TOKENS = -1 # Token budget for source facts in search_observations (-1 = disabled)
|
||||
DEFAULT_RECALL_INCLUDE_CHUNKS = True # Whether internal recall (e.g. mental model refresh) returns raw chunks
|
||||
DEFAULT_RECALL_MAX_TOKENS = 2048 # Token budget for facts returned by internal recall
|
||||
DEFAULT_RECALL_CHUNKS_MAX_TOKENS = 1000 # Token budget for raw chunks returned by internal recall
|
||||
|
||||
# Recall budget mapping
|
||||
# "fixed": thinking_budget = recall_budget_fixed_<level> (preserves legacy behavior)
|
||||
# "adaptive": thinking_budget = round(max_tokens * recall_budget_adaptive_<level>),
|
||||
# clamped to [recall_budget_min, recall_budget_max]
|
||||
RECALL_BUDGET_FUNCTIONS = ("fixed", "adaptive")
|
||||
DEFAULT_RECALL_BUDGET_FUNCTION = "fixed"
|
||||
DEFAULT_RECALL_BUDGET_FIXED_LOW = 100
|
||||
DEFAULT_RECALL_BUDGET_FIXED_MID = 300
|
||||
DEFAULT_RECALL_BUDGET_FIXED_HIGH = 1000
|
||||
# Adaptive defaults chosen to roughly match fixed defaults at max_tokens=4096
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_LOW = 0.025
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_MID = 0.075
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_HIGH = 0.25
|
||||
DEFAULT_RECALL_BUDGET_MIN = 20 # Floor for the adaptive function
|
||||
DEFAULT_RECALL_BUDGET_MAX = 2000 # Ceiling for the adaptive function
|
||||
|
||||
# Disposition defaults (None = not set, fall back to bank DB value or 3)
|
||||
DEFAULT_DISPOSITION_SKEPTICISM = None
|
||||
@@ -509,6 +643,12 @@ DEFAULT_DISPOSITION_EMPATHY = None
|
||||
DEFAULT_OTEL_TRACES_ENABLED = False # Disabled by default for backward compatibility
|
||||
DEFAULT_OTEL_SERVICE_NAME = "hindsight-api"
|
||||
DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT = "development"
|
||||
DEFAULT_METRICS_INCLUDE_BANK_ID = False # Disabled by default to avoid high-cardinality OTel metric growth
|
||||
|
||||
# Audit log defaults
|
||||
DEFAULT_AUDIT_LOG_ENABLED = False # Disabled by default
|
||||
DEFAULT_AUDIT_LOG_ACTIONS = "" # Empty = audit all eligible actions
|
||||
DEFAULT_AUDIT_LOG_RETENTION_DAYS = -1 # -1 = keep forever
|
||||
|
||||
# Default MCP tool descriptions (can be customized via env vars)
|
||||
DEFAULT_MCP_RETAIN_DESCRIPTION = """Store important information to long-term memory.
|
||||
@@ -555,6 +695,10 @@ class JsonFormatter(logging.Formatter):
|
||||
logging.CRITICAL: "CRITICAL",
|
||||
}
|
||||
|
||||
def __init__(self, allowed_fields: frozenset[str] | None = None):
|
||||
super().__init__()
|
||||
self._allowed_fields = allowed_fields
|
||||
|
||||
def format(self, record: logging.LogRecord) -> str:
|
||||
log_entry = {
|
||||
"severity": self.SEVERITY_MAP.get(record.levelno, "DEFAULT"),
|
||||
@@ -563,10 +707,20 @@ class JsonFormatter(logging.Formatter):
|
||||
"logger": record.name,
|
||||
}
|
||||
|
||||
# Lazy import to avoid circular dependency (engine imports from config).
|
||||
from hindsight_api.engine.memory_engine import _current_schema
|
||||
|
||||
tenant = _current_schema.get()
|
||||
if tenant:
|
||||
log_entry["tenant"] = tenant
|
||||
|
||||
# Add exception info if present
|
||||
if record.exc_info:
|
||||
log_entry["exception"] = self.formatException(record.exc_info)
|
||||
|
||||
if self._allowed_fields is not None:
|
||||
log_entry = {k: v for k, v in log_entry.items() if k in self._allowed_fields}
|
||||
|
||||
return json.dumps(log_entry)
|
||||
|
||||
|
||||
@@ -575,6 +729,25 @@ def _parse_str_list(value: str) -> list[str]:
|
||||
return [v.strip() for v in value.split(",") if v.strip()]
|
||||
|
||||
|
||||
def _parse_positive_int(name: str, raw: str | None, default: int) -> int:
|
||||
"""
|
||||
Parse an env var that must be a positive integer (>= 1).
|
||||
|
||||
Falls back to ``default`` when unset/empty. Raises ValueError on non-integer
|
||||
or non-positive values so misconfiguration fails fast instead of triggering
|
||||
infinite loops or zero-step range() calls downstream.
|
||||
"""
|
||||
if raw is None or raw == "":
|
||||
return default
|
||||
try:
|
||||
parsed = int(raw)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"{name} must be an integer, got {raw!r}") from e
|
||||
if parsed < 1:
|
||||
raise ValueError(f"{name} must be >= 1, got {parsed}")
|
||||
return parsed
|
||||
|
||||
|
||||
def _validate_extraction_mode(mode: str) -> str:
|
||||
"""Validate and normalize extraction mode."""
|
||||
mode_lower = mode.lower()
|
||||
@@ -587,17 +760,50 @@ def _validate_extraction_mode(mode: str) -> str:
|
||||
return mode_lower
|
||||
|
||||
|
||||
def _validate_recall_budget_function(function: str) -> str:
|
||||
"""Validate and normalize recall budget function."""
|
||||
function_lower = function.lower()
|
||||
if function_lower not in RECALL_BUDGET_FUNCTIONS:
|
||||
logger.warning(
|
||||
f"Invalid recall budget function '{function}', must be one of {RECALL_BUDGET_FUNCTIONS}. "
|
||||
f"Defaulting to '{DEFAULT_RECALL_BUDGET_FUNCTION}'."
|
||||
)
|
||||
return DEFAULT_RECALL_BUDGET_FUNCTION
|
||||
return function_lower
|
||||
|
||||
|
||||
def _get_default_model_for_provider(provider: str) -> str:
|
||||
"""Get the default model for a given provider."""
|
||||
return PROVIDER_DEFAULT_MODELS.get(provider.lower(), DEFAULT_LLM_MODEL)
|
||||
|
||||
|
||||
def _parse_default_bank_template(raw: str | None) -> dict | None:
|
||||
"""
|
||||
Parse HINDSIGHT_API_DEFAULT_BANK_TEMPLATE as JSON.
|
||||
|
||||
The env var holds a BankTemplateManifest (JSON object) applied verbatim to
|
||||
every newly-created bank. Full Pydantic validation is deferred to bank
|
||||
creation time (to avoid pulling API models into config.py), but we fail
|
||||
fast here if the value is not valid JSON or not a JSON object.
|
||||
"""
|
||||
if raw is None or raw.strip() == "":
|
||||
return DEFAULT_DEFAULT_BANK_TEMPLATE
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"Invalid {ENV_DEFAULT_BANK_TEMPLATE}: expected a JSON object, got invalid JSON: {e}") from e
|
||||
if not isinstance(parsed, dict):
|
||||
raise ValueError(f"Invalid {ENV_DEFAULT_BANK_TEMPLATE}: expected a JSON object, got {type(parsed).__name__}")
|
||||
return parsed
|
||||
|
||||
|
||||
@dataclass
|
||||
class HindsightConfig:
|
||||
"""Configuration container for Hindsight API."""
|
||||
|
||||
# Database
|
||||
database_url: str
|
||||
migration_database_url: str | None
|
||||
database_schema: str
|
||||
vector_extension: str # "pgvector" or "vchord"
|
||||
text_search_extension: str # "native" or "vchord"
|
||||
@@ -614,6 +820,9 @@ class HindsightConfig:
|
||||
llm_timeout: float
|
||||
llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto"
|
||||
llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper)
|
||||
llm_extra_body: (
|
||||
dict | None
|
||||
) # Extra body params merged into OpenAI-compatible API calls (e.g. {"chat_template_kwargs": {"enable_thinking": true}})
|
||||
|
||||
# Vertex AI configuration
|
||||
llm_vertexai_project_id: str | None
|
||||
@@ -623,6 +832,14 @@ class HindsightConfig:
|
||||
# Gemini safety settings (None = use Gemini defaults; list of dicts with category/threshold)
|
||||
llm_gemini_safety_settings: list | None
|
||||
|
||||
# Built-in llama.cpp configuration (for provider=llamacpp)
|
||||
llamacpp_model_path: str | None # Path to GGUF file (None = auto-download default)
|
||||
llamacpp_gpu_layers: int # -1 = all layers on GPU, 0 = CPU only
|
||||
llamacpp_context_size: int # Context window size
|
||||
llamacpp_chat_format: str | None # Chat template format (None = auto-detect from GGUF)
|
||||
llamacpp_no_grammar: bool # Disable JSON grammar enforcement (faster, less reliable)
|
||||
llamacpp_extra_args: str | None # Space-separated extra CLI args for llama.cpp server
|
||||
|
||||
# Per-operation LLM configuration (None = use default LLM config)
|
||||
retain_llm_provider: str | None
|
||||
retain_llm_api_key: str | None
|
||||
@@ -664,12 +881,23 @@ class HindsightConfig:
|
||||
embeddings_cohere_api_key: str | None
|
||||
embeddings_cohere_model: str
|
||||
embeddings_cohere_base_url: str | None
|
||||
embeddings_openrouter_api_key: str | None
|
||||
embeddings_openrouter_model: str
|
||||
embeddings_litellm_api_base: str
|
||||
embeddings_litellm_api_key: str | None
|
||||
embeddings_litellm_model: str
|
||||
embeddings_litellm_sdk_api_key: str | None
|
||||
embeddings_litellm_sdk_model: str
|
||||
embeddings_litellm_sdk_api_base: str | None
|
||||
embeddings_litellm_sdk_output_dimensions: int | None
|
||||
embeddings_litellm_sdk_encoding_format: str | None
|
||||
# Gemini/Vertex AI embeddings
|
||||
embeddings_gemini_api_key: str | None
|
||||
embeddings_gemini_model: str
|
||||
embeddings_gemini_output_dimensionality: int | None
|
||||
embeddings_vertexai_project_id: str | None
|
||||
embeddings_vertexai_region: str | None
|
||||
embeddings_vertexai_service_account_key: str | None
|
||||
|
||||
# Reranker
|
||||
reranker_provider: str
|
||||
@@ -683,10 +911,13 @@ class HindsightConfig:
|
||||
reranker_tei_url: str | None
|
||||
reranker_tei_batch_size: int
|
||||
reranker_tei_max_concurrent: int
|
||||
reranker_tei_http_timeout: float
|
||||
reranker_max_candidates: int
|
||||
reranker_cohere_api_key: str | None
|
||||
reranker_cohere_model: str
|
||||
reranker_cohere_base_url: str | None
|
||||
reranker_openrouter_api_key: str | None
|
||||
reranker_openrouter_model: str
|
||||
reranker_litellm_api_base: str
|
||||
reranker_litellm_api_key: str | None
|
||||
reranker_litellm_model: str
|
||||
@@ -696,6 +927,13 @@ class HindsightConfig:
|
||||
reranker_litellm_sdk_api_base: str | None
|
||||
reranker_zeroentropy_api_key: str | None
|
||||
reranker_zeroentropy_model: str
|
||||
reranker_zeroentropy_base_url: str | None
|
||||
reranker_siliconflow_api_key: str | None
|
||||
reranker_siliconflow_model: str
|
||||
reranker_siliconflow_base_url: str
|
||||
reranker_google_model: str
|
||||
reranker_google_project_id: str | None
|
||||
reranker_google_service_account_key: str | None
|
||||
|
||||
# Server
|
||||
host: str
|
||||
@@ -703,17 +941,23 @@ class HindsightConfig:
|
||||
base_path: str
|
||||
log_level: str
|
||||
log_format: str
|
||||
log_json_fields: list[str] | None # None = all fields; explicit list = allowlist
|
||||
mcp_enabled: bool
|
||||
mcp_enabled_tools: list[str] | None # None = all tools; explicit list = allowlist
|
||||
mcp_stateless: bool # True = stateless HTTP (POST-only); False = stateful (supports GET/SSE)
|
||||
enable_bank_config_api: bool
|
||||
# Default bank template (static, server-level only). When set, the manifest is applied
|
||||
# to every newly-created bank, overriding the env/config defaults for any fields it sets.
|
||||
default_bank_template: dict | None
|
||||
|
||||
# Recall
|
||||
graph_retriever: str
|
||||
mpfp_top_k_neighbors: int
|
||||
recall_max_concurrent: int
|
||||
recall_connection_budget: int
|
||||
recall_max_query_tokens: int
|
||||
mental_model_refresh_concurrency: int
|
||||
link_expansion_per_entity_limit: int
|
||||
link_expansion_timeout: float
|
||||
|
||||
# Retain settings
|
||||
retain_max_completion_tokens: int
|
||||
@@ -728,6 +972,7 @@ class HindsightConfig:
|
||||
retain_batch_enabled: bool
|
||||
retain_batch_poll_interval_seconds: int
|
||||
retain_entity_lookup: str # "full" or "trigram"
|
||||
retain_chunk_batch_size: int # Max chunks per streaming batch (0 = disabled)
|
||||
|
||||
# File storage (static - server-level only)
|
||||
file_storage_type: str # "native" (PostgreSQL) or "s3" (S3-compatible)
|
||||
@@ -755,11 +1000,14 @@ class HindsightConfig:
|
||||
enable_observation_history: bool
|
||||
enable_mental_model_history: bool
|
||||
consolidation_batch_size: int
|
||||
consolidation_max_memories_per_round: int
|
||||
consolidation_llm_batch_size: int
|
||||
consolidation_max_tokens: int
|
||||
consolidation_source_facts_max_tokens: int
|
||||
consolidation_source_facts_max_tokens_per_observation: int
|
||||
consolidation_max_attempts: int
|
||||
observations_mission: str | None
|
||||
max_observations_per_scope: int
|
||||
|
||||
# Entity labels (controlled vocabulary of key:value classification labels extracted at retain time)
|
||||
# List of label group dicts: [{key, description, type, optional, values: [{value, description}]}]
|
||||
@@ -770,6 +1018,26 @@ class HindsightConfig:
|
||||
|
||||
# Reflect agent settings
|
||||
reflect_mission: str | None
|
||||
reflect_source_facts_max_tokens: int
|
||||
|
||||
# Recall settings (used by internal recall, e.g. during mental model refresh)
|
||||
recall_include_chunks: bool
|
||||
recall_max_tokens: int
|
||||
recall_chunks_max_tokens: int
|
||||
|
||||
# Recall budget mapping: how the Budget enum (LOW/MID/HIGH) maps to thinking_budget integer.
|
||||
# function="fixed": use the recall_budget_fixed_* values directly (legacy behavior).
|
||||
# function="adaptive": compute round(max_tokens * recall_budget_adaptive_*),
|
||||
# clamped to [recall_budget_min, recall_budget_max].
|
||||
recall_budget_function: str
|
||||
recall_budget_fixed_low: int
|
||||
recall_budget_fixed_mid: int
|
||||
recall_budget_fixed_high: int
|
||||
recall_budget_adaptive_low: float
|
||||
recall_budget_adaptive_mid: float
|
||||
recall_budget_adaptive_high: float
|
||||
recall_budget_min: int
|
||||
recall_budget_max: int
|
||||
|
||||
# Disposition settings (hierarchical - can be overridden per bank; None = fall back to DB)
|
||||
disposition_skepticism: int | None
|
||||
@@ -788,6 +1056,7 @@ class HindsightConfig:
|
||||
db_pool_max_size: int
|
||||
db_command_timeout: int
|
||||
db_acquire_timeout: int
|
||||
db_statement_timeout: int
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
worker_enabled: bool
|
||||
@@ -797,10 +1066,12 @@ class HindsightConfig:
|
||||
worker_http_port: int
|
||||
worker_max_slots: int
|
||||
worker_consolidation_max_slots: int
|
||||
retain_max_concurrent: int
|
||||
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations: int
|
||||
reflect_max_context_tokens: int
|
||||
reflect_wall_timeout: int
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
otel_traces_enabled: bool
|
||||
@@ -808,6 +1079,12 @@ class HindsightConfig:
|
||||
otel_exporter_otlp_headers: str | None
|
||||
otel_service_name: str
|
||||
otel_deployment_environment: str
|
||||
metrics_include_bank_id: bool
|
||||
|
||||
# Audit log configuration (static - server-level only)
|
||||
audit_log_enabled: bool # Master switch for audit logging
|
||||
audit_log_actions: list[str] # Allowlist of action types (empty = all)
|
||||
audit_log_retention_days: int # -1 = keep forever, >0 = delete after N days
|
||||
|
||||
# Webhook configuration (static - server-level only, not per-bank)
|
||||
webhook_url: str | None # Global webhook URL (None = disabled)
|
||||
@@ -815,6 +1092,10 @@ class HindsightConfig:
|
||||
webhook_event_types: list[str] # Event types to deliver globally
|
||||
webhook_delivery_poll_interval_seconds: int # How often the delivery worker polls
|
||||
|
||||
# Defaulted fields (source-compatible additions — existing direct constructor callers keep working).
|
||||
# Keep at the end of the dataclass; Python forbids non-default fields after default fields.
|
||||
embeddings_openai_batch_size: int = DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
# CREDENTIAL_FIELDS: Never exposed via API, never configurable per-tenant/bank
|
||||
@@ -832,8 +1113,14 @@ class HindsightConfig:
|
||||
"embeddings_tei_base_url",
|
||||
"reranker_tei_base_url",
|
||||
"reranker_cohere_base_url",
|
||||
"reranker_zeroentropy_base_url",
|
||||
"reranker_siliconflow_base_url",
|
||||
# Service Account Keys
|
||||
"llm_vertexai_service_account_key",
|
||||
"embeddings_vertexai_service_account_key",
|
||||
"reranker_google_service_account_key",
|
||||
# Embeddings API keys
|
||||
"embeddings_gemini_api_key",
|
||||
# File storage credentials
|
||||
"file_storage_s3_access_key_id",
|
||||
"file_storage_s3_secret_access_key",
|
||||
@@ -856,17 +1143,35 @@ class HindsightConfig:
|
||||
"retain_custom_instructions",
|
||||
"retain_default_strategy",
|
||||
"retain_strategies",
|
||||
"retain_chunk_batch_size",
|
||||
# Entity labels (controlled vocabulary for entity classification)
|
||||
"entity_labels",
|
||||
"entities_allow_free_form",
|
||||
# Consolidation settings
|
||||
"enable_observations",
|
||||
"consolidation_llm_batch_size",
|
||||
"consolidation_max_memories_per_round",
|
||||
"consolidation_source_facts_max_tokens",
|
||||
"consolidation_source_facts_max_tokens_per_observation",
|
||||
"observations_mission",
|
||||
"max_observations_per_scope",
|
||||
# Reflect settings
|
||||
"reflect_mission",
|
||||
"reflect_source_facts_max_tokens",
|
||||
# Recall settings (used by internal recall, e.g. mental model refresh)
|
||||
"recall_include_chunks",
|
||||
"recall_max_tokens",
|
||||
"recall_chunks_max_tokens",
|
||||
# Recall budget mapping (Budget enum -> thinking_budget integer)
|
||||
"recall_budget_function",
|
||||
"recall_budget_fixed_low",
|
||||
"recall_budget_fixed_mid",
|
||||
"recall_budget_fixed_high",
|
||||
"recall_budget_adaptive_low",
|
||||
"recall_budget_adaptive_mid",
|
||||
"recall_budget_adaptive_high",
|
||||
"recall_budget_min",
|
||||
"recall_budget_max",
|
||||
# Disposition settings
|
||||
"disposition_skepticism",
|
||||
"disposition_literalism",
|
||||
@@ -949,9 +1254,19 @@ class HindsightConfig:
|
||||
f"Invalid text_search_extension: {self.text_search_extension}. Must be one of: {', '.join(valid_text_search)}"
|
||||
)
|
||||
|
||||
# When LLM provider is "none", force chunks-only mode and disable LLM-dependent features
|
||||
if self.llm_provider == "none":
|
||||
self.retain_extraction_mode = "chunks"
|
||||
self.enable_observations = False
|
||||
logger.info(
|
||||
"LLM provider set to 'none': forcing retain_extraction_mode='chunks', "
|
||||
"disabling observations/consolidation. Reflect will return HTTP 400."
|
||||
)
|
||||
|
||||
# RETAIN_MAX_COMPLETION_TOKENS must be greater than RETAIN_CHUNK_SIZE
|
||||
# to ensure the LLM has enough output capacity to extract facts from chunks
|
||||
if self.retain_max_completion_tokens <= self.retain_chunk_size:
|
||||
# (not applicable when provider is "none" since no LLM calls are made)
|
||||
if self.llm_provider != "none" and self.retain_max_completion_tokens <= self.retain_chunk_size:
|
||||
raise ValueError(
|
||||
f"Invalid configuration: HINDSIGHT_API_RETAIN_MAX_COMPLETION_TOKENS "
|
||||
f"({self.retain_max_completion_tokens}) must be greater than "
|
||||
@@ -973,6 +1288,7 @@ class HindsightConfig:
|
||||
config = cls(
|
||||
# Database
|
||||
database_url=os.getenv(ENV_DATABASE_URL, DEFAULT_DATABASE_URL),
|
||||
migration_database_url=os.getenv(ENV_MIGRATION_DATABASE_URL) or None,
|
||||
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
|
||||
vector_extension=os.getenv(ENV_VECTOR_EXTENSION, DEFAULT_VECTOR_EXTENSION).lower(),
|
||||
text_search_extension=os.getenv(ENV_TEXT_SEARCH_EXTENSION, DEFAULT_TEXT_SEARCH_EXTENSION).lower(),
|
||||
@@ -988,6 +1304,7 @@ class HindsightConfig:
|
||||
llm_timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))),
|
||||
llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
|
||||
llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
|
||||
llm_extra_body=json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null")),
|
||||
# Vertex AI
|
||||
llm_vertexai_project_id=os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID) or DEFAULT_LLM_VERTEXAI_PROJECT_ID,
|
||||
llm_vertexai_region=os.getenv(ENV_LLM_VERTEXAI_REGION, DEFAULT_LLM_VERTEXAI_REGION),
|
||||
@@ -995,6 +1312,14 @@ class HindsightConfig:
|
||||
or DEFAULT_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
|
||||
# Gemini safety settings (JSON-encoded list of {category, threshold} dicts)
|
||||
llm_gemini_safety_settings=json.loads(os.getenv(ENV_LLM_GEMINI_SAFETY_SETTINGS, "null")),
|
||||
# Built-in llama.cpp configuration
|
||||
llamacpp_model_path=os.getenv(ENV_LLAMACPP_MODEL_PATH) or None,
|
||||
llamacpp_gpu_layers=int(os.getenv(ENV_LLAMACPP_GPU_LAYERS, str(DEFAULT_LLAMACPP_GPU_LAYERS))),
|
||||
llamacpp_context_size=int(os.getenv(ENV_LLAMACPP_CONTEXT_SIZE, str(DEFAULT_LLAMACPP_CONTEXT_SIZE))),
|
||||
llamacpp_chat_format=os.getenv(ENV_LLAMACPP_CHAT_FORMAT) or DEFAULT_LLAMACPP_CHAT_FORMAT,
|
||||
llamacpp_no_grammar=os.getenv(ENV_LLAMACPP_NO_GRAMMAR, str(DEFAULT_LLAMACPP_NO_GRAMMAR)).lower()
|
||||
in ("true", "1"),
|
||||
llamacpp_extra_args=os.getenv(ENV_LLAMACPP_EXTRA_ARGS) or DEFAULT_LLAMACPP_EXTRA_ARGS,
|
||||
# Per-operation LLM config (None = use default)
|
||||
retain_llm_provider=os.getenv(ENV_RETAIN_LLM_PROVIDER) or None,
|
||||
retain_llm_api_key=os.getenv(ENV_RETAIN_LLM_API_KEY) or None,
|
||||
@@ -1079,10 +1404,20 @@ class HindsightConfig:
|
||||
in ("true", "1"),
|
||||
embeddings_tei_url=os.getenv(ENV_EMBEDDINGS_TEI_URL),
|
||||
embeddings_openai_base_url=os.getenv(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None,
|
||||
embeddings_openai_batch_size=_parse_positive_int(
|
||||
ENV_EMBEDDINGS_OPENAI_BATCH_SIZE,
|
||||
os.getenv(ENV_EMBEDDINGS_OPENAI_BATCH_SIZE),
|
||||
DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE,
|
||||
),
|
||||
# Cohere embeddings (with backward-compatible fallback to shared API key)
|
||||
embeddings_cohere_api_key=os.getenv(ENV_EMBEDDINGS_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
embeddings_cohere_model=os.getenv(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL),
|
||||
embeddings_cohere_base_url=os.getenv(ENV_EMBEDDINGS_COHERE_BASE_URL) or None,
|
||||
# OpenRouter embeddings (with fallback to shared OpenRouter key, then LLM key)
|
||||
embeddings_openrouter_api_key=os.getenv(ENV_EMBEDDINGS_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_LLM_API_KEY),
|
||||
embeddings_openrouter_model=os.getenv(ENV_EMBEDDINGS_OPENROUTER_MODEL, DEFAULT_EMBEDDINGS_OPENROUTER_MODEL),
|
||||
# LiteLLM embeddings (with backward-compatible fallback to shared config)
|
||||
embeddings_litellm_api_base=os.getenv(ENV_EMBEDDINGS_LITELLM_API_BASE)
|
||||
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
|
||||
@@ -1094,6 +1429,26 @@ class HindsightConfig:
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_MODEL, DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL
|
||||
),
|
||||
embeddings_litellm_sdk_api_base=os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_API_BASE) or None,
|
||||
embeddings_litellm_sdk_output_dimensions=int(v)
|
||||
if (v := os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS))
|
||||
else None,
|
||||
embeddings_litellm_sdk_encoding_format=os.getenv(
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT, DEFAULT_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT
|
||||
),
|
||||
# Gemini/Vertex AI embeddings (with fallback to LLM keys)
|
||||
embeddings_gemini_api_key=os.getenv(ENV_EMBEDDINGS_GEMINI_API_KEY) or os.getenv(ENV_LLM_API_KEY),
|
||||
embeddings_gemini_model=os.getenv(ENV_EMBEDDINGS_GEMINI_MODEL, DEFAULT_EMBEDDINGS_GEMINI_MODEL),
|
||||
embeddings_gemini_output_dimensionality=int(
|
||||
os.getenv(
|
||||
ENV_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY,
|
||||
str(DEFAULT_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY),
|
||||
)
|
||||
),
|
||||
embeddings_vertexai_project_id=os.getenv(ENV_EMBEDDINGS_VERTEXAI_PROJECT_ID)
|
||||
or os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID),
|
||||
embeddings_vertexai_region=os.getenv(ENV_EMBEDDINGS_VERTEXAI_REGION) or os.getenv(ENV_LLM_VERTEXAI_REGION),
|
||||
embeddings_vertexai_service_account_key=os.getenv(ENV_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY)
|
||||
or os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY),
|
||||
# Reranker
|
||||
reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER),
|
||||
reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL),
|
||||
@@ -1122,11 +1477,19 @@ class HindsightConfig:
|
||||
reranker_tei_max_concurrent=int(
|
||||
os.getenv(ENV_RERANKER_TEI_MAX_CONCURRENT, str(DEFAULT_RERANKER_TEI_MAX_CONCURRENT))
|
||||
),
|
||||
reranker_tei_http_timeout=float(
|
||||
os.getenv(ENV_RERANKER_TEI_HTTP_TIMEOUT, str(DEFAULT_RERANKER_TEI_HTTP_TIMEOUT))
|
||||
),
|
||||
reranker_max_candidates=int(os.getenv(ENV_RERANKER_MAX_CANDIDATES, str(DEFAULT_RERANKER_MAX_CANDIDATES))),
|
||||
# Cohere reranker (with backward-compatible fallback to shared API key)
|
||||
reranker_cohere_api_key=os.getenv(ENV_RERANKER_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
reranker_cohere_model=os.getenv(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL),
|
||||
reranker_cohere_base_url=os.getenv(ENV_RERANKER_COHERE_BASE_URL) or None,
|
||||
# OpenRouter reranker (with fallback to shared OpenRouter key, then LLM key)
|
||||
reranker_openrouter_api_key=os.getenv(ENV_RERANKER_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_LLM_API_KEY),
|
||||
reranker_openrouter_model=os.getenv(ENV_RERANKER_OPENROUTER_MODEL, DEFAULT_RERANKER_OPENROUTER_MODEL),
|
||||
# LiteLLM reranker (with backward-compatible fallback to shared config)
|
||||
reranker_litellm_api_base=os.getenv(ENV_RERANKER_LITELLM_API_BASE)
|
||||
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
|
||||
@@ -1142,21 +1505,36 @@ class HindsightConfig:
|
||||
# ZeroEntropy reranker
|
||||
reranker_zeroentropy_api_key=os.getenv(ENV_RERANKER_ZEROENTROPY_API_KEY),
|
||||
reranker_zeroentropy_model=os.getenv(ENV_RERANKER_ZEROENTROPY_MODEL, DEFAULT_RERANKER_ZEROENTROPY_MODEL),
|
||||
reranker_zeroentropy_base_url=os.getenv(ENV_RERANKER_ZEROENTROPY_BASE_URL) or None,
|
||||
# SiliconFlow reranker (Cohere-compatible /rerank endpoint)
|
||||
reranker_siliconflow_api_key=os.getenv(ENV_RERANKER_SILICONFLOW_API_KEY),
|
||||
reranker_siliconflow_model=os.getenv(ENV_RERANKER_SILICONFLOW_MODEL, DEFAULT_RERANKER_SILICONFLOW_MODEL),
|
||||
reranker_siliconflow_base_url=os.getenv(
|
||||
ENV_RERANKER_SILICONFLOW_BASE_URL, DEFAULT_RERANKER_SILICONFLOW_BASE_URL
|
||||
),
|
||||
# Google Discovery Engine reranker (with fallback to LLM Vertex AI keys)
|
||||
reranker_google_model=os.getenv(ENV_RERANKER_GOOGLE_MODEL, DEFAULT_RERANKER_GOOGLE_MODEL),
|
||||
reranker_google_project_id=os.getenv(ENV_RERANKER_GOOGLE_PROJECT_ID)
|
||||
or os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID),
|
||||
reranker_google_service_account_key=os.getenv(ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY)
|
||||
or os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY),
|
||||
# Server
|
||||
host=os.getenv(ENV_HOST, DEFAULT_HOST),
|
||||
port=int(os.getenv(ENV_PORT, DEFAULT_PORT)),
|
||||
base_path=os.getenv(ENV_BASE_PATH, DEFAULT_BASE_PATH),
|
||||
log_level=os.getenv(ENV_LOG_LEVEL, DEFAULT_LOG_LEVEL),
|
||||
log_format=os.getenv(ENV_LOG_FORMAT, DEFAULT_LOG_FORMAT).lower(),
|
||||
log_json_fields=_parse_str_list(os.getenv(ENV_LOG_JSON_FIELDS, "")) or None,
|
||||
mcp_enabled=os.getenv(ENV_MCP_ENABLED, str(DEFAULT_MCP_ENABLED)).lower() == "true",
|
||||
mcp_enabled_tools=[t.strip() for t in os.getenv(ENV_MCP_ENABLED_TOOLS).split(",") if t.strip()]
|
||||
if os.getenv(ENV_MCP_ENABLED_TOOLS)
|
||||
else DEFAULT_MCP_ENABLED_TOOLS,
|
||||
mcp_stateless=os.getenv(ENV_MCP_STATELESS, str(DEFAULT_MCP_STATELESS)).lower() == "true",
|
||||
enable_bank_config_api=os.getenv(ENV_ENABLE_BANK_CONFIG_API, str(DEFAULT_ENABLE_BANK_CONFIG_API)).lower()
|
||||
== "true",
|
||||
default_bank_template=_parse_default_bank_template(os.getenv(ENV_DEFAULT_BANK_TEMPLATE)),
|
||||
# Recall
|
||||
graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER),
|
||||
mpfp_top_k_neighbors=int(os.getenv(ENV_MPFP_TOP_K_NEIGHBORS, str(DEFAULT_MPFP_TOP_K_NEIGHBORS))),
|
||||
recall_max_concurrent=int(os.getenv(ENV_RECALL_MAX_CONCURRENT, str(DEFAULT_RECALL_MAX_CONCURRENT))),
|
||||
recall_connection_budget=int(
|
||||
os.getenv(ENV_RECALL_CONNECTION_BUDGET, str(DEFAULT_RECALL_CONNECTION_BUDGET))
|
||||
@@ -1165,6 +1543,10 @@ class HindsightConfig:
|
||||
mental_model_refresh_concurrency=int(
|
||||
os.getenv(ENV_MENTAL_MODEL_REFRESH_CONCURRENCY, str(DEFAULT_MENTAL_MODEL_REFRESH_CONCURRENCY))
|
||||
),
|
||||
link_expansion_per_entity_limit=int(
|
||||
os.getenv(ENV_LINK_EXPANSION_PER_ENTITY_LIMIT, str(DEFAULT_LINK_EXPANSION_PER_ENTITY_LIMIT))
|
||||
),
|
||||
link_expansion_timeout=float(os.getenv(ENV_LINK_EXPANSION_TIMEOUT, str(DEFAULT_LINK_EXPANSION_TIMEOUT))),
|
||||
# Optimization flags
|
||||
skip_llm_verification=os.getenv(ENV_SKIP_LLM_VERIFICATION, "false").lower() == "true",
|
||||
lazy_reranker=os.getenv(ENV_LAZY_RERANKER, "false").lower() == "true",
|
||||
@@ -1191,6 +1573,7 @@ class HindsightConfig:
|
||||
retain_batch_poll_interval_seconds=int(
|
||||
os.getenv(ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS, str(DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS))
|
||||
),
|
||||
retain_chunk_batch_size=int(os.getenv(ENV_RETAIN_CHUNK_BATCH_SIZE, str(DEFAULT_RETAIN_CHUNK_BATCH_SIZE))),
|
||||
# File storage
|
||||
file_storage_type=os.getenv(ENV_FILE_STORAGE_TYPE, DEFAULT_FILE_STORAGE_TYPE),
|
||||
file_storage_s3_bucket=os.getenv(ENV_FILE_STORAGE_S3_BUCKET) or None,
|
||||
@@ -1234,6 +1617,12 @@ class HindsightConfig:
|
||||
consolidation_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_BATCH_SIZE))
|
||||
),
|
||||
consolidation_max_memories_per_round=int(
|
||||
os.getenv(
|
||||
ENV_CONSOLIDATION_MAX_MEMORIES_PER_ROUND,
|
||||
str(DEFAULT_CONSOLIDATION_MAX_MEMORIES_PER_ROUND),
|
||||
)
|
||||
),
|
||||
consolidation_llm_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_LLM_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE))
|
||||
),
|
||||
@@ -1249,7 +1638,13 @@ class HindsightConfig:
|
||||
str(DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION),
|
||||
)
|
||||
),
|
||||
consolidation_max_attempts=int(
|
||||
os.getenv(ENV_CONSOLIDATION_MAX_ATTEMPTS, str(DEFAULT_CONSOLIDATION_MAX_ATTEMPTS))
|
||||
),
|
||||
observations_mission=os.getenv(ENV_OBSERVATIONS_MISSION) or DEFAULT_OBSERVATIONS_MISSION,
|
||||
max_observations_per_scope=int(
|
||||
os.getenv(ENV_MAX_OBSERVATIONS_PER_SCOPE, str(DEFAULT_MAX_OBSERVATIONS_PER_SCOPE))
|
||||
),
|
||||
entity_labels=None,
|
||||
entities_allow_free_form=True,
|
||||
# Database migrations
|
||||
@@ -1259,6 +1654,7 @@ class HindsightConfig:
|
||||
db_pool_max_size=int(os.getenv(ENV_DB_POOL_MAX_SIZE, str(DEFAULT_DB_POOL_MAX_SIZE))),
|
||||
db_command_timeout=int(os.getenv(ENV_DB_COMMAND_TIMEOUT, str(DEFAULT_DB_COMMAND_TIMEOUT))),
|
||||
db_acquire_timeout=int(os.getenv(ENV_DB_ACQUIRE_TIMEOUT, str(DEFAULT_DB_ACQUIRE_TIMEOUT))),
|
||||
db_statement_timeout=int(os.getenv(ENV_DB_STATEMENT_TIMEOUT, str(DEFAULT_DB_STATEMENT_TIMEOUT))),
|
||||
# Worker configuration
|
||||
worker_enabled=os.getenv(ENV_WORKER_ENABLED, str(DEFAULT_WORKER_ENABLED)).lower() == "true",
|
||||
worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID,
|
||||
@@ -1269,12 +1665,42 @@ class HindsightConfig:
|
||||
worker_consolidation_max_slots=int(
|
||||
os.getenv(ENV_WORKER_CONSOLIDATION_MAX_SLOTS, str(DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS))
|
||||
),
|
||||
retain_max_concurrent=int(os.getenv(ENV_RETAIN_MAX_CONCURRENT, str(DEFAULT_RETAIN_MAX_CONCURRENT))),
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
|
||||
reflect_max_context_tokens=int(
|
||||
os.getenv(ENV_REFLECT_MAX_CONTEXT_TOKENS, str(DEFAULT_REFLECT_MAX_CONTEXT_TOKENS))
|
||||
),
|
||||
reflect_wall_timeout=int(os.getenv(ENV_REFLECT_WALL_TIMEOUT, str(DEFAULT_REFLECT_WALL_TIMEOUT))),
|
||||
reflect_mission=os.getenv(ENV_REFLECT_MISSION) or None,
|
||||
reflect_source_facts_max_tokens=int(
|
||||
os.getenv(ENV_REFLECT_SOURCE_FACTS_MAX_TOKENS, str(DEFAULT_REFLECT_SOURCE_FACTS_MAX_TOKENS))
|
||||
),
|
||||
recall_include_chunks=os.getenv(ENV_RECALL_INCLUDE_CHUNKS, str(DEFAULT_RECALL_INCLUDE_CHUNKS)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
recall_max_tokens=int(os.getenv(ENV_RECALL_MAX_TOKENS, str(DEFAULT_RECALL_MAX_TOKENS))),
|
||||
recall_chunks_max_tokens=int(
|
||||
os.getenv(ENV_RECALL_CHUNKS_MAX_TOKENS, str(DEFAULT_RECALL_CHUNKS_MAX_TOKENS))
|
||||
),
|
||||
recall_budget_function=_validate_recall_budget_function(
|
||||
os.getenv(ENV_RECALL_BUDGET_FUNCTION, DEFAULT_RECALL_BUDGET_FUNCTION)
|
||||
),
|
||||
recall_budget_fixed_low=int(os.getenv(ENV_RECALL_BUDGET_FIXED_LOW, str(DEFAULT_RECALL_BUDGET_FIXED_LOW))),
|
||||
recall_budget_fixed_mid=int(os.getenv(ENV_RECALL_BUDGET_FIXED_MID, str(DEFAULT_RECALL_BUDGET_FIXED_MID))),
|
||||
recall_budget_fixed_high=int(
|
||||
os.getenv(ENV_RECALL_BUDGET_FIXED_HIGH, str(DEFAULT_RECALL_BUDGET_FIXED_HIGH))
|
||||
),
|
||||
recall_budget_adaptive_low=float(
|
||||
os.getenv(ENV_RECALL_BUDGET_ADAPTIVE_LOW, str(DEFAULT_RECALL_BUDGET_ADAPTIVE_LOW))
|
||||
),
|
||||
recall_budget_adaptive_mid=float(
|
||||
os.getenv(ENV_RECALL_BUDGET_ADAPTIVE_MID, str(DEFAULT_RECALL_BUDGET_ADAPTIVE_MID))
|
||||
),
|
||||
recall_budget_adaptive_high=float(
|
||||
os.getenv(ENV_RECALL_BUDGET_ADAPTIVE_HIGH, str(DEFAULT_RECALL_BUDGET_ADAPTIVE_HIGH))
|
||||
),
|
||||
recall_budget_min=int(os.getenv(ENV_RECALL_BUDGET_MIN, str(DEFAULT_RECALL_BUDGET_MIN))),
|
||||
recall_budget_max=int(os.getenv(ENV_RECALL_BUDGET_MAX, str(DEFAULT_RECALL_BUDGET_MAX))),
|
||||
# Disposition settings (None = fall back to DB value)
|
||||
disposition_skepticism=int(os.getenv(ENV_DISPOSITION_SKEPTICISM))
|
||||
if os.getenv(ENV_DISPOSITION_SKEPTICISM)
|
||||
@@ -1292,6 +1718,16 @@ class HindsightConfig:
|
||||
otel_exporter_otlp_headers=os.getenv(ENV_OTEL_EXPORTER_OTLP_HEADERS) or None,
|
||||
otel_service_name=os.getenv(ENV_OTEL_SERVICE_NAME, DEFAULT_OTEL_SERVICE_NAME),
|
||||
otel_deployment_environment=os.getenv(ENV_OTEL_DEPLOYMENT_ENVIRONMENT, DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT),
|
||||
metrics_include_bank_id=os.getenv(ENV_METRICS_INCLUDE_BANK_ID, str(DEFAULT_METRICS_INCLUDE_BANK_ID)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
# Audit log configuration (static, server-level only)
|
||||
audit_log_enabled=os.getenv(ENV_AUDIT_LOG_ENABLED, str(DEFAULT_AUDIT_LOG_ENABLED)).lower() == "true",
|
||||
audit_log_actions=[
|
||||
a.strip() for a in os.getenv(ENV_AUDIT_LOG_ACTIONS, DEFAULT_AUDIT_LOG_ACTIONS).split(",") if a.strip()
|
||||
],
|
||||
audit_log_retention_days=int(
|
||||
os.getenv(ENV_AUDIT_LOG_RETENTION_DAYS, str(DEFAULT_AUDIT_LOG_RETENTION_DAYS))
|
||||
),
|
||||
# Webhook configuration (static, server-level only)
|
||||
webhook_url=os.getenv(ENV_WEBHOOK_URL) or DEFAULT_WEBHOOK_URL,
|
||||
webhook_secret=os.getenv(ENV_WEBHOOK_SECRET) or DEFAULT_WEBHOOK_SECRET,
|
||||
@@ -1355,15 +1791,21 @@ class HindsightConfig:
|
||||
handler.setLevel(self.get_python_log_level())
|
||||
|
||||
if self.log_format == "json":
|
||||
handler.setFormatter(JsonFormatter())
|
||||
allowed = frozenset(self.log_json_fields) if self.log_json_fields else None
|
||||
handler.setFormatter(JsonFormatter(allowed_fields=allowed))
|
||||
else:
|
||||
handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s"))
|
||||
|
||||
root_logger.addHandler(handler)
|
||||
|
||||
# Silence noisy third-party loggers
|
||||
logging.getLogger("google_genai.models").setLevel(logging.WARNING)
|
||||
|
||||
def log_config(self) -> None:
|
||||
"""Log the current configuration (without sensitive values)."""
|
||||
logger.info(f"Database: {self.database_url} (schema: {self.database_schema})")
|
||||
if self.migration_database_url:
|
||||
logger.info(f"Migration database: {self.migration_database_url}")
|
||||
logger.info(f"LLM: provider={self.llm_provider}, model={self.llm_model}")
|
||||
if self.retain_llm_provider or self.retain_llm_model:
|
||||
retain_provider = self.retain_llm_provider or self.llm_provider
|
||||
|
||||
@@ -15,7 +15,12 @@ from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from hindsight_api.config import HindsightConfig, _get_raw_config, normalize_config_dict
|
||||
from hindsight_api.config import (
|
||||
RECALL_BUDGET_FUNCTIONS,
|
||||
HindsightConfig,
|
||||
_get_raw_config,
|
||||
normalize_config_dict,
|
||||
)
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
from hindsight_api.models import RequestContext
|
||||
@@ -239,6 +244,15 @@ class ConfigResolver:
|
||||
logger.warning(f"Failed to check permissions for bank {bank_id}: {e}")
|
||||
# Continue without permission check (fail open for backward compatibility)
|
||||
|
||||
# Validate entity_labels structure
|
||||
if "entity_labels" in normalized_updates and normalized_updates["entity_labels"] is not None:
|
||||
from .engine.retain.entity_labels import parse_entity_labels
|
||||
|
||||
try:
|
||||
parse_entity_labels(normalized_updates["entity_labels"])
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid entity_labels format: {e}")
|
||||
|
||||
# Validate retain_strategies: reject empty string keys
|
||||
if "retain_strategies" in normalized_updates and normalized_updates["retain_strategies"]:
|
||||
empty_keys = [k for k in normalized_updates["retain_strategies"] if not str(k).strip()]
|
||||
@@ -247,6 +261,9 @@ class ConfigResolver:
|
||||
"Strategy names must not be empty strings. Remove entries with empty names before saving."
|
||||
)
|
||||
|
||||
# Validate recall budget fields
|
||||
_validate_recall_budget_updates(normalized_updates)
|
||||
|
||||
# Merge with existing config (JSONB || operator)
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
@@ -283,6 +300,53 @@ class ConfigResolver:
|
||||
logger.info(f"Reset bank config for {bank_id} to defaults")
|
||||
|
||||
|
||||
_RECALL_BUDGET_FIXED_KEYS = (
|
||||
"recall_budget_fixed_low",
|
||||
"recall_budget_fixed_mid",
|
||||
"recall_budget_fixed_high",
|
||||
)
|
||||
_RECALL_BUDGET_ADAPTIVE_KEYS = (
|
||||
"recall_budget_adaptive_low",
|
||||
"recall_budget_adaptive_mid",
|
||||
"recall_budget_adaptive_high",
|
||||
)
|
||||
|
||||
|
||||
def _validate_recall_budget_updates(updates: dict[str, Any]) -> None:
|
||||
"""Validate recall budget config updates. Raises ValueError on invalid input."""
|
||||
if "recall_budget_function" in updates:
|
||||
function = updates["recall_budget_function"]
|
||||
if not isinstance(function, str) or function.lower() not in RECALL_BUDGET_FUNCTIONS:
|
||||
raise ValueError(
|
||||
f"recall_budget_function must be one of {sorted(RECALL_BUDGET_FUNCTIONS)}, got {function!r}"
|
||||
)
|
||||
|
||||
for key in _RECALL_BUDGET_FIXED_KEYS:
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value < 1:
|
||||
raise ValueError(f"{key} must be a positive integer, got {value!r}")
|
||||
|
||||
for key in _RECALL_BUDGET_ADAPTIVE_KEYS:
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0:
|
||||
raise ValueError(f"{key} must be a positive number, got {value!r}")
|
||||
|
||||
for key in ("recall_budget_min", "recall_budget_max"):
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value < 1:
|
||||
raise ValueError(f"{key} must be a positive integer, got {value!r}")
|
||||
|
||||
if "recall_budget_min" in updates and "recall_budget_max" in updates:
|
||||
if updates["recall_budget_min"] > updates["recall_budget_max"]:
|
||||
raise ValueError(
|
||||
f"recall_budget_min ({updates['recall_budget_min']}) must be <= "
|
||||
f"recall_budget_max ({updates['recall_budget_max']})"
|
||||
)
|
||||
|
||||
|
||||
def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConfig:
|
||||
"""
|
||||
Apply a named retain strategy's overrides on top of a resolved config.
|
||||
|
||||
@@ -0,0 +1,209 @@
|
||||
"""Audit logging for feature usage tracking.
|
||||
|
||||
Provides fire-and-forget audit logging of all mutating and core operations
|
||||
(retain, recall, reflect, bank CRUD, etc.) across HTTP, MCP, and system transports.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from ..engine.db_utils import acquire_with_retry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AuditEntry:
|
||||
"""A single audit log entry."""
|
||||
|
||||
action: str
|
||||
transport: str # "http", "mcp", "system"
|
||||
bank_id: str | None = None
|
||||
started_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
ended_at: datetime | None = None
|
||||
request: dict[str, Any] | None = None
|
||||
response: dict[str, Any] | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
def _json_default(obj: Any) -> str:
|
||||
"""JSON serializer for objects not serializable by default."""
|
||||
if isinstance(obj, datetime):
|
||||
return obj.isoformat()
|
||||
if isinstance(obj, uuid.UUID):
|
||||
return str(obj)
|
||||
if isinstance(obj, bytes):
|
||||
return "<bytes>"
|
||||
if isinstance(obj, set):
|
||||
return list(obj)
|
||||
return str(obj)
|
||||
|
||||
|
||||
def _safe_json(data: Any) -> str | None:
|
||||
"""Serialize data to JSON string, returning None on failure."""
|
||||
if data is None:
|
||||
return None
|
||||
try:
|
||||
return json.dumps(data, default=_json_default)
|
||||
except Exception:
|
||||
logger.debug("Failed to serialize audit data", exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
_SWEEP_INTERVAL_SECONDS = 3600 # Run retention sweep every hour
|
||||
|
||||
|
||||
class AuditLogger:
|
||||
"""Fire-and-forget audit log writer with optional retention sweep."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool_getter: Callable[[], asyncpg.Pool | None],
|
||||
schema_getter: Callable[[], str],
|
||||
enabled: bool,
|
||||
allowed_actions: list[str],
|
||||
retention_days: int = -1,
|
||||
) -> None:
|
||||
self._pool_getter = pool_getter
|
||||
self._schema_getter = schema_getter
|
||||
self._enabled = enabled
|
||||
self._allowed_actions: frozenset[str] | None = frozenset(allowed_actions) if allowed_actions else None
|
||||
self._retention_days = retention_days
|
||||
self._sweep_task: asyncio.Task | None = None
|
||||
|
||||
def is_enabled(self, action: str) -> bool:
|
||||
"""Check if audit logging is enabled for this action."""
|
||||
if not self._enabled:
|
||||
return False
|
||||
if self._allowed_actions is not None:
|
||||
return action in self._allowed_actions
|
||||
return True
|
||||
|
||||
def log_fire_and_forget(self, entry: AuditEntry) -> None:
|
||||
"""Schedule an audit write as a background task."""
|
||||
if not self.is_enabled(entry.action):
|
||||
return
|
||||
try:
|
||||
asyncio.create_task(self._safe_log(entry))
|
||||
except RuntimeError:
|
||||
# No running event loop (e.g. during shutdown)
|
||||
logger.debug("Cannot schedule audit log write: no running event loop")
|
||||
|
||||
async def _safe_log(self, entry: AuditEntry) -> None:
|
||||
"""Write audit entry to DB. Errors are logged, never raised."""
|
||||
pool = self._pool_getter()
|
||||
if pool is None:
|
||||
logger.debug("Audit log skipped: pool not available")
|
||||
return
|
||||
try:
|
||||
schema = self._schema_getter()
|
||||
table = f"{schema}.audit_log"
|
||||
async with acquire_with_retry(pool, max_retries=1) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table}
|
||||
(id, action, transport, bank_id, started_at, ended_at, request, response, metadata)
|
||||
VALUES
|
||||
($1, $2, $3, $4, $5, $6, $7::jsonb, $8::jsonb, $9::jsonb)
|
||||
""",
|
||||
uuid.uuid4(),
|
||||
entry.action,
|
||||
entry.transport,
|
||||
entry.bank_id,
|
||||
entry.started_at,
|
||||
entry.ended_at,
|
||||
_safe_json(entry.request),
|
||||
_safe_json(entry.response),
|
||||
_safe_json(entry.metadata) or "{}",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Audit log write failed for action={entry.action}: {e}")
|
||||
|
||||
def start_retention_sweep(self) -> None:
|
||||
"""Start the periodic retention sweep if retention is configured."""
|
||||
if self._retention_days <= 0 or not self._enabled:
|
||||
return
|
||||
try:
|
||||
self._sweep_task = asyncio.create_task(self._sweep_loop())
|
||||
except RuntimeError:
|
||||
logger.debug("Cannot start retention sweep: no running event loop")
|
||||
|
||||
async def stop_retention_sweep(self) -> None:
|
||||
"""Stop the periodic retention sweep."""
|
||||
if self._sweep_task and not self._sweep_task.done():
|
||||
self._sweep_task.cancel()
|
||||
try:
|
||||
await self._sweep_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._sweep_task = None
|
||||
|
||||
async def _sweep_loop(self) -> None:
|
||||
"""Periodically delete audit log entries older than retention_days."""
|
||||
while True:
|
||||
await self._run_sweep()
|
||||
await asyncio.sleep(_SWEEP_INTERVAL_SECONDS)
|
||||
|
||||
async def _run_sweep(self) -> None:
|
||||
"""Delete expired audit log entries. Concurrent-safe via row-level deletes."""
|
||||
pool = self._pool_getter()
|
||||
if pool is None:
|
||||
return
|
||||
try:
|
||||
schema = self._schema_getter()
|
||||
table = f"{schema}.audit_log"
|
||||
async with acquire_with_retry(pool, max_retries=1) as conn:
|
||||
result = await conn.execute(
|
||||
f"DELETE FROM {table} WHERE started_at < NOW() - INTERVAL '{self._retention_days} days'"
|
||||
)
|
||||
if result and result != "DELETE 0":
|
||||
logger.info(f"Audit log retention sweep: {result}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Audit log retention sweep failed: {e}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def audit_context(
|
||||
audit_logger: AuditLogger | None,
|
||||
action: str,
|
||||
transport: str,
|
||||
bank_id: str | None = None,
|
||||
request: dict[str, Any] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
):
|
||||
"""Async context manager that times the operation and writes audit on exit.
|
||||
|
||||
Usage:
|
||||
async with audit_context(logger, "retain", "http", bank_id, request_dict) as entry:
|
||||
result = await do_work()
|
||||
entry.response = result_dict
|
||||
"""
|
||||
if audit_logger is None or not audit_logger.is_enabled(action):
|
||||
entry = AuditEntry(action=action, transport=transport, bank_id=bank_id)
|
||||
yield entry
|
||||
return
|
||||
|
||||
entry = AuditEntry(
|
||||
action=action,
|
||||
transport=transport,
|
||||
bank_id=bank_id,
|
||||
started_at=datetime.now(timezone.utc),
|
||||
request=request,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
try:
|
||||
yield entry
|
||||
finally:
|
||||
entry.ended_at = datetime.now(timezone.utc)
|
||||
audit_logger.log_fire_and_forget(entry)
|
||||
@@ -42,6 +42,34 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _filter_live_source_memories(
|
||||
conn: "Connection",
|
||||
bank_id: str,
|
||||
source_memory_ids: list[uuid.UUID],
|
||||
) -> list[uuid.UUID]:
|
||||
"""Return only the source memory ids that still exist in the bank.
|
||||
|
||||
Uses FOR SHARE to block concurrent deletes from removing a row between the
|
||||
check and the subsequent insert/update. Combined with the delete path running
|
||||
its stale-observation sweep *after* deleting the source row, this closes the
|
||||
race window where consolidation would otherwise produce an orphan observation.
|
||||
"""
|
||||
if not source_memory_ids:
|
||||
return []
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[]) AND bank_id = $2
|
||||
FOR SHARE
|
||||
""",
|
||||
source_memory_ids,
|
||||
bank_id,
|
||||
)
|
||||
live = {row["id"] for row in rows}
|
||||
return [mid for mid in source_memory_ids if mid in live]
|
||||
|
||||
|
||||
class _CreateAction(BaseModel):
|
||||
text: str
|
||||
source_fact_ids: list[str] # memory UUIDs from the NEW FACTS list
|
||||
@@ -119,6 +147,39 @@ def _aggregate_source_fields(source_mems: list[dict[str, Any]], tags: list[str]
|
||||
)
|
||||
|
||||
|
||||
async def _count_observations_for_scope(
|
||||
conn: "Connection",
|
||||
bank_id: str,
|
||||
tags: list[str],
|
||||
) -> int:
|
||||
"""Count existing observations matching the given tag scope.
|
||||
|
||||
Returns the count of observations whose tags contain all specified tags.
|
||||
Observations with no tags are not counted (the limit does not apply to them).
|
||||
"""
|
||||
return await conn.fetchval(
|
||||
f"SELECT COUNT(*) FROM {fq_table('memory_units')} "
|
||||
f"WHERE bank_id = $1 AND fact_type = 'observation' AND tags @> $2::varchar[]",
|
||||
bank_id,
|
||||
tags,
|
||||
)
|
||||
|
||||
|
||||
def _build_response_model(max_creates: int | None = None) -> type[_ConsolidationBatchResponse]:
|
||||
"""Build a response model, optionally constraining max creates via JSON schema."""
|
||||
if max_creates is None or max_creates < 0:
|
||||
return _ConsolidationBatchResponse
|
||||
|
||||
from pydantic import Field as PydanticField
|
||||
|
||||
clamped = max(max_creates, 0)
|
||||
|
||||
class _ConstrainedConsolidationBatchResponse(_ConsolidationBatchResponse):
|
||||
creates: list[_CreateAction] = PydanticField(default=[], max_length=clamped)
|
||||
|
||||
return _ConstrainedConsolidationBatchResponse
|
||||
|
||||
|
||||
class ConsolidationPerfLog:
|
||||
"""Performance logging for consolidation operations."""
|
||||
|
||||
@@ -186,6 +247,7 @@ async def run_consolidation_job(
|
||||
|
||||
perf = ConsolidationPerfLog(bank_id)
|
||||
max_memories_per_batch = config.consolidation_batch_size
|
||||
max_memories_per_round = config.consolidation_max_memories_per_round
|
||||
llm_batch_size = max(1, config.consolidation_llm_batch_size)
|
||||
|
||||
# Check if consolidation is enabled
|
||||
@@ -248,8 +310,17 @@ async def run_consolidation_job(
|
||||
# Track all unique tags from consolidated memories for mental model refresh filtering
|
||||
consolidated_tags: set[str] = set()
|
||||
|
||||
round_limit_enabled = max_memories_per_round > 0
|
||||
round_remaining = max_memories_per_round if round_limit_enabled else float("inf")
|
||||
hit_round_limit = False
|
||||
|
||||
llm_batch_num = 0
|
||||
while True:
|
||||
# Cap fetch size by remaining round budget
|
||||
fetch_limit = (
|
||||
min(max_memories_per_batch, int(round_remaining)) if round_limit_enabled else max_memories_per_batch
|
||||
)
|
||||
|
||||
# Fetch next batch of unconsolidated memories
|
||||
async with pool.acquire() as conn:
|
||||
t0 = time.time()
|
||||
@@ -266,7 +337,7 @@ async def run_consolidation_job(
|
||||
LIMIT $2
|
||||
""",
|
||||
bank_id,
|
||||
max_memories_per_batch,
|
||||
fetch_limit,
|
||||
)
|
||||
perf.record_timing("fetch_memories", time.time() - t0)
|
||||
|
||||
@@ -491,6 +562,25 @@ async def run_consolidation_job(
|
||||
f" | avg={llm_batch_time / len(llm_batch):.3f}s/memory"
|
||||
)
|
||||
|
||||
# Update round budget after processing this DB fetch batch
|
||||
if round_limit_enabled:
|
||||
round_remaining -= len(memories)
|
||||
if round_remaining <= 0:
|
||||
hit_round_limit = True
|
||||
break
|
||||
|
||||
# Re-submit consolidation if we hit the round limit and there's likely more work
|
||||
if hit_round_limit:
|
||||
remaining = total_count - stats["memories_processed"]
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] bank={bank_id} hit round limit of {max_memories_per_round} memories,"
|
||||
f" ~{remaining} remaining. Re-queuing consolidation."
|
||||
)
|
||||
try:
|
||||
await memory_engine.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
|
||||
except Exception as e:
|
||||
logger.warning(f"[CONSOLIDATION] bank={bank_id} failed to re-queue consolidation: {e}")
|
||||
|
||||
# Build summary
|
||||
perf.log(
|
||||
f"[3] Results: {stats['memories_processed']} memories -> "
|
||||
@@ -519,16 +609,21 @@ async def run_consolidation_job(
|
||||
if timing_parts:
|
||||
perf.log(f"[4] Timing breakdown: {', '.join(timing_parts)}")
|
||||
|
||||
# Trigger mental model refreshes for models with refresh_after_consolidation=true
|
||||
# SECURITY: Only refresh mental models with matching tags (or all if no tags were consolidated)
|
||||
mental_models_refreshed = await _trigger_mental_model_refreshes(
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
consolidated_tags=list(consolidated_tags) if consolidated_tags else None,
|
||||
perf=perf,
|
||||
)
|
||||
stats["mental_models_refreshed"] = mental_models_refreshed
|
||||
# Trigger mental model refreshes only on the final round (when all memories are processed).
|
||||
# If we hit the round limit and re-queued, skip MM refresh — the next round will handle it.
|
||||
if hit_round_limit:
|
||||
stats["mental_models_refreshed"] = 0
|
||||
logger.info(f"[CONSOLIDATION] bank={bank_id} skipping mental model refresh (round limit hit, re-queued)")
|
||||
else:
|
||||
# SECURITY: Only refresh mental models with matching tags (or all if no tags were consolidated)
|
||||
mental_models_refreshed = await _trigger_mental_model_refreshes(
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
consolidated_tags=list(consolidated_tags) if consolidated_tags else None,
|
||||
perf=perf,
|
||||
)
|
||||
stats["mental_models_refreshed"] = mental_models_refreshed
|
||||
|
||||
perf.flush()
|
||||
|
||||
@@ -560,17 +655,15 @@ async def _trigger_mental_model_refreshes(
|
||||
"""
|
||||
pool = memory_engine._pool
|
||||
|
||||
# Find mental models with refresh_after_consolidation=true
|
||||
# SECURITY: Control which mental models get refreshed based on tags
|
||||
# Find mental models with refresh_after_consolidation=true that are actually stale.
|
||||
# The tag filter on the SELECT enforces the security boundary (never look outside the
|
||||
# relevant tag scope); compute_mental_model_is_stale then verifies that new memories
|
||||
# in the MM's scope really were ingested since its last refresh.
|
||||
async with pool.acquire() as conn:
|
||||
if consolidated_tags:
|
||||
# Tagged memories were consolidated - refresh:
|
||||
# 1. Mental models with overlapping tags (security boundary)
|
||||
# 2. Untagged mental models (they're "global" and available to all contexts)
|
||||
# DO NOT refresh mental models with different tags
|
||||
rows = await conn.fetch(
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags
|
||||
SELECT id, name, tags, last_refreshed_at, trigger
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
@@ -583,11 +676,9 @@ async def _trigger_mental_model_refreshes(
|
||||
consolidated_tags,
|
||||
)
|
||||
else:
|
||||
# Untagged memories were consolidated - only refresh untagged mental models
|
||||
# SECURITY: Tagged mental models are NOT refreshed when untagged memories are consolidated
|
||||
rows = await conn.fetch(
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags
|
||||
SELECT id, name, tags, last_refreshed_at, trigger
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
@@ -596,6 +687,11 @@ async def _trigger_mental_model_refreshes(
|
||||
bank_id,
|
||||
)
|
||||
|
||||
rows = []
|
||||
for candidate in candidates:
|
||||
if await memory_engine.compute_mental_model_is_stale(conn, bank_id, candidate):
|
||||
rows.append(candidate)
|
||||
|
||||
if not rows:
|
||||
return 0
|
||||
|
||||
@@ -698,24 +794,6 @@ async def _process_memory_batch(
|
||||
if recall_result.source_facts:
|
||||
union_source_facts.update(recall_result.source_facts)
|
||||
|
||||
# 3. Single LLM call
|
||||
t0 = time.time()
|
||||
llm_result = await _consolidate_batch_with_llm(
|
||||
llm_config=llm_config,
|
||||
memories=memories,
|
||||
union_observations=union_observations,
|
||||
union_source_facts=union_source_facts,
|
||||
config=config,
|
||||
)
|
||||
if perf:
|
||||
perf.record_timing("llm", time.time() - t0)
|
||||
perf.record_llm_call(llm_result.obs_count, llm_result.prompt_chars)
|
||||
|
||||
# 4. Sequential execution of creates / updates / deletes
|
||||
# Track which memory indices participated so we can build per-memory results for stats
|
||||
per_memory_created: set[str] = set()
|
||||
per_memory_updated: set[str] = set()
|
||||
|
||||
# Determine effective tag scope for observations.
|
||||
# When obs_tags_override is set, use it; otherwise use the memory's own tags.
|
||||
if obs_tags_override is not None:
|
||||
@@ -724,28 +802,52 @@ async def _process_memory_batch(
|
||||
# All memories in the batch share the same tag set (enforced by batching)
|
||||
fact_tags = memories[0].get("tags") or [] if memories else []
|
||||
|
||||
# 2b. Compute remaining observation slots for this scope (if limit configured)
|
||||
max_obs = config.max_observations_per_scope if config is not None else -1
|
||||
remaining_observation_slots: int | None = None
|
||||
if max_obs > 0 and fact_tags:
|
||||
current_count = await _count_observations_for_scope(conn, bank_id, fact_tags)
|
||||
remaining_observation_slots = max(max_obs - current_count, 0)
|
||||
if remaining_observation_slots == 0:
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] bank={bank_id} scope={fact_tags} at observation limit "
|
||||
f"({current_count}/{max_obs}), only updates/deletes allowed"
|
||||
)
|
||||
|
||||
# 3. Single LLM call
|
||||
t0 = time.time()
|
||||
llm_result = await _consolidate_batch_with_llm(
|
||||
llm_config=llm_config,
|
||||
memories=memories,
|
||||
union_observations=union_observations,
|
||||
union_source_facts=union_source_facts,
|
||||
config=config,
|
||||
remaining_observation_slots=remaining_observation_slots,
|
||||
max_observations_per_scope=max_obs,
|
||||
)
|
||||
if perf:
|
||||
perf.record_timing("llm", time.time() - t0)
|
||||
perf.record_llm_call(llm_result.obs_count, llm_result.prompt_chars)
|
||||
|
||||
# 4. Sequential execution of deletes / updates / creates
|
||||
# Deletes run first to free observation slots before creates consume them.
|
||||
# Track which memory indices participated so we can build per-memory results for stats
|
||||
per_memory_created: set[str] = set()
|
||||
per_memory_updated: set[str] = set()
|
||||
|
||||
mem_by_id = {str(m["id"]): m for m in memories}
|
||||
|
||||
for create in llm_result.creates:
|
||||
source_mems = [mem_by_id[fid] for fid in create.source_fact_ids if fid in mem_by_id]
|
||||
if not source_mems:
|
||||
# Execute deletes first to free observation slots before creates consume them
|
||||
deleted_count = 0
|
||||
for delete in llm_result.deletes:
|
||||
# Security: the observation must be present in the unioned recall
|
||||
if not any(str(obs.id) == delete.observation_id for obs in union_observations):
|
||||
logger.debug(
|
||||
f"Batch consolidation: rejected delete — observation {delete.observation_id} not in unioned recall"
|
||||
)
|
||||
continue
|
||||
agg = _aggregate_source_fields(source_mems, tags=fact_tags)
|
||||
await _execute_create_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[m["id"] for m in source_mems],
|
||||
text=create.text,
|
||||
source_fact_tags=agg.tags,
|
||||
event_date=agg.event_date,
|
||||
occurred_start=agg.occurred_start,
|
||||
occurred_end=agg.occurred_end,
|
||||
mentioned_at=agg.mentioned_at,
|
||||
perf=perf,
|
||||
)
|
||||
for m in source_mems:
|
||||
per_memory_created.add(str(m["id"]))
|
||||
await _execute_delete_action(conn=conn, bank_id=bank_id, observation_id=delete.observation_id)
|
||||
deleted_count += 1
|
||||
|
||||
for update in llm_result.updates:
|
||||
source_mems = [mem_by_id[fid] for fid in update.source_fact_ids if fid in mem_by_id]
|
||||
@@ -776,16 +878,26 @@ async def _process_memory_batch(
|
||||
for m in source_mems:
|
||||
per_memory_updated.add(str(m["id"]))
|
||||
|
||||
deleted_count = 0
|
||||
for delete in llm_result.deletes:
|
||||
# Security: the observation must be present in the unioned recall
|
||||
if not any(str(obs.id) == delete.observation_id for obs in union_observations):
|
||||
logger.debug(
|
||||
f"Batch consolidation: rejected delete — observation {delete.observation_id} not in unioned recall"
|
||||
)
|
||||
for create in llm_result.creates:
|
||||
source_mems = [mem_by_id[fid] for fid in create.source_fact_ids if fid in mem_by_id]
|
||||
if not source_mems:
|
||||
continue
|
||||
await _execute_delete_action(conn=conn, bank_id=bank_id, observation_id=delete.observation_id)
|
||||
deleted_count += 1
|
||||
agg = _aggregate_source_fields(source_mems, tags=fact_tags)
|
||||
await _execute_create_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[m["id"] for m in source_mems],
|
||||
text=create.text,
|
||||
source_fact_tags=agg.tags,
|
||||
event_date=agg.event_date,
|
||||
occurred_start=agg.occurred_start,
|
||||
occurred_end=agg.occurred_end,
|
||||
mentioned_at=agg.mentioned_at,
|
||||
perf=perf,
|
||||
)
|
||||
for m in source_mems:
|
||||
per_memory_created.add(str(m["id"]))
|
||||
|
||||
# Build per-memory result dicts for the stats tracker in the outer loop
|
||||
results: list[dict[str, Any]] = []
|
||||
@@ -840,6 +952,15 @@ async def _execute_update_action(
|
||||
logger.debug(f"Update skipped: observation {observation_id} not found in recall results")
|
||||
return
|
||||
|
||||
live_source_memory_ids = await _filter_live_source_memories(conn, bank_id, source_memory_ids)
|
||||
if not live_source_memory_ids:
|
||||
logger.debug(
|
||||
f"Update skipped: all {len(source_memory_ids)} source memories for observation "
|
||||
f"{observation_id} were deleted concurrently"
|
||||
)
|
||||
return
|
||||
source_memory_ids = live_source_memory_ids
|
||||
|
||||
from ...config import get_config
|
||||
|
||||
history_entry = {
|
||||
@@ -1082,12 +1203,16 @@ async def _consolidate_batch_with_llm(
|
||||
memories: list[dict[str, Any]],
|
||||
union_observations: "list[MemoryFact]",
|
||||
union_source_facts: "dict[str, MemoryFact]",
|
||||
config: Any = None,
|
||||
config: Any,
|
||||
remaining_observation_slots: int | None = None,
|
||||
max_observations_per_scope: int = -1,
|
||||
) -> _BatchLLMResult:
|
||||
"""Single LLM call for a batch of facts against a pooled set of observations."""
|
||||
if config is None:
|
||||
raise ValueError("config is required for _consolidate_batch_with_llm")
|
||||
if union_observations:
|
||||
obs_list = _build_observations_for_llm(union_observations, union_source_facts)
|
||||
observations_text = json.dumps(obs_list, indent=2)
|
||||
observations_text = json.dumps(obs_list, indent=2, ensure_ascii=False)
|
||||
else:
|
||||
observations_text = "[]"
|
||||
|
||||
@@ -1106,24 +1231,64 @@ async def _consolidate_batch_with_llm(
|
||||
|
||||
facts_lines = "\n".join(_fact_line(m) for m in memories)
|
||||
|
||||
observations_mission = config.observations_mission if config is not None else None
|
||||
prompt_template = build_batch_consolidation_prompt(observations_mission)
|
||||
# Build capacity note for the prompt when observation limit is configured
|
||||
observation_capacity_note: str | None = None
|
||||
if remaining_observation_slots is not None and max_observations_per_scope > 0:
|
||||
if remaining_observation_slots == 0:
|
||||
observation_capacity_note = (
|
||||
f"OBSERVATION LIMIT REACHED ({max_observations_per_scope}/{max_observations_per_scope}). "
|
||||
"Only UPDATE or DELETE existing observations. Do NOT create new ones — "
|
||||
"merge new knowledge into existing observations via UPDATE."
|
||||
)
|
||||
elif remaining_observation_slots <= len(memories):
|
||||
observation_capacity_note = (
|
||||
f"This scope has {remaining_observation_slots} observation slot(s) remaining "
|
||||
f"(out of {max_observations_per_scope}). Prefer UPDATE over CREATE when possible."
|
||||
)
|
||||
|
||||
prompt_template = build_batch_consolidation_prompt(config.observations_mission, observation_capacity_note)
|
||||
prompt = prompt_template.format(
|
||||
facts_text=facts_lines,
|
||||
observations_text=observations_text,
|
||||
)
|
||||
|
||||
max_attempts = 3
|
||||
# Use a constrained response model when observation limit is active
|
||||
response_model = _build_response_model(max_creates=remaining_observation_slots)
|
||||
|
||||
max_attempts = config.consolidation_max_attempts
|
||||
inner_max_retries = config.consolidation_llm_max_retries
|
||||
last_exc: Exception | None = None
|
||||
# Pre-compute a stable identifier set for the batch so failure logs name the
|
||||
# exact memories whose consolidation is failing — without this, an opaque
|
||||
# "LLM batch call failed" line gives operators no way to find the offending
|
||||
# input until adaptive bisection narrows the batch down to a single memory.
|
||||
memory_ids = [str(m.get("id")) for m in memories]
|
||||
if len(memory_ids) <= 5:
|
||||
ids_label = ", ".join(memory_ids)
|
||||
else:
|
||||
ids_label = f"{', '.join(memory_ids[:3])}, ... +{len(memory_ids) - 3} more"
|
||||
batch_label = f"{len(memory_ids)} memories [{ids_label}]"
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
try:
|
||||
response: _ConsolidationBatchResponse = await llm_config.call(
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
response_format=_ConsolidationBatchResponse,
|
||||
scope="consolidation",
|
||||
)
|
||||
call_kwargs: dict[str, Any] = {
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"response_format": response_model,
|
||||
"scope": "consolidation",
|
||||
}
|
||||
if inner_max_retries is not None:
|
||||
call_kwargs["max_retries"] = inner_max_retries
|
||||
response: _ConsolidationBatchResponse = await llm_config.call(**call_kwargs)
|
||||
# Defensive truncation: some LLM providers may not enforce JSON schema max_length
|
||||
creates = response.creates
|
||||
if remaining_observation_slots is not None and remaining_observation_slots >= 0:
|
||||
if len(creates) > remaining_observation_slots:
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] Truncating {len(creates)} creates to {remaining_observation_slots} "
|
||||
f"(max_observations_per_scope={max_observations_per_scope})"
|
||||
)
|
||||
creates = creates[:remaining_observation_slots]
|
||||
return _BatchLLMResult(
|
||||
creates=response.creates,
|
||||
creates=creates,
|
||||
updates=response.updates,
|
||||
deletes=response.deletes,
|
||||
obs_count=len(union_observations),
|
||||
@@ -1131,10 +1296,13 @@ async def _consolidate_batch_with_llm(
|
||||
)
|
||||
except Exception as exc:
|
||||
last_exc = exc
|
||||
logger.warning(f"[CONSOLIDATION] LLM batch call failed (attempt {attempt}/{max_attempts}): {exc}")
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] LLM batch call failed (attempt {attempt}/{max_attempts}) for {batch_label}: {exc}"
|
||||
)
|
||||
|
||||
logger.error(
|
||||
f"[CONSOLIDATION] LLM batch call failed after {max_attempts} attempts, skipping batch. Last error: {last_exc}"
|
||||
f"[CONSOLIDATION] LLM batch call failed after {max_attempts} attempts for {batch_label}, "
|
||||
f"skipping batch. Last error: {last_exc}"
|
||||
)
|
||||
return _BatchLLMResult(obs_count=len(union_observations), prompt_chars=len(prompt), failed=True)
|
||||
|
||||
@@ -1153,6 +1321,12 @@ async def _create_observation_directly(
|
||||
perf: ConsolidationPerfLog | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Create an observation from one or more source memories with pre-processed text."""
|
||||
live_source_memory_ids = await _filter_live_source_memories(conn, bank_id, source_memory_ids)
|
||||
if not live_source_memory_ids:
|
||||
logger.debug(f"Create skipped: all {len(source_memory_ids)} source memories were deleted concurrently")
|
||||
return {"action": "skipped", "reason": "sources_deleted"}
|
||||
source_memory_ids = live_source_memory_ids
|
||||
|
||||
# Generate embedding for the observation (convert to string for pgvector)
|
||||
t0 = time.time()
|
||||
embeddings = await embedding_utils.generate_embeddings_batch(memory_engine.embeddings, [observation_text])
|
||||
|
||||
@@ -5,10 +5,24 @@ _DEFAULT_MISSION = "Track every detail: names, numbers, dates, places, and relat
|
||||
|
||||
# Processing rules — always present regardless of mission
|
||||
_PROCESSING_RULES = """Processing rules (always apply):
|
||||
- REDUNDANT: same info worded differently → UPDATE the existing observation.
|
||||
- CONTRADICTION/UPDATE: capture both states with temporal markers ("used to X, now Y").
|
||||
- RESOLVE REFERENCES: when a new fact provides a concrete value resolving a vague placeholder in an existing observation (e.g. "home country", "hometown", "birthplace", "native language", "her ex", "that city"), UPDATE the observation to embed the resolved value explicitly. Example: new fact says "grandma in Sweden" + existing observation says "moved from her home country" → update to "home country is Sweden".
|
||||
- NEVER merge observations about different people or unrelated topics."""
|
||||
|
||||
1. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), etc. Never merge different facets into one observation.
|
||||
|
||||
2. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
|
||||
|
||||
3. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
|
||||
|
||||
4. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
|
||||
|
||||
5. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
|
||||
|
||||
6. SAME FACET → UPDATE, NOT CREATE: a new count supersedes the old count — UPDATE the existing count observation, don't create a second one. If there's an existing observation for the same specific facet, always UPDATE it rather than creating a duplicate.
|
||||
|
||||
7. PRESERVE HISTORY: observations that record significant events (sold, died, moved, changed) are important history — never DELETE them. Only delete an observation when it is restated identically or truly meaningless. Be very conservative with deletes.
|
||||
|
||||
8. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country" → "Sweden"), UPDATE to embed the resolved value.
|
||||
|
||||
9. NEVER merge observations about different people or unrelated topics."""
|
||||
|
||||
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
|
||||
_BATCH_DATA_SECTION = """
|
||||
@@ -26,8 +40,8 @@ Each observation includes:
|
||||
- source_memories: array of supporting facts with their text and dates
|
||||
|
||||
Compare the facts against existing observations:
|
||||
- Same topic as an existing observation → UPDATE it (observation_id + source_fact_ids)
|
||||
- New topic with durable knowledge → CREATE a new observation (source_fact_ids)
|
||||
- Same facet as an existing observation → UPDATE it (observation_id + source_fact_ids)
|
||||
- New facet with durable knowledge → CREATE a new observation (source_fact_ids)
|
||||
- Cross-reference facts within the batch: a later fact may resolve a vague reference in an earlier one
|
||||
- Purely ephemeral facts → omit them unless the MISSION above explicitly targets such data (e.g. timestamped events, session state, screen content)"""
|
||||
|
||||
@@ -66,7 +80,10 @@ Rules:
|
||||
- Return {{"creates": [], "updates": [], "deletes": []}} if nothing durable is found."""
|
||||
|
||||
|
||||
def build_batch_consolidation_prompt(observations_mission: str | None = None) -> str:
|
||||
def build_batch_consolidation_prompt(
|
||||
observations_mission: str | None = None,
|
||||
observation_capacity_note: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build the consolidation prompt for batch mode (multiple facts per LLM call).
|
||||
|
||||
@@ -75,9 +92,13 @@ def build_batch_consolidation_prompt(observations_mission: str | None = None) ->
|
||||
"""
|
||||
mission = observations_mission or _DEFAULT_MISSION
|
||||
|
||||
capacity_section = ""
|
||||
if observation_capacity_note:
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n{observation_capacity_note}"
|
||||
|
||||
return (
|
||||
"You are a memory consolidation system. Synthesize facts into observations "
|
||||
"and merge with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n{mission}\n\n"
|
||||
f"## MISSION\n{mission}{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}" + _BATCH_DATA_SECTION + _BATCH_OUTPUT_FORMAT
|
||||
)
|
||||
|
||||
@@ -20,6 +20,7 @@ from ..config import (
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL,
|
||||
DEFAULT_RERANKER_GOOGLE_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
|
||||
DEFAULT_RERANKER_LITELLM_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_SDK_MODEL,
|
||||
@@ -29,20 +30,26 @@ from ..config import (
|
||||
DEFAULT_RERANKER_LOCAL_MODEL,
|
||||
DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_RERANKER_PROVIDER,
|
||||
DEFAULT_RERANKER_SILICONFLOW_BASE_URL,
|
||||
DEFAULT_RERANKER_SILICONFLOW_MODEL,
|
||||
DEFAULT_RERANKER_TEI_BATCH_SIZE,
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL,
|
||||
ENV_RERANKER_COHERE_API_KEY,
|
||||
ENV_RERANKER_COHERE_MODEL,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
ENV_RERANKER_FLASHRANK_MODEL,
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID,
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY,
|
||||
ENV_RERANKER_LOCAL_FORCE_CPU,
|
||||
ENV_RERANKER_LOCAL_MAX_CONCURRENT,
|
||||
ENV_RERANKER_LOCAL_MODEL,
|
||||
ENV_RERANKER_LOCAL_TRUST_REMOTE_CODE,
|
||||
ENV_RERANKER_PROVIDER,
|
||||
ENV_RERANKER_SILICONFLOW_API_KEY,
|
||||
ENV_RERANKER_TEI_BATCH_SIZE,
|
||||
ENV_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
ENV_RERANKER_TEI_MAX_CONCURRENT,
|
||||
ENV_RERANKER_TEI_URL,
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY,
|
||||
@@ -516,6 +523,84 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
|
||||
return await self._predict_async(pairs)
|
||||
|
||||
|
||||
class _CohereCompatibleRerankClient:
|
||||
"""
|
||||
Internal HTTP client for Cohere-compatible /rerank endpoints.
|
||||
|
||||
Shared by all providers that speak the Cohere rerank wire format —
|
||||
{model, query, documents[, top_n]} request and
|
||||
{results: [{index, relevance_score}, ...]} response. This covers
|
||||
SiliconFlow, ZeroEntropy, Jina, Voyage, BGE self-hosted, and Cohere
|
||||
itself when reached via a custom base_url (e.g. Azure AI Foundry).
|
||||
|
||||
Not a CrossEncoderModel — providers compose it and expose their own
|
||||
provider_name / initialization logging.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str,
|
||||
rerank_url: str,
|
||||
timeout: float = 60.0,
|
||||
include_top_n: bool = True,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.rerank_url = rerank_url
|
||||
self.timeout = timeout
|
||||
self.include_top_n = include_top_n
|
||||
self._async_client: httpx.AsyncClient | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
if self._async_client is not None:
|
||||
return
|
||||
self._async_client = httpx.AsyncClient(
|
||||
timeout=self.timeout,
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
if self._async_client is None:
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
query_groups.setdefault(query, []).append((idx, text))
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
texts = [text for _, text in indexed_texts]
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
body: dict[str, object] = {
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": texts,
|
||||
"return_documents": False,
|
||||
}
|
||||
if self.include_top_n:
|
||||
body["top_n"] = len(texts)
|
||||
|
||||
response = await self._async_client.post(self.rerank_url, json=body)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
for item in result.get("results", []):
|
||||
original_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
all_scores[indices[original_idx]] = score
|
||||
|
||||
return all_scores
|
||||
|
||||
|
||||
class CohereCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Cohere cross-encoder implementation using the Cohere Rerank API.
|
||||
@@ -544,6 +629,20 @@ class CohereCrossEncoder(CrossEncoderModel):
|
||||
self.base_url = base_url
|
||||
self.timeout = timeout
|
||||
self._client = None
|
||||
# Used when base_url is set (Azure AI Foundry and other Cohere-compatible hosts).
|
||||
# Azure endpoints already include the full invoke path, so rerank_url == base_url
|
||||
# and top_n is omitted to match the existing Azure contract.
|
||||
self._http_client: _CohereCompatibleRerankClient | None = (
|
||||
_CohereCompatibleRerankClient(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
rerank_url=base_url,
|
||||
timeout=timeout,
|
||||
include_top_n=False,
|
||||
)
|
||||
if base_url
|
||||
else None
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
@@ -551,23 +650,24 @@ class CohereCrossEncoder(CrossEncoderModel):
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the Cohere client."""
|
||||
if self._client is not None:
|
||||
if self._client is not None or (self._http_client and self._http_client._async_client):
|
||||
return
|
||||
|
||||
try:
|
||||
import cohere
|
||||
except ImportError:
|
||||
raise ImportError("cohere is required for CohereCrossEncoder. Install it with: pip install cohere")
|
||||
|
||||
base_url_msg = f" at {self.base_url}" if self.base_url else ""
|
||||
logger.info(f"Reranker: initializing Cohere provider with model {self.model}{base_url_msg}")
|
||||
|
||||
# Build client kwargs, only including base_url if set (for Azure or custom endpoints)
|
||||
client_kwargs = {"api_key": self.api_key, "timeout": self.timeout}
|
||||
if self.base_url:
|
||||
client_kwargs["base_url"] = self.base_url
|
||||
self._client = cohere.Client(**client_kwargs)
|
||||
logger.info("Reranker: Cohere provider initialized")
|
||||
if self._http_client is not None:
|
||||
await self._http_client.initialize()
|
||||
logger.info("Reranker: Cohere provider initialized (Cohere-compatible HTTP endpoint)")
|
||||
else:
|
||||
# For native Cohere API, use the official SDK
|
||||
try:
|
||||
import cohere
|
||||
except ImportError:
|
||||
raise ImportError("cohere is required for CohereCrossEncoder. Install it with: pip install cohere")
|
||||
|
||||
self._client = cohere.Client(api_key=self.api_key, timeout=self.timeout)
|
||||
logger.info("Reranker: Cohere provider initialized")
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -579,25 +679,24 @@ class CohereCrossEncoder(CrossEncoderModel):
|
||||
Returns:
|
||||
List of relevance scores
|
||||
"""
|
||||
if self._client is None:
|
||||
if self._client is None and self._http_client is None:
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
# Run sync Cohere API calls in thread pool
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
if self._http_client is not None:
|
||||
return await self._http_client.predict(pairs)
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Synchronous predict implementation for Cohere API."""
|
||||
# Group pairs by query for efficient batching
|
||||
# Cohere rerank expects one query with multiple documents
|
||||
# Run sync Cohere SDK calls in thread pool
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, self._predict_sync_sdk, pairs)
|
||||
|
||||
def _predict_sync_sdk(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Synchronous predict using the native Cohere SDK."""
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
query_groups.setdefault(query, []).append((idx, text))
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
@@ -612,7 +711,6 @@ class CohereCrossEncoder(CrossEncoderModel):
|
||||
return_documents=False,
|
||||
)
|
||||
|
||||
# Map scores back to original positions
|
||||
for result in response.results:
|
||||
original_idx = result.index
|
||||
score = result.relevance_score
|
||||
@@ -629,94 +727,80 @@ class ZeroEntropyCrossEncoder(CrossEncoderModel):
|
||||
See: https://docs.zeroentropy.dev/models
|
||||
"""
|
||||
|
||||
RERANK_URL = "https://api.zeroentropy.dev/v1/models/rerank"
|
||||
DEFAULT_BASE_URL = "https://api.zeroentropy.dev"
|
||||
RERANK_PATH = "/v1/models/rerank"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_RERANKER_ZEROENTROPY_MODEL,
|
||||
base_url: str | None = None,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
"""
|
||||
Initialize ZeroEntropy cross-encoder client.
|
||||
|
||||
Args:
|
||||
api_key: ZeroEntropy API key
|
||||
model: ZeroEntropy rerank model name (default: zerank-2)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
self._async_client: httpx.AsyncClient | None = None
|
||||
self.base_url = base_url.rstrip("/") if base_url else self.DEFAULT_BASE_URL
|
||||
self._client = _CohereCompatibleRerankClient(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
rerank_url=f"{self.base_url}{self.RERANK_PATH}",
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "zeroentropy"
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the async HTTP client."""
|
||||
if self._async_client is not None:
|
||||
if self._client._async_client is not None:
|
||||
return
|
||||
|
||||
logger.info(f"Reranker: initializing ZeroEntropy provider with model {self.model}")
|
||||
self._async_client = httpx.AsyncClient(
|
||||
timeout=self.timeout,
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
await self._client.initialize()
|
||||
logger.info("Reranker: ZeroEntropy provider initialized")
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
Score query-document pairs using the ZeroEntropy Rerank API.
|
||||
return await self._client.predict(pairs)
|
||||
|
||||
Args:
|
||||
pairs: List of (query, document) tuples to score
|
||||
|
||||
Returns:
|
||||
List of relevance scores
|
||||
"""
|
||||
if self._async_client is None:
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
class SiliconFlowCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
SiliconFlow cross-encoder implementation.
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
SiliconFlow (https://siliconflow.cn) exposes a Cohere-compatible /rerank
|
||||
endpoint. Shares the HTTP client with ZeroEntropy/Cohere-custom-endpoint
|
||||
via _CohereCompatibleRerankClient.
|
||||
"""
|
||||
|
||||
# Group pairs by query for efficient batching
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
RERANK_PATH = "/rerank"
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_RERANKER_SILICONFLOW_MODEL,
|
||||
base_url: str = DEFAULT_RERANKER_SILICONFLOW_BASE_URL,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
self.model = model
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self._client = _CohereCompatibleRerankClient(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
rerank_url=f"{self.base_url}{self.RERANK_PATH}",
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
texts = [text for _, text in indexed_texts]
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "siliconflow"
|
||||
|
||||
response = await self._async_client.post(
|
||||
self.RERANK_URL,
|
||||
json={
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": texts,
|
||||
"top_n": len(texts),
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
async def initialize(self) -> None:
|
||||
if self._client._async_client is not None:
|
||||
return
|
||||
logger.info(f"Reranker: initializing SiliconFlow provider at {self.base_url} with model {self.model}")
|
||||
await self._client.initialize()
|
||||
logger.info("Reranker: SiliconFlow provider initialized")
|
||||
|
||||
# Map scores back to original positions
|
||||
for item in result.get("results", []):
|
||||
original_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
all_scores[indices[original_idx]] = score
|
||||
|
||||
return all_scores
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
return await self._client.predict(pairs)
|
||||
|
||||
|
||||
class RRFPassthroughCrossEncoder(CrossEncoderModel):
|
||||
@@ -1168,14 +1252,31 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
if self._reranker is not None:
|
||||
return
|
||||
|
||||
# Pre-warm transformers.AutoTokenizer to fully populate the transformers
|
||||
# namespace before mlx_lm imports it. transformers 5.x uses _LazyModule,
|
||||
# which has an unguarded window where `from transformers import AutoTokenizer`
|
||||
# raises ImportError if another thread is concurrently initializing the
|
||||
# namespace (e.g. embeddings init in an executor thread).
|
||||
# See: https://github.com/vectorize-io/hindsight/issues/994
|
||||
import transformers
|
||||
|
||||
_ = transformers.AutoTokenizer
|
||||
|
||||
try:
|
||||
import mlx.core # noqa: F401
|
||||
import mlx_lm # noqa: F401
|
||||
except ImportError:
|
||||
except ImportError as exc:
|
||||
# Only swallow "package not installed" errors. Anything else (e.g. a
|
||||
# transitive import failure inside mlx_lm) must surface verbatim so
|
||||
# the real cause is debuggable instead of being masked by a generic
|
||||
# "install mlx" message.
|
||||
msg = str(exc)
|
||||
if "mlx" not in msg and "mlx_lm" not in msg:
|
||||
raise
|
||||
raise ImportError(
|
||||
"mlx and mlx-lm are required for JinaMLXCrossEncoder. "
|
||||
"Install with: pip install mlx>=0.31.0 mlx-lm>=0.31.1 safetensors>=0.6.2"
|
||||
)
|
||||
) from exc
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, self._load_model)
|
||||
@@ -1183,6 +1284,7 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
def _load_model(self) -> None:
|
||||
"""Download (if needed) and load the MLX reranker. Runs in a thread."""
|
||||
import os
|
||||
import threading
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
@@ -1198,6 +1300,10 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
model_path=model_path,
|
||||
projector_path=os.path.join(model_path, "projector.safetensors"),
|
||||
)
|
||||
# MLX Metal GPU ops are not thread-safe — concurrent calls to
|
||||
# Device::end_encoding() crash with SIGSEGV (NULL deref).
|
||||
# Serialize all reranker inference through this lock.
|
||||
self._mlx_lock = threading.Lock()
|
||||
logger.info("Reranker: jina-mlx provider initialized")
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
@@ -1211,13 +1317,14 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_docs in query_groups.items():
|
||||
docs = [doc for _, doc in indexed_docs]
|
||||
indices = [idx for idx, _ in indexed_docs]
|
||||
results = self._reranker.rerank(query, docs)
|
||||
for result in results:
|
||||
original_idx = result["index"]
|
||||
all_scores[indices[original_idx]] = result["relevance_score"]
|
||||
with self._mlx_lock:
|
||||
for query, indexed_docs in query_groups.items():
|
||||
docs = [doc for _, doc in indexed_docs]
|
||||
indices = [idx for idx, _ in indexed_docs]
|
||||
results = self._reranker.rerank(query, docs)
|
||||
for result in results:
|
||||
original_idx = result["index"]
|
||||
all_scores[indices[original_idx]] = result["relevance_score"]
|
||||
|
||||
return all_scores
|
||||
|
||||
@@ -1229,6 +1336,164 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
|
||||
|
||||
class GoogleCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Google Discovery Engine cross-encoder using the Ranking REST API.
|
||||
|
||||
Uses httpx + google-auth for lightweight REST calls (no gRPC/protobuf).
|
||||
Supports ADC (Application Default Credentials) or service account key file.
|
||||
|
||||
Available models:
|
||||
- semantic-ranker-default-004: Best quality, 1024 tokens/record (recommended)
|
||||
- semantic-ranker-fast-004: Lower latency, 1024 tokens/record
|
||||
|
||||
Max 200 records per API request. Location is always "global".
|
||||
"""
|
||||
|
||||
MAX_RECORDS_PER_REQUEST = 200
|
||||
API_BASE = "https://discoveryengine.googleapis.com/v1"
|
||||
SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
project_id: str,
|
||||
model: str = DEFAULT_RERANKER_GOOGLE_MODEL,
|
||||
service_account_key: str | None = None,
|
||||
location: str = "global",
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
"""
|
||||
Initialize Google Discovery Engine cross-encoder.
|
||||
|
||||
Args:
|
||||
project_id: Google Cloud project ID
|
||||
model: Ranking model name (default: semantic-ranker-default-004)
|
||||
service_account_key: Path to service account JSON key file.
|
||||
If None, uses Application Default Credentials (ADC).
|
||||
location: API location (default: "global")
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
"""
|
||||
self.project_id = project_id
|
||||
self.model = model
|
||||
self.service_account_key = service_account_key
|
||||
self.location = location
|
||||
self.timeout = timeout
|
||||
self._credentials = None
|
||||
self._client: httpx.Client | None = None
|
||||
self._rank_url: str | None = None
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "google"
|
||||
|
||||
def _get_auth_headers(self) -> dict[str, str]:
|
||||
"""Get Authorization header with a fresh access token."""
|
||||
import google.auth.transport.requests
|
||||
|
||||
if not self._credentials.valid:
|
||||
self._credentials.refresh(google.auth.transport.requests.Request())
|
||||
return {"Authorization": f"Bearer {self._credentials.token}"}
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize credentials and HTTP client."""
|
||||
if self._client is not None:
|
||||
return
|
||||
|
||||
auth_method = "ADC" if not self.service_account_key else "service_account"
|
||||
logger.info(
|
||||
f"Reranker: initializing Google Discovery Engine provider "
|
||||
f"(project={self.project_id}, model={self.model}, auth={auth_method})"
|
||||
)
|
||||
if self.service_account_key:
|
||||
try:
|
||||
from google.oauth2 import service_account
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"google-auth is required for GoogleCrossEncoder. Install it with: pip install google-auth"
|
||||
)
|
||||
self._credentials = service_account.Credentials.from_service_account_file(
|
||||
self.service_account_key,
|
||||
scopes=self.SCOPES,
|
||||
)
|
||||
else:
|
||||
try:
|
||||
import google.auth
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"google-auth is required for GoogleCrossEncoder. Install it with: pip install google-auth"
|
||||
)
|
||||
self._credentials, _ = google.auth.default(scopes=self.SCOPES)
|
||||
|
||||
ranking_config = f"projects/{self.project_id}/locations/{self.location}/rankingConfigs/default_ranking_config"
|
||||
self._rank_url = f"{self.API_BASE}/{ranking_config}:rank"
|
||||
self._client = httpx.Client(timeout=self.timeout)
|
||||
|
||||
logger.info("Reranker: Google Discovery Engine provider initialized")
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Synchronous predict via REST API."""
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
texts = [text for _, text in indexed_texts]
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Process in batches of MAX_RECORDS_PER_REQUEST
|
||||
for batch_start in range(0, len(texts), self.MAX_RECORDS_PER_REQUEST):
|
||||
batch_texts = texts[batch_start : batch_start + self.MAX_RECORDS_PER_REQUEST]
|
||||
batch_indices = indices[batch_start : batch_start + self.MAX_RECORDS_PER_REQUEST]
|
||||
|
||||
records = [{"id": str(i), "content": text} for i, text in enumerate(batch_texts)]
|
||||
|
||||
response = self._client.post(
|
||||
self._rank_url,
|
||||
headers=self._get_auth_headers(),
|
||||
json={
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"records": records,
|
||||
"topN": len(records),
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
for record in result.get("records", []):
|
||||
local_idx = int(record["id"])
|
||||
all_scores[batch_indices[local_idx]] = record["score"]
|
||||
|
||||
return all_scores
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
Score query-document pairs using Google Discovery Engine Ranking API.
|
||||
|
||||
Args:
|
||||
pairs: List of (query, document) tuples to score
|
||||
|
||||
Returns:
|
||||
List of relevance scores (0-1, higher = more relevant)
|
||||
"""
|
||||
if self._client is None:
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel instance based on configuration.
|
||||
@@ -1249,6 +1514,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
raise ValueError(f"{ENV_RERANKER_TEI_URL} is required when {ENV_RERANKER_PROVIDER} is 'tei'")
|
||||
return RemoteTEICrossEncoder(
|
||||
base_url=url,
|
||||
timeout=config.reranker_tei_http_timeout,
|
||||
batch_size=config.reranker_tei_batch_size,
|
||||
max_concurrent=config.reranker_tei_max_concurrent,
|
||||
)
|
||||
@@ -1271,6 +1537,18 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
model=config.reranker_cohere_model,
|
||||
base_url=config.reranker_cohere_base_url,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.reranker_openrouter_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_RERANKER_OPENROUTER_API_KEY, HINDSIGHT_API_OPENROUTER_API_KEY, "
|
||||
f"or HINDSIGHT_API_LLM_API_KEY is required when {ENV_RERANKER_PROVIDER} is 'openrouter'"
|
||||
)
|
||||
return CohereCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1/rerank",
|
||||
)
|
||||
elif provider == "flashrank":
|
||||
model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL)
|
||||
cache_dir = os.environ.get(ENV_RERANKER_FLASHRANK_CACHE_DIR, DEFAULT_RERANKER_FLASHRANK_CACHE_DIR)
|
||||
@@ -1304,11 +1582,34 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_zeroentropy_model,
|
||||
)
|
||||
elif provider == "siliconflow":
|
||||
api_key = config.reranker_siliconflow_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_RERANKER_SILICONFLOW_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'siliconflow'"
|
||||
)
|
||||
return SiliconFlowCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_siliconflow_model,
|
||||
base_url=config.reranker_siliconflow_base_url,
|
||||
)
|
||||
elif provider == "google":
|
||||
project_id = config.reranker_google_project_id
|
||||
if not project_id:
|
||||
raise ValueError(
|
||||
f"{ENV_RERANKER_GOOGLE_PROJECT_ID} (or HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID) "
|
||||
f"is required when {ENV_RERANKER_PROVIDER} is 'google'"
|
||||
)
|
||||
return GoogleCrossEncoder(
|
||||
project_id=project_id,
|
||||
model=config.reranker_google_model,
|
||||
service_account_key=config.reranker_google_service_account_key,
|
||||
)
|
||||
elif provider == "rrf":
|
||||
return RRFPassthroughCrossEncoder()
|
||||
elif provider == "jina-mlx":
|
||||
return JinaMLXCrossEncoder()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
)
|
||||
|
||||
@@ -13,11 +13,13 @@ import logging
|
||||
import os
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL,
|
||||
DEFAULT_EMBEDDINGS_GEMINI_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LITELLM_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
@@ -27,6 +29,7 @@ from ..config import (
|
||||
DEFAULT_EMBEDDINGS_PROVIDER,
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY,
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY,
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY,
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL,
|
||||
@@ -426,9 +429,19 @@ class OpenAIEmbeddings(Embeddings):
|
||||
logger.info(f"Embeddings: initializing OpenAI provider with model {self.model}{base_url_msg}")
|
||||
|
||||
# Build client kwargs, only including base_url if set (for Azure or custom endpoints)
|
||||
# Parse query parameters from base_url (e.g. ?api-version=xxx for Azure OpenAI)
|
||||
# and pass them as default_query so they're included in every request.
|
||||
client_kwargs = {"api_key": self.api_key, "max_retries": self.max_retries}
|
||||
if self.base_url:
|
||||
client_kwargs["base_url"] = self.base_url
|
||||
parsed = urlparse(self.base_url)
|
||||
if parsed.query:
|
||||
clean_url = urlunparse(parsed._replace(query=""))
|
||||
client_kwargs["base_url"] = clean_url
|
||||
default_query = {k: v[0] for k, v in parse_qs(parsed.query).items()}
|
||||
client_kwargs["default_query"] = default_query
|
||||
self.base_url = clean_url
|
||||
else:
|
||||
client_kwargs["base_url"] = self.base_url
|
||||
self._client = OpenAI(**client_kwargs)
|
||||
|
||||
# Try to get dimension from known models, otherwise do a test embedding
|
||||
@@ -741,8 +754,10 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
output_dimensions: int | None = None,
|
||||
batch_size: int = 100,
|
||||
timeout: float = 60.0,
|
||||
encoding_format: str | None = "float",
|
||||
):
|
||||
"""
|
||||
Initialize LiteLLM SDK embeddings client.
|
||||
@@ -751,14 +766,19 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
api_key: API key for the embedding provider
|
||||
model: Model name with provider prefix (e.g., "cohere/embed-english-v3.0")
|
||||
api_base: Custom base URL for API (optional)
|
||||
output_dimensions: Optional output embedding dimensions (provider-dependent)
|
||||
batch_size: Maximum batch size for embedding requests (default: 100)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
encoding_format: Encoding format for embeddings (default: "float").
|
||||
Set to None or empty string to omit (needed for Voyage AI, Gemini).
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.api_base = api_base
|
||||
self.output_dimensions = output_dimensions
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self.encoding_format = encoding_format or None
|
||||
self._litellm = None # Will be set during initialization
|
||||
self._dimension: int | None = None
|
||||
|
||||
@@ -794,10 +814,13 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
"model": self.model,
|
||||
"input": ["test"],
|
||||
"api_key": self.api_key,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
if self.output_dimensions is not None:
|
||||
embed_kwargs["dimensions"] = self.output_dimensions
|
||||
|
||||
# Use async embedding method (standard in litellm)
|
||||
response = await self._litellm.aembedding(**embed_kwargs)
|
||||
@@ -841,10 +864,13 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
"api_key": self.api_key,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
if self.output_dimensions is not None:
|
||||
embed_kwargs["dimensions"] = self.output_dimensions
|
||||
|
||||
# Use sync embedding (litellm doesn't have async in thread-safe way)
|
||||
response = self._litellm.embedding(**embed_kwargs)
|
||||
@@ -866,6 +892,179 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class GeminiEmbeddings(Embeddings):
|
||||
"""
|
||||
Google embeddings via the google.genai SDK.
|
||||
|
||||
Supports both:
|
||||
1. Gemini API (api.generativeai.google.com) with API key authentication
|
||||
2. Vertex AI with service account or Application Default Credentials (ADC)
|
||||
|
||||
Uses the embed_content API: client.models.embed_content(model, contents)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = DEFAULT_EMBEDDINGS_GEMINI_MODEL,
|
||||
api_key: str | None = None,
|
||||
vertexai_project_id: str | None = None,
|
||||
vertexai_region: str | None = None,
|
||||
vertexai_service_account_key: str | None = None,
|
||||
output_dimensionality: int | None = None,
|
||||
batch_size: int = 100,
|
||||
):
|
||||
self.model = model
|
||||
self.api_key = api_key
|
||||
self.vertexai_project_id = vertexai_project_id
|
||||
self.vertexai_region = vertexai_region or "us-central1"
|
||||
self.vertexai_service_account_key = vertexai_service_account_key
|
||||
self.output_dimensionality = output_dimensionality
|
||||
self.batch_size = batch_size
|
||||
self._client = None
|
||||
self._dimension: int | None = None
|
||||
self._is_vertexai = vertexai_project_id is not None
|
||||
self._embed_config = None # EmbedContentConfig, built during initialize()
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "google"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
if self._dimension is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
return self._dimension
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the Google genai client and detect embedding dimension."""
|
||||
if self._client is not None:
|
||||
return
|
||||
|
||||
from google import genai
|
||||
from google.genai import types as genai_types
|
||||
|
||||
if self._is_vertexai:
|
||||
self._init_vertexai(genai)
|
||||
else:
|
||||
self._init_gemini(genai)
|
||||
|
||||
# Build EmbedContentConfig if output_dimensionality is set
|
||||
if self.output_dimensionality is not None:
|
||||
self._embed_config = genai_types.EmbedContentConfig(
|
||||
output_dimensionality=self.output_dimensionality,
|
||||
)
|
||||
|
||||
# Detect dimension via a test embedding (respects output_dimensionality)
|
||||
embed_kwargs = {"model": self.model, "contents": ["test"]}
|
||||
if self._embed_config is not None:
|
||||
embed_kwargs["config"] = self._embed_config
|
||||
|
||||
result = self._client.models.embed_content(**embed_kwargs) # type: ignore[union-attr]
|
||||
if result.embeddings and len(result.embeddings) > 0:
|
||||
self._dimension = len(result.embeddings[0].values)
|
||||
|
||||
auth_mode = "vertex_ai" if self._is_vertexai else "api_key"
|
||||
logger.info(
|
||||
f"Embeddings: google provider initialized (auth: {auth_mode}, model: {self.model}, dim: {self._dimension})"
|
||||
)
|
||||
|
||||
def _init_gemini(self, genai) -> None:
|
||||
"""Initialize Gemini API client with API key."""
|
||||
if not self.api_key:
|
||||
raise ValueError("Gemini embeddings provider requires an API key")
|
||||
|
||||
self._client = genai.Client(api_key=self.api_key)
|
||||
logger.info(f"Embeddings: initializing Gemini provider with model {self.model}")
|
||||
|
||||
def _init_vertexai(self, genai) -> None:
|
||||
"""Initialize Vertex AI client with project, region, and credentials."""
|
||||
if not self.vertexai_project_id:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_EMBEDDINGS_VERTEXAI_PROJECT_ID (or HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID) "
|
||||
"is required for Vertex AI embeddings provider."
|
||||
)
|
||||
|
||||
auth_method = "ADC"
|
||||
credentials = None
|
||||
|
||||
if self.vertexai_service_account_key:
|
||||
try:
|
||||
from google.oauth2 import service_account
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Vertex AI service account auth requires 'google-auth' package. "
|
||||
"Install with: pip install google-auth"
|
||||
)
|
||||
credentials = service_account.Credentials.from_service_account_file(
|
||||
self.vertexai_service_account_key,
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
)
|
||||
auth_method = "service_account"
|
||||
logger.info(f"Embeddings: Vertex AI using service account key: {self.vertexai_service_account_key}")
|
||||
|
||||
# Strip google/ prefix from model name — native SDK uses bare names
|
||||
if self.model.startswith("google/"):
|
||||
self.model = self.model[len("google/") :]
|
||||
|
||||
client_kwargs = {
|
||||
"vertexai": True,
|
||||
"project": self.vertexai_project_id,
|
||||
"location": self.vertexai_region,
|
||||
}
|
||||
if credentials is not None:
|
||||
client_kwargs["credentials"] = credentials
|
||||
|
||||
self._client = genai.Client(**client_kwargs)
|
||||
logger.info(
|
||||
f"Embeddings: initializing Vertex AI provider "
|
||||
f"(project={self.vertexai_project_id}, region={self.vertexai_region}, "
|
||||
f"model={self.model}, auth={auth_method})"
|
||||
)
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
Generate embeddings using the Google genai SDK.
|
||||
|
||||
Args:
|
||||
texts: List of text strings to encode
|
||||
|
||||
Returns:
|
||||
List of embedding vectors
|
||||
"""
|
||||
if self._client is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
all_embeddings = []
|
||||
|
||||
# Process in batches
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
embed_kwargs = {"model": self.model, "contents": batch}
|
||||
if self._embed_config is not None:
|
||||
embed_kwargs["config"] = self._embed_config
|
||||
|
||||
result = self._client.models.embed_content(**embed_kwargs)
|
||||
|
||||
all_embeddings.extend([emb.values for emb in result.embeddings])
|
||||
|
||||
# L2-normalize when output_dimensionality is set — Gemini only returns
|
||||
# normalized vectors at full 3072 dims; truncated dims need re-normalization
|
||||
# for accurate cosine similarity.
|
||||
if self.output_dimensionality is not None:
|
||||
import numpy as np
|
||||
|
||||
arr = np.array(all_embeddings)
|
||||
norms = np.linalg.norm(arr, axis=1, keepdims=True)
|
||||
norms[norms == 0] = 1
|
||||
all_embeddings = (arr / norms).tolist()
|
||||
|
||||
return all_embeddings
|
||||
|
||||
|
||||
def create_embeddings_from_env() -> Embeddings:
|
||||
"""
|
||||
Create an Embeddings instance based on configuration.
|
||||
@@ -901,7 +1100,25 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
)
|
||||
model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL)
|
||||
base_url = os.environ.get(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None
|
||||
return OpenAIEmbeddings(api_key=api_key, model=model, base_url=base_url)
|
||||
return OpenAIEmbeddings(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
base_url=base_url,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.embeddings_openrouter_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_EMBEDDINGS_OPENROUTER_API_KEY, HINDSIGHT_API_OPENROUTER_API_KEY, "
|
||||
f"or {ENV_LLM_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'openrouter'"
|
||||
)
|
||||
return OpenAIEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = config.embeddings_cohere_api_key
|
||||
if not api_key:
|
||||
@@ -927,9 +1144,30 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
api_key=api_key,
|
||||
model=config.embeddings_litellm_sdk_model,
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
|
||||
encoding_format=config.embeddings_litellm_sdk_encoding_format,
|
||||
)
|
||||
elif provider == "google":
|
||||
vertexai_project_id = config.embeddings_vertexai_project_id
|
||||
if vertexai_project_id:
|
||||
api_key = None # Vertex AI uses ADC or service account
|
||||
else:
|
||||
api_key = config.embeddings_gemini_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_GEMINI_API_KEY} or {ENV_LLM_API_KEY} is required "
|
||||
f"when {ENV_EMBEDDINGS_PROVIDER} is 'google' (set VERTEXAI_PROJECT_ID for Vertex AI auth instead)"
|
||||
)
|
||||
return GeminiEmbeddings(
|
||||
model=config.embeddings_gemini_model,
|
||||
api_key=api_key,
|
||||
vertexai_project_id=vertexai_project_id,
|
||||
vertexai_region=config.embeddings_vertexai_region,
|
||||
vertexai_service_account_key=config.embeddings_vertexai_service_account_key,
|
||||
output_dimensionality=config.embeddings_gemini_output_dimensionality,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. "
|
||||
f"Supported: 'local', 'tei', 'openai', 'cohere', 'litellm', 'litellm-sdk'"
|
||||
f"Supported: 'local', 'tei', 'openai', 'cohere', 'google', 'litellm', 'litellm-sdk'"
|
||||
)
|
||||
|
||||
@@ -75,6 +75,7 @@ class EntityResolver:
|
||||
"""
|
||||
self.pool = pool
|
||||
self.entity_lookup = entity_lookup
|
||||
self._pg_trgm_checked = False
|
||||
# Keyed by asyncio task id so concurrent retain batches never mix their
|
||||
# pending updates. flush_pending_stats() pops only the calling task's items.
|
||||
self._pending_stats: dict[int, list[_EntityStat]] = {}
|
||||
@@ -85,6 +86,19 @@ class EntityResolver:
|
||||
task = asyncio.current_task()
|
||||
return id(task) if task is not None else 0
|
||||
|
||||
def discard_pending_stats(self) -> None:
|
||||
"""
|
||||
Discard accumulated entity stats and co-occurrence counts for the current task.
|
||||
|
||||
Call this on any exception path between resolve_entities_batch /
|
||||
link_units_to_entities_batch and flush_pending_stats() to prevent the
|
||||
per-task dicts from growing unbounded when tasks fail before flushing.
|
||||
Safe to call even if no entries exist for the current task.
|
||||
"""
|
||||
key = self._task_key()
|
||||
self._pending_stats.pop(key, None)
|
||||
self._pending_cooccurrences.pop(key, None)
|
||||
|
||||
async def flush_pending_stats(self) -> None:
|
||||
"""
|
||||
Flush accumulated entity stats and co-occurrence counts for the current task.
|
||||
@@ -202,6 +216,20 @@ class EntityResolver:
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
) -> list[str]:
|
||||
if self.entity_lookup == "trigram":
|
||||
# Auto-detect pg_trgm availability on first call and fall back to
|
||||
# "full" strategy if the extension is not installed. See #626.
|
||||
if not self._pg_trgm_checked:
|
||||
self._pg_trgm_checked = True
|
||||
has_trgm = await conn.fetchval("SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'pg_trgm')")
|
||||
if not has_trgm:
|
||||
logger.warning(
|
||||
"pg_trgm extension is not available — falling back to 'full' "
|
||||
"entity lookup strategy. Install pg_trgm for faster entity "
|
||||
"resolution on large banks. See: "
|
||||
"https://github.com/vectorize-io/hindsight/issues/626"
|
||||
)
|
||||
self.entity_lookup = "full"
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_trigram(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
|
||||
@@ -289,8 +317,13 @@ class EntityResolver:
|
||||
entity_texts = list(set(e["text"] for e in entities_data))
|
||||
|
||||
# Fetch candidates for all unique entity texts in a single batched query.
|
||||
# The trigram % operator uses the GIN index; the substring conditions cover
|
||||
# exact prefix/suffix matches that trigrams might miss at low similarity.
|
||||
# Uses the GIN trigram index on LOWER(canonical_name) for case-insensitive
|
||||
# similarity lookup. Previous version also had LIKE '%...' substring fallbacks,
|
||||
# but those forced full sequential scans of the entities table and caused
|
||||
# TimeoutErrors on banks with 10k+ entities. Lowering the similarity threshold
|
||||
# to 0.15 (from default 0.3) catches most substring relationships while
|
||||
# staying fully index-based.
|
||||
await conn.execute("SET pg_trgm.similarity_threshold = 0.15")
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT ON (e.id)
|
||||
@@ -299,16 +332,13 @@ class EntityResolver:
|
||||
FROM unnest($2::text[]) AS q(query_text)
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
e.bank_id = $1
|
||||
AND (
|
||||
e.canonical_name % q.query_text
|
||||
OR LOWER(e.canonical_name) LIKE '%' || LOWER(q.query_text) || '%'
|
||||
OR LOWER(q.query_text) LIKE '%' || LOWER(e.canonical_name) || '%'
|
||||
)
|
||||
AND LOWER(e.canonical_name) % LOWER(q.query_text)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_texts,
|
||||
)
|
||||
await conn.execute("RESET pg_trgm.similarity_threshold")
|
||||
|
||||
# Group candidates by query_text
|
||||
all_candidates: dict[str, list] = {t: [] for t in entity_texts}
|
||||
@@ -477,19 +507,42 @@ class EntityResolver:
|
||||
id_by_name: dict[str, str] = {row["name_lower"]: row["id"] for row in inserted_rows}
|
||||
|
||||
# Fallback SELECT for names that conflicted (another worker won the race).
|
||||
missing = [n for n, _ in sorted_groups if n not in id_by_name]
|
||||
if missing:
|
||||
#
|
||||
# IMPORTANT: we must let PostgreSQL do the lowercasing on BOTH sides of the
|
||||
# comparison. Python's str.lower() and PostgreSQL's LOWER() differ for some
|
||||
# Unicode characters — most notably Turkish İ (U+0130):
|
||||
# Python: 'İstanbul'.lower() == 'i\u0307stanbul' (i + combining dot, 2 chars)
|
||||
# PostgreSQL: LOWER('İstanbul') == 'istanbul' (plain i, 1 char)
|
||||
# Passing a Python-lowercased name to "LOWER(canonical_name) = ANY($2::text[])"
|
||||
# would fail to match the stored entity, leaving entity_id as None and causing
|
||||
# a NOT NULL constraint violation on unit_entities.entity_id.
|
||||
#
|
||||
# Fix: pass the original (mixed-case) input names and use
|
||||
# "LOWER(canonical_name) = ANY(SELECT LOWER(n) FROM unnest($2) AS n)" so
|
||||
# PostgreSQL lowercases both sides identically. The query also returns the
|
||||
# original input_name so we can index id_by_name by Python's lower() of that
|
||||
# name, which is what the assignment loop below uses as its lookup key.
|
||||
missing_original = [g.name for name_lower, g in sorted_groups if name_lower not in id_by_name]
|
||||
if missing_original:
|
||||
existing_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, LOWER(canonical_name) AS name_lower
|
||||
FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1 AND LOWER(canonical_name) = ANY($2::text[])
|
||||
SELECT e.id, LOWER(e.canonical_name) AS name_lower, inputs.input_name
|
||||
FROM {fq_table("entities")} e
|
||||
JOIN (
|
||||
SELECT LOWER(n) AS input_name_lower, n AS input_name
|
||||
FROM unnest($2::text[]) AS n
|
||||
) AS inputs ON LOWER(e.canonical_name) = inputs.input_name_lower
|
||||
WHERE e.bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
missing,
|
||||
missing_original,
|
||||
)
|
||||
for row in existing_rows:
|
||||
id_by_name[row["name_lower"]] = row["id"]
|
||||
# Also index by Python's lower() of the original input name so the
|
||||
# assignment loop (which uses Python-lowercased keys) finds it even
|
||||
# when Python and PostgreSQL produce different lowercase strings.
|
||||
id_by_name[row["input_name"].lower()] = row["id"]
|
||||
|
||||
# Assign entity IDs back and queue one stat per original mention so that
|
||||
# flush_pending_stats() increments mention_count by the true mention count,
|
||||
@@ -757,14 +810,19 @@ class EntityResolver:
|
||||
return await self._link_units_to_entities_batch_impl(conn, unit_entity_pairs)
|
||||
|
||||
async def _link_units_to_entities_batch_impl(self, conn, unit_entity_pairs: list[tuple[str, str]]):
|
||||
# Batch insert all unit-entity links
|
||||
await conn.executemany(
|
||||
# Sorted bulk insert to prevent deadlocks from inconsistent lock ordering
|
||||
# across concurrent transactions on the unit_entities unique index.
|
||||
sorted_pairs = sorted(unit_entity_pairs)
|
||||
unit_ids = [p[0] for p in sorted_pairs]
|
||||
entity_ids = [p[1] for p in sorted_pairs]
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||||
VALUES ($1, $2)
|
||||
SELECT u, e FROM unnest($1::uuid[], $2::uuid[]) AS t(u, e)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
unit_entity_pairs,
|
||||
unit_ids,
|
||||
entity_ids,
|
||||
)
|
||||
|
||||
# Build map of unit -> entities for co-occurrence calculation
|
||||
|
||||
@@ -240,6 +240,7 @@ class MemoryEngineInterface(ABC):
|
||||
bank_id: str,
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
delete_bank_profile: bool = True,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, int]:
|
||||
"""
|
||||
@@ -248,6 +249,8 @@ class MemoryEngineInterface(ABC):
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: If specified, only delete memories of this type.
|
||||
delete_bank_profile: If True, also delete the bank profile row itself.
|
||||
If False, only delete memories/entities/documents but preserve the bank.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
|
||||
@@ -122,10 +122,14 @@ _PROVIDERS_WITHOUT_API_KEY = frozenset(
|
||||
{
|
||||
"ollama",
|
||||
"lmstudio",
|
||||
"llamacpp",
|
||||
"openai-codex",
|
||||
"claude-code",
|
||||
"mock",
|
||||
"none",
|
||||
"vertexai",
|
||||
"litellm",
|
||||
"bedrock",
|
||||
}
|
||||
)
|
||||
|
||||
@@ -143,6 +147,7 @@ def create_llm_provider(
|
||||
reasoning_effort: str,
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
vertexai_project_id: str | None = None,
|
||||
vertexai_region: str | None = None,
|
||||
vertexai_credentials: Any = None,
|
||||
@@ -159,6 +164,7 @@ def create_llm_provider(
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
groq_service_tier: Groq service tier (for Groq provider) - "on_demand", "flex", or "auto".
|
||||
openai_service_tier: OpenAI service tier (for OpenAI provider) - None (default) or "flex" (50% cheaper).
|
||||
extra_body: Extra body params merged into OpenAI-compatible API calls.
|
||||
vertexai_project_id: Vertex AI project ID (for VertexAI provider).
|
||||
vertexai_region: Vertex AI region (for VertexAI provider).
|
||||
vertexai_credentials: Vertex AI credentials object (for VertexAI provider).
|
||||
@@ -172,7 +178,10 @@ def create_llm_provider(
|
||||
ClaudeCodeLLM,
|
||||
CodexLLM,
|
||||
GeminiLLM,
|
||||
LiteLLMLLM,
|
||||
LlamaCppLLM,
|
||||
MockLLM,
|
||||
NoneLLM,
|
||||
OpenAICompatibleLLM,
|
||||
)
|
||||
|
||||
@@ -205,6 +214,15 @@ def create_llm_provider(
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
|
||||
elif provider_lower == "none":
|
||||
return NoneLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
|
||||
elif provider_lower in ("gemini", "vertexai"):
|
||||
return GeminiLLM(
|
||||
provider=provider,
|
||||
@@ -227,7 +245,45 @@ def create_llm_provider(
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
|
||||
elif provider_lower in ("openai", "groq", "ollama", "lmstudio", "minimax"):
|
||||
elif provider_lower == "litellm":
|
||||
return LiteLLMLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
|
||||
elif provider_lower == "bedrock":
|
||||
# Bedrock is a first-class alias backed by LiteLLM with auto-prefixed model names
|
||||
bedrock_model = model if model.startswith("bedrock/") else f"bedrock/{model}"
|
||||
return LiteLLMLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=bedrock_model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
|
||||
elif provider_lower == "llamacpp":
|
||||
from ..config import get_config
|
||||
|
||||
config = get_config()
|
||||
return LlamaCppLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
model_path=config.llamacpp_model_path,
|
||||
gpu_layers=config.llamacpp_gpu_layers,
|
||||
context_size=config.llamacpp_context_size,
|
||||
chat_format=config.llamacpp_chat_format,
|
||||
no_grammar=config.llamacpp_no_grammar,
|
||||
extra_args=config.llamacpp_extra_args,
|
||||
)
|
||||
|
||||
elif provider_lower in ("openai", "groq", "ollama", "lmstudio", "minimax", "volcano", "openrouter"):
|
||||
return OpenAICompatibleLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
@@ -236,6 +292,7 @@ def create_llm_provider(
|
||||
reasoning_effort=reasoning_effort,
|
||||
groq_service_tier=groq_service_tier,
|
||||
openai_service_tier=openai_service_tier,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
else:
|
||||
@@ -259,6 +316,7 @@ class LLMProvider:
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
gemini_safety_settings: list | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize LLM provider.
|
||||
@@ -272,6 +330,7 @@ class LLMProvider:
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config.
|
||||
openai_service_tier: OpenAI service tier (None or "flex") - from config.
|
||||
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
|
||||
extra_body: Extra body params merged into OpenAI-compatible API calls.
|
||||
"""
|
||||
self.provider = provider.lower()
|
||||
self.api_key = api_key
|
||||
@@ -283,6 +342,8 @@ class LLMProvider:
|
||||
self.openai_service_tier = openai_service_tier
|
||||
# Gemini safety settings (instance default; can be overridden per-request via context var)
|
||||
self.gemini_safety_settings = gemini_safety_settings
|
||||
# Extra body params for OpenAI-compatible providers (e.g. chat_template_kwargs)
|
||||
self.extra_body = extra_body
|
||||
|
||||
# Validate provider
|
||||
valid_providers = [
|
||||
@@ -292,11 +353,17 @@ class LLMProvider:
|
||||
"gemini",
|
||||
"anthropic",
|
||||
"lmstudio",
|
||||
"llamacpp",
|
||||
"vertexai",
|
||||
"openai-codex",
|
||||
"claude-code",
|
||||
"mock",
|
||||
"none",
|
||||
"minimax",
|
||||
"litellm",
|
||||
"bedrock",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}")
|
||||
@@ -311,6 +378,8 @@ class LLMProvider:
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
self.base_url = "https://api.minimax.io/v1"
|
||||
elif self.provider == "openrouter":
|
||||
self.base_url = "https://openrouter.ai/api/v1"
|
||||
|
||||
# Prepare Vertex AI config (if applicable)
|
||||
vertexai_project_id = None
|
||||
@@ -375,6 +444,7 @@ class LLMProvider:
|
||||
reasoning_effort=self.reasoning_effort,
|
||||
groq_service_tier=self.groq_service_tier,
|
||||
openai_service_tier=self.openai_service_tier,
|
||||
extra_body=self.extra_body,
|
||||
vertexai_project_id=vertexai_project_id,
|
||||
vertexai_region=vertexai_region,
|
||||
vertexai_credentials=vertexai_credentials,
|
||||
@@ -466,6 +536,15 @@ class LLMProvider:
|
||||
OutputTooLongError: If output exceeds token limits.
|
||||
Exception: Re-raises API errors after retries exhausted.
|
||||
"""
|
||||
# Stage breadcrumb so the worker log shows which LLM call a task is
|
||||
# currently inside; the stage_age field then reveals long JSON-schema
|
||||
# retry loops (e.g. a small model that can't satisfy strict_schema).
|
||||
# No-op outside a worker context.
|
||||
from ..worker.stage import set_stage
|
||||
|
||||
structured = "+structured" if response_format is not None else ""
|
||||
set_stage(f"llm.{self.provider}.{scope}{structured}")
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call(
|
||||
@@ -522,6 +601,10 @@ class LLMProvider:
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
from ..worker.stage import set_stage
|
||||
|
||||
set_stage(f"llm.{self.provider}.{scope}+tools")
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call_with_tools(
|
||||
@@ -665,64 +748,45 @@ class LLMProvider:
|
||||
return ConfiguredLLMProvider(self, config.llm_gemini_safety_settings)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources."""
|
||||
pass
|
||||
"""Clean up resources (e.g. stop llamacpp subprocess)."""
|
||||
if self._provider_impl:
|
||||
await self._provider_impl.cleanup()
|
||||
|
||||
@classmethod
|
||||
def for_memory(cls) -> "LLMProvider":
|
||||
"""Create provider for memory operations from environment variables."""
|
||||
provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq")
|
||||
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY", "")
|
||||
def from_env(cls) -> "LLMProvider":
|
||||
"""Create provider from environment variables using config.py constants."""
|
||||
from ..config import (
|
||||
DEFAULT_LLM_MODEL,
|
||||
DEFAULT_LLM_PROVIDER,
|
||||
ENV_LLM_API_KEY,
|
||||
ENV_LLM_BASE_URL,
|
||||
ENV_LLM_EXTRA_BODY,
|
||||
ENV_LLM_MODEL,
|
||||
ENV_LLM_PROVIDER,
|
||||
)
|
||||
|
||||
# API key not needed for openai-codex (uses OAuth), claude-code (uses Keychain OAuth),
|
||||
# ollama (local), or vertexai (uses GCP service account credentials)
|
||||
if not api_key and provider not in ("openai-codex", "claude-code", "ollama", "vertexai"):
|
||||
provider = os.getenv(ENV_LLM_PROVIDER, DEFAULT_LLM_PROVIDER)
|
||||
api_key = os.getenv(ENV_LLM_API_KEY, "")
|
||||
|
||||
if not api_key and not requires_api_key(provider):
|
||||
pass # Provider handles its own auth
|
||||
elif not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_API_KEY environment variable is required (unless using openai-codex or claude-code)"
|
||||
f"{ENV_LLM_API_KEY} environment variable is required (unless using openai-codex, claude-code, or litellm)"
|
||||
)
|
||||
|
||||
base_url = os.getenv("HINDSIGHT_API_LLM_BASE_URL", "")
|
||||
model = os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b")
|
||||
base_url = os.getenv(ENV_LLM_BASE_URL, "")
|
||||
model = os.getenv(ENV_LLM_MODEL, DEFAULT_LLM_MODEL)
|
||||
extra_body = json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null"))
|
||||
|
||||
return cls(provider=provider, api_key=api_key, base_url=base_url, model=model, reasoning_effort="low")
|
||||
|
||||
@classmethod
|
||||
def for_answer_generation(cls) -> "LLMProvider":
|
||||
"""Create provider for answer generation. Falls back to memory config if not set."""
|
||||
provider = os.getenv("HINDSIGHT_API_ANSWER_LLM_PROVIDER", os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"))
|
||||
api_key = os.getenv("HINDSIGHT_API_ANSWER_LLM_API_KEY", os.getenv("HINDSIGHT_API_LLM_API_KEY", ""))
|
||||
|
||||
# API key not needed for openai-codex (uses OAuth), claude-code (uses Keychain OAuth),
|
||||
# ollama (local), or vertexai (uses GCP service account credentials)
|
||||
if not api_key and provider not in ("openai-codex", "claude-code", "ollama", "vertexai"):
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_ANSWER_LLM_API_KEY environment variable is required "
|
||||
"(unless using openai-codex or claude-code)"
|
||||
)
|
||||
|
||||
base_url = os.getenv("HINDSIGHT_API_ANSWER_LLM_BASE_URL", os.getenv("HINDSIGHT_API_LLM_BASE_URL", ""))
|
||||
model = os.getenv("HINDSIGHT_API_ANSWER_LLM_MODEL", os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"))
|
||||
|
||||
return cls(provider=provider, api_key=api_key, base_url=base_url, model=model, reasoning_effort="high")
|
||||
|
||||
@classmethod
|
||||
def for_judge(cls) -> "LLMProvider":
|
||||
"""Create provider for judge/evaluator operations. Falls back to memory config if not set."""
|
||||
provider = os.getenv("HINDSIGHT_API_JUDGE_LLM_PROVIDER", os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"))
|
||||
api_key = os.getenv("HINDSIGHT_API_JUDGE_LLM_API_KEY", os.getenv("HINDSIGHT_API_LLM_API_KEY", ""))
|
||||
|
||||
# API key not needed for openai-codex (uses OAuth), claude-code (uses Keychain OAuth),
|
||||
# ollama (local), or vertexai (uses GCP service account credentials)
|
||||
if not api_key and provider not in ("openai-codex", "claude-code", "ollama", "vertexai"):
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_JUDGE_LLM_API_KEY environment variable is required "
|
||||
"(unless using openai-codex or claude-code)"
|
||||
)
|
||||
|
||||
base_url = os.getenv("HINDSIGHT_API_JUDGE_LLM_BASE_URL", os.getenv("HINDSIGHT_API_LLM_BASE_URL", ""))
|
||||
model = os.getenv("HINDSIGHT_API_JUDGE_LLM_MODEL", os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"))
|
||||
|
||||
return cls(provider=provider, api_key=api_key, base_url=base_url, model=model, reasoning_effort="high")
|
||||
return cls(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort="low",
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
|
||||
class ConfiguredLLMProvider:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -8,7 +8,20 @@ from .anthropic_llm import AnthropicLLM
|
||||
from .claude_code_llm import ClaudeCodeLLM
|
||||
from .codex_llm import CodexLLM
|
||||
from .gemini_llm import GeminiLLM
|
||||
from .litellm_llm import LiteLLMLLM
|
||||
from .llamacpp_llm import LlamaCppLLM
|
||||
from .mock_llm import MockLLM
|
||||
from .none_llm import NoneLLM
|
||||
from .openai_compatible_llm import OpenAICompatibleLLM
|
||||
|
||||
__all__ = ["AnthropicLLM", "ClaudeCodeLLM", "CodexLLM", "GeminiLLM", "MockLLM", "OpenAICompatibleLLM"]
|
||||
__all__ = [
|
||||
"AnthropicLLM",
|
||||
"ClaudeCodeLLM",
|
||||
"CodexLLM",
|
||||
"GeminiLLM",
|
||||
"LlamaCppLLM",
|
||||
"LiteLLMLLM",
|
||||
"MockLLM",
|
||||
"NoneLLM",
|
||||
"OpenAICompatibleLLM",
|
||||
]
|
||||
|
||||
@@ -153,7 +153,7 @@ class AnthropicLLM(LLMInterface):
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
if system_prompt:
|
||||
system_prompt += schema_msg
|
||||
else:
|
||||
|
||||
@@ -171,7 +171,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_instruction = (
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}\n\n"
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}\n\n"
|
||||
"Respond with ONLY the JSON, no markdown formatting."
|
||||
)
|
||||
user_content += schema_instruction
|
||||
@@ -331,7 +331,11 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools (not used by Claude Agent SDK).
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or specific function dict.
|
||||
- "auto": Model decides whether to call tools (default)
|
||||
- "required": Model must call at least one tool
|
||||
- "none": Model must not call any tools
|
||||
- {"type": "function", "function": {"name": "..."}}: Force specific tool call
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
@@ -410,16 +414,57 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
tool_call_id = msg.get("tool_call_id", "")
|
||||
user_content += f"\n\n[Tool result for {tool_call_id}: {content}]"
|
||||
|
||||
# Handle tool_choice parameter to filter tools and adjust instructions
|
||||
# The Claude Agent SDK doesn't have a native tool_choice parameter, so we
|
||||
# enforce it via allowed_tools filtering and system prompt instructions.
|
||||
|
||||
# Format tool names for SDK MCP servers: mcp__{server_name}__{tool_name}
|
||||
# This is required by the Claude Agent SDK for MCP server tools
|
||||
allowed_tool_names = [f"mcp__hindsight_tools__{name}" for name in tool_names]
|
||||
mcp_servers_config = {"hindsight_tools": mcp_server} if sdk_tools else {}
|
||||
|
||||
# Process tool_choice
|
||||
if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
||||
# Force a specific tool: filter allowed_tools to only that tool and add instruction
|
||||
forced_name = tool_choice.get("function", {}).get("name")
|
||||
if forced_name:
|
||||
# Filter to only the forced tool (with MCP prefix)
|
||||
forced_tool_mcp_name = f"mcp__hindsight_tools__{forced_name}"
|
||||
if forced_tool_mcp_name in allowed_tool_names:
|
||||
allowed_tool_names = [forced_tool_mcp_name]
|
||||
# Add strong instruction to system prompt
|
||||
force_instruction = (
|
||||
f"\n\nIMPORTANT: You MUST call the '{forced_name}' tool. Do not respond with text only."
|
||||
)
|
||||
system_prompt += force_instruction
|
||||
logger.debug(f"Claude Code: Forcing tool call to '{forced_name}'")
|
||||
else:
|
||||
logger.warning(f"Claude Code: Forced tool '{forced_name}' not found in available tools")
|
||||
elif tool_choice == "required":
|
||||
# Must call at least one tool
|
||||
tool_instruction = (
|
||||
"\n\nIMPORTANT: You MUST call at least one of the available tools. Do not respond with text only."
|
||||
)
|
||||
system_prompt += tool_instruction
|
||||
logger.debug("Claude Code: Tool call required")
|
||||
elif tool_choice == "none":
|
||||
# No tools should be called - disable all tools
|
||||
allowed_tool_names = []
|
||||
mcp_servers_config = {}
|
||||
logger.debug("Claude Code: Tools disabled (tool_choice=none)")
|
||||
# else: tool_choice == "auto" or unspecified - use default behavior (no changes needed)
|
||||
|
||||
# Configure SDK options with MCP server
|
||||
# tools=[] disables built-in CLI tools (Read, Write, Bash, ToolSearch, etc.)
|
||||
# Without this, Claude Code CLI defers MCP tools when too many built-in tools
|
||||
# are loaded, forcing Claude to use ToolSearch first — which wastes the max_turns
|
||||
# budget and prevents direct MCP tool calls.
|
||||
options = ClaudeAgentOptions(
|
||||
system_prompt=system_prompt if system_prompt else None,
|
||||
max_turns=1, # Single-turn for API-style interactions
|
||||
mcp_servers={"hindsight_tools": mcp_server} if sdk_tools else {},
|
||||
allowed_tools=allowed_tool_names if allowed_tool_names else [],
|
||||
tools=[], # Disable built-in tools so MCP tools load eagerly
|
||||
max_turns=2, # Allow tool call + tool result round-trip
|
||||
mcp_servers=mcp_servers_config,
|
||||
allowed_tools=allowed_tool_names,
|
||||
)
|
||||
|
||||
# Call Claude Agent SDK with retry logic
|
||||
|
||||
@@ -126,6 +126,32 @@ class CodexLLM(LLMInterface):
|
||||
}
|
||||
return mapping.get(effort.lower(), "auto")
|
||||
|
||||
def _normalize_tool_choice(self, tool_choice: str | dict[str, Any]) -> str | dict[str, Any]:
|
||||
"""Normalize forced function tool choice for the Codex Responses API.
|
||||
|
||||
Older agent paths may still pass OpenAI chat-completions style named
|
||||
tool choice payloads such as:
|
||||
|
||||
{"type": "function", "function": {"name": "recall"}}
|
||||
|
||||
Codex Responses expects the named function at the top level instead:
|
||||
|
||||
{"type": "function", "name": "recall"}
|
||||
"""
|
||||
if not isinstance(tool_choice, dict):
|
||||
return tool_choice
|
||||
if str(tool_choice.get("type") or "").strip() != "function":
|
||||
return tool_choice
|
||||
function_payload = tool_choice.get("function")
|
||||
if isinstance(function_payload, dict):
|
||||
function_name = str(function_payload.get("name") or "").strip()
|
||||
if function_name:
|
||||
return {"type": "function", "name": function_name}
|
||||
function_name = str(tool_choice.get("name") or "").strip()
|
||||
if function_name:
|
||||
return {"type": "function", "name": function_name}
|
||||
return tool_choice
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""Verify Codex connection by making a simple test call."""
|
||||
try:
|
||||
@@ -140,6 +166,10 @@ class CodexLLM(LLMInterface):
|
||||
)
|
||||
logger.info(f"Codex LLM verified: {self.model}")
|
||||
except Exception as e:
|
||||
# 429 means quota exhausted, not a configuration error — warn but allow startup
|
||||
if "429" in str(e) or "usage_limit_reached" in str(e):
|
||||
logger.warning(f"Codex LLM quota exhausted for {self.model}, continuing startup: {e}")
|
||||
return
|
||||
raise RuntimeError(f"Codex LLM connection verification failed for {self.model}: {e}") from e
|
||||
|
||||
async def call(
|
||||
@@ -175,7 +205,7 @@ class CodexLLM(LLMInterface):
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
system_instruction += schema_msg
|
||||
|
||||
# gpt-5.2-codex only supports "detailed" reasoning summary
|
||||
@@ -263,24 +293,27 @@ class CodexLLM(LLMInterface):
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
try:
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
|
||||
# Estimate tokens for tracing
|
||||
estimated_input = sum(len(m.get("content", "")) for m in messages) // 4
|
||||
estimated_output = len(content) // 4
|
||||
span_recorder = get_span_recorder()
|
||||
span_recorder.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=result if isinstance(result, str) else json.dumps(result),
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
duration=duration,
|
||||
finish_reason=None,
|
||||
error=None,
|
||||
)
|
||||
# Estimate tokens for tracing
|
||||
estimated_input = sum(len(m.get("content", "")) for m in messages) // 4
|
||||
estimated_output = len(content) // 4
|
||||
span_recorder = get_span_recorder()
|
||||
span_recorder.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=result if isinstance(result, str) else result.model_dump_json(),
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
duration=duration,
|
||||
finish_reason=None,
|
||||
error=None,
|
||||
)
|
||||
except Exception:
|
||||
pass # logging failure must never affect the operation
|
||||
|
||||
if return_usage:
|
||||
# Codex doesn't provide token counts, estimate based on content
|
||||
@@ -422,7 +455,7 @@ class CodexLLM(LLMInterface):
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or a specific function.
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
@@ -479,7 +512,7 @@ class CodexLLM(LLMInterface):
|
||||
"instructions": system_instruction,
|
||||
"input": user_messages,
|
||||
"tools": codex_tools,
|
||||
"tool_choice": tool_choice,
|
||||
"tool_choice": self._normalize_tool_choice(tool_choice),
|
||||
"parallel_tool_calls": True,
|
||||
"reasoning": {"summary": reasoning_summary},
|
||||
"store": False,
|
||||
@@ -526,26 +559,31 @@ class CodexLLM(LLMInterface):
|
||||
)
|
||||
|
||||
# Record OpenTelemetry span
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
try:
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
|
||||
span_recorder = get_span_recorder()
|
||||
# Convert LLMToolCall objects to dicts for span recording
|
||||
tool_calls_dict = (
|
||||
[{"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in tool_calls] if tool_calls else None
|
||||
)
|
||||
span_recorder.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=content,
|
||||
input_tokens=0, # Codex doesn't provide token counts
|
||||
output_tokens=0,
|
||||
duration=duration,
|
||||
finish_reason="tool_calls" if tool_calls else "stop",
|
||||
error=None,
|
||||
tool_calls=tool_calls_dict,
|
||||
)
|
||||
span_recorder = get_span_recorder()
|
||||
# Convert LLMToolCall objects to dicts for span recording
|
||||
tool_calls_dict = (
|
||||
[{"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in tool_calls]
|
||||
if tool_calls
|
||||
else None
|
||||
)
|
||||
span_recorder.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=content,
|
||||
input_tokens=0, # Codex doesn't provide token counts
|
||||
output_tokens=0,
|
||||
duration=duration,
|
||||
finish_reason="tool_calls" if tool_calls else "stop",
|
||||
error=None,
|
||||
tool_calls=tool_calls_dict,
|
||||
)
|
||||
except Exception:
|
||||
pass # logging failure must never affect the operation
|
||||
|
||||
return LLMToolCallResult(
|
||||
content=content,
|
||||
|
||||
@@ -7,6 +7,7 @@ This provider supports both:
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
@@ -22,6 +23,7 @@ from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.llm_wrapper import parse_llm_json
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -173,7 +175,7 @@ class GeminiLLM(LLMInterface):
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
response_format: Optional Pydantic model for structured output.
|
||||
max_completion_tokens: Maximum tokens in response (not supported by Gemini).
|
||||
max_completion_tokens: Maximum tokens in response (mapped to Gemini's max_output_tokens).
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
@@ -210,7 +212,7 @@ class GeminiLLM(LLMInterface):
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
if system_instruction:
|
||||
system_instruction += schema_msg
|
||||
else:
|
||||
@@ -225,6 +227,11 @@ class GeminiLLM(LLMInterface):
|
||||
config_kwargs["response_schema"] = response_format
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
# Gemini's equivalent of OpenAI-style max_completion_tokens is max_output_tokens.
|
||||
# Without it the model can produce arbitrarily long responses, ignoring the
|
||||
# caller's intended cap (e.g. mental_models max_tokens during refresh).
|
||||
if max_completion_tokens is not None:
|
||||
config_kwargs["max_output_tokens"] = max_completion_tokens
|
||||
|
||||
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
|
||||
effective_safety_settings = _safety_settings_ctx.get()
|
||||
@@ -241,6 +248,8 @@ class GeminiLLM(LLMInterface):
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.gemini.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
@@ -397,7 +406,7 @@ class GeminiLLM(LLMInterface):
|
||||
Args:
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Maximum tokens (not supported by Gemini).
|
||||
max_completion_tokens: Maximum tokens (mapped to Gemini's max_output_tokens).
|
||||
temperature: Sampling temperature.
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
@@ -472,9 +481,10 @@ class GeminiLLM(LLMInterface):
|
||||
fn_args = parse_llm_json(fn_args_str)
|
||||
thought_signature = tc.get("thought_signature")
|
||||
fc_kwargs: dict[str, Any] = {"name": fn_name, "args": fn_args}
|
||||
part_kwargs: dict[str, Any] = {"function_call": genai_types.FunctionCall(**fc_kwargs)}
|
||||
if thought_signature:
|
||||
fc_kwargs["thought_signature"] = thought_signature
|
||||
parts.append(genai_types.Part(function_call=genai_types.FunctionCall(**fc_kwargs)))
|
||||
part_kwargs["thought_signature"] = base64.b64decode(thought_signature)
|
||||
parts.append(genai_types.Part(**part_kwargs))
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=parts))
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
|
||||
@@ -488,6 +498,10 @@ class GeminiLLM(LLMInterface):
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
# See note in `call`: Gemini's max_output_tokens is the equivalent of
|
||||
# OpenAI-style max_completion_tokens.
|
||||
if max_completion_tokens is not None:
|
||||
config_kwargs["max_output_tokens"] = max_completion_tokens
|
||||
|
||||
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
|
||||
if tool_choice == "required":
|
||||
@@ -525,6 +539,8 @@ class GeminiLLM(LLMInterface):
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.gemini.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
@@ -547,7 +563,10 @@ class GeminiLLM(LLMInterface):
|
||||
content = part.text
|
||||
if hasattr(part, "function_call") and part.function_call:
|
||||
fc = part.function_call
|
||||
thought_signature = getattr(fc, "thought_signature", None)
|
||||
_raw_ts = getattr(part, "thought_signature", None)
|
||||
thought_signature = (
|
||||
base64.b64encode(_raw_ts).decode("ascii") if isinstance(_raw_ts, bytes) else _raw_ts
|
||||
)
|
||||
tool_calls.append(
|
||||
LLMToolCall(
|
||||
id=f"gemini_{len(tool_calls)}",
|
||||
|
||||
@@ -0,0 +1,385 @@
|
||||
"""
|
||||
LiteLLM LLM provider for universal model support.
|
||||
|
||||
This provider enables using 100+ LLM providers via the LiteLLM SDK, including:
|
||||
- AWS Bedrock (bedrock/anthropic.claude-3-5-sonnet-...)
|
||||
- Azure OpenAI (azure/gpt-4o)
|
||||
- Together AI (together_ai/meta-llama/...)
|
||||
- Any other LiteLLM-supported provider
|
||||
|
||||
Uses litellm.acompletion() for async chat completions.
|
||||
Authentication for cloud providers (e.g., AWS Bedrock via boto3 credential chain)
|
||||
is handled automatically by LiteLLM.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LiteLLMLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using the LiteLLM SDK for universal model support.
|
||||
|
||||
Supports any model accessible via litellm.acompletion(), including AWS Bedrock,
|
||||
Azure OpenAI, Together AI, Fireworks AI, and more.
|
||||
|
||||
Model names follow LiteLLM conventions with provider prefixes:
|
||||
- bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0
|
||||
- azure/gpt-4o
|
||||
- together_ai/meta-llama/Llama-3-70b-chat-hf
|
||||
- fireworks_ai/accounts/fireworks/models/llama-v3p1-70b-instruct
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
self.timeout = timeout
|
||||
self._litellm: Any = None
|
||||
|
||||
try:
|
||||
import litellm
|
||||
|
||||
self._litellm = litellm
|
||||
# Suppress LiteLLM's verbose logging
|
||||
litellm.suppress_debug_info = True # type: ignore[assignment]
|
||||
# Drop unsupported params instead of raising errors (e.g. tool_choice on some Bedrock models)
|
||||
litellm.drop_params = True # type: ignore[assignment]
|
||||
logging.getLogger("LiteLLM").setLevel(logging.WARNING)
|
||||
logger.info(f"LiteLLM SDK initialized for model: {self.model}")
|
||||
except ImportError as e:
|
||||
raise RuntimeError("LiteLLM SDK not installed. Run: uv add litellm or pip install litellm") from e
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
try:
|
||||
test_messages = [{"role": "user", "content": "test"}]
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=50,
|
||||
temperature=0.0,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
logger.info("LiteLLM connection verified successfully")
|
||||
except OutputTooLongError:
|
||||
# Truncation is fine for verification — it means the connection works
|
||||
logger.info("LiteLLM connection verified successfully (response truncated)")
|
||||
except Exception as e:
|
||||
logger.error(f"LiteLLM connection verification failed: {e}")
|
||||
raise RuntimeError(f"Failed to verify LiteLLM connection: {e}") from e
|
||||
|
||||
def _build_common_kwargs(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build common kwargs for litellm calls."""
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"timeout": self.timeout,
|
||||
}
|
||||
|
||||
if self.api_key:
|
||||
kwargs["api_key"] = self.api_key
|
||||
if self.base_url:
|
||||
kwargs["api_base"] = self.base_url
|
||||
if max_completion_tokens is not None:
|
||||
kwargs["max_completion_tokens"] = max_completion_tokens
|
||||
if temperature is not None:
|
||||
kwargs["temperature"] = temperature
|
||||
|
||||
return kwargs
|
||||
|
||||
async def call(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
response_format: Any | None = None,
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "memory",
|
||||
max_retries: int = 10,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 60.0,
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
start_time = time.time()
|
||||
|
||||
call_kwargs = self._build_common_kwargs(messages, max_completion_tokens, temperature)
|
||||
|
||||
# Add JSON schema response format if provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
call_kwargs["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": response_format.__name__ if hasattr(response_format, "__name__") else "response",
|
||||
"schema": schema,
|
||||
"strict": strict_schema,
|
||||
},
|
||||
}
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.litellm.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._litellm.acompletion(**call_kwargs)
|
||||
|
||||
content = response.choices[0].message.content or ""
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
|
||||
# Check for length-limited output
|
||||
if finish_reason == "length":
|
||||
raise OutputTooLongError("LiteLLM response was truncated due to token limit")
|
||||
|
||||
if response_format is not None:
|
||||
# Strip markdown code fences if present
|
||||
clean_content = content
|
||||
if "```json" in content:
|
||||
clean_content = content.split("```json")[1].split("```")[0].strip()
|
||||
elif "```" in content:
|
||||
clean_content = content.split("```")[1].split("```")[0].strip()
|
||||
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
json_data = json.loads(content)
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
result = content
|
||||
|
||||
# Extract usage
|
||||
input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
|
||||
output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
|
||||
total_tokens = input_tokens + output_tokens
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
|
||||
|
||||
span_recorder = get_span_recorder()
|
||||
span_recorder.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=_serialize_for_span(result),
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
)
|
||||
|
||||
if duration > 10.0:
|
||||
logger.info(
|
||||
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, "
|
||||
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
|
||||
f"time={duration:.3f}s"
|
||||
)
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
|
||||
except OutputTooLongError:
|
||||
raise
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning("LiteLLM returned invalid JSON, retrying...")
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"LiteLLM returned invalid JSON after {max_retries + 1} attempts")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
error_str = str(e).lower()
|
||||
# Fast fail on auth errors
|
||||
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
|
||||
logger.error(f"LiteLLM auth error, not retrying: {e}")
|
||||
raise
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
# Retry on rate limits, connection errors, server errors
|
||||
is_retryable = any(
|
||||
keyword in error_str
|
||||
for keyword in ("rate", "limit", "timeout", "connection", "500", "502", "503", "529")
|
||||
)
|
||||
if is_retryable:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
|
||||
await asyncio.sleep(backoff + jitter)
|
||||
continue
|
||||
|
||||
logger.error(f"LiteLLM API error after {attempt + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("LiteLLM call failed after all retries")
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]],
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "tools",
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
) -> LLMToolCallResult:
|
||||
start_time = time.time()
|
||||
|
||||
call_kwargs = self._build_common_kwargs(messages, max_completion_tokens, temperature)
|
||||
call_kwargs["tools"] = tools
|
||||
call_kwargs["tool_choice"] = tool_choice
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.litellm.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._litellm.acompletion(**call_kwargs)
|
||||
|
||||
message = response.choices[0].message
|
||||
content = message.content
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
|
||||
# Extract tool calls
|
||||
tool_calls: list[LLMToolCall] = []
|
||||
if message.tool_calls:
|
||||
for tc in message.tool_calls:
|
||||
arguments = tc.function.arguments
|
||||
if isinstance(arguments, str):
|
||||
arguments = json.loads(arguments)
|
||||
tool_calls.append(
|
||||
LLMToolCall(
|
||||
id=tc.id,
|
||||
name=tc.function.name,
|
||||
arguments=arguments,
|
||||
)
|
||||
)
|
||||
|
||||
# Extract usage
|
||||
input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
|
||||
output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
|
||||
span_recorder = get_span_recorder()
|
||||
tool_calls_dict = (
|
||||
[{"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in tool_calls]
|
||||
if tool_calls
|
||||
else None
|
||||
)
|
||||
span_recorder.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=content,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
tool_calls=tool_calls_dict,
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason=finish_reason or ("tool_calls" if tool_calls else "stop"),
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_str = str(e).lower()
|
||||
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
|
||||
raise
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
is_retryable = any(
|
||||
keyword in error_str
|
||||
for keyword in ("rate", "limit", "timeout", "connection", "500", "502", "503", "529")
|
||||
)
|
||||
if is_retryable:
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
|
||||
logger.error(f"LiteLLM tool call error after {attempt + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("LiteLLM tool call failed after all retries")
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources."""
|
||||
pass
|
||||
@@ -0,0 +1,428 @@
|
||||
"""
|
||||
Built-in llama.cpp LLM provider for fully offline operation.
|
||||
|
||||
Manages a llama-cpp-python server as a subprocess, downloads GGUF models
|
||||
from HuggingFace on first use, and delegates inference to the OpenAI-compatible API.
|
||||
|
||||
Usage:
|
||||
HINDSIGHT_API_LLM_PROVIDER=llamacpp
|
||||
HINDSIGHT_API_LLAMACPP_MODEL_PATH=~/.hindsight/models/gemma-4-E2B-it-Q4_K_M.gguf
|
||||
HINDSIGHT_API_LLAMACPP_GPU_LAYERS=-1 # -1 = all layers on GPU
|
||||
HINDSIGHT_API_LLAMACPP_CONTEXT_SIZE=8192
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.response_models import LLMToolCallResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default GGUF model for offline mode
|
||||
DEFAULT_LLAMACPP_HF_REPO = "bartowski/google_gemma-4-E2B-it-GGUF"
|
||||
DEFAULT_LLAMACPP_HF_FILENAME = "google_gemma-4-E2B-it-Q4_K_M.gguf"
|
||||
DEFAULT_LLAMACPP_MODEL_ALIAS = "gemma-4-e2b-it"
|
||||
|
||||
MODELS_DIR = Path.home() / ".hindsight" / "models"
|
||||
|
||||
# Singleton server instance — shared across all LlamaCppLLM instances
|
||||
# (retain, reflect, consolidation each create their own LLMProvider,
|
||||
# but they should all share one llama.cpp server process)
|
||||
_shared_server: "LlamaCppServer | None" = None
|
||||
_shared_server_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
"""Find a free TCP port on localhost."""
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def _download_default_model() -> Path:
|
||||
"""Download the default GGUF model from HuggingFace if not already cached.
|
||||
|
||||
Returns:
|
||||
Path to the downloaded GGUF file.
|
||||
"""
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"huggingface-hub is required for automatic model download. "
|
||||
"Install with: pip install 'hindsight-api-slim[local-llm]'"
|
||||
)
|
||||
|
||||
MODELS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
target = MODELS_DIR / DEFAULT_LLAMACPP_HF_FILENAME
|
||||
|
||||
if target.exists():
|
||||
logger.info(f"Using cached model: {target}")
|
||||
return target
|
||||
|
||||
logger.info(
|
||||
f"Downloading {DEFAULT_LLAMACPP_HF_FILENAME} from {DEFAULT_LLAMACPP_HF_REPO} (~3.5 GB, first run only)..."
|
||||
)
|
||||
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=DEFAULT_LLAMACPP_HF_REPO,
|
||||
filename=DEFAULT_LLAMACPP_HF_FILENAME,
|
||||
local_dir=str(MODELS_DIR),
|
||||
)
|
||||
|
||||
logger.info(f"Model downloaded: {downloaded}")
|
||||
return Path(downloaded)
|
||||
|
||||
|
||||
def _resolve_model_path(model_path: str | None) -> Path:
|
||||
"""Resolve the model path, downloading the default if needed.
|
||||
|
||||
Args:
|
||||
model_path: Explicit path to a GGUF file, or None to use the default.
|
||||
|
||||
Returns:
|
||||
Resolved Path to the GGUF file.
|
||||
"""
|
||||
if model_path:
|
||||
p = Path(model_path).expanduser()
|
||||
if not p.exists():
|
||||
raise FileNotFoundError(
|
||||
f"GGUF model not found: {p}\n"
|
||||
f"Set HINDSIGHT_API_LLAMACPP_MODEL_PATH to a valid .gguf file, "
|
||||
f"or remove the setting to auto-download the default model."
|
||||
)
|
||||
return p
|
||||
|
||||
return _download_default_model()
|
||||
|
||||
|
||||
class LlamaCppServer:
|
||||
"""Manages a llama-cpp-python OpenAI-compatible server as a subprocess."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: Path,
|
||||
port: int,
|
||||
gpu_layers: int = -1,
|
||||
context_size: int = 8192,
|
||||
chat_format: str | None = None,
|
||||
extra_args: str | None = None,
|
||||
):
|
||||
self.model_path = model_path
|
||||
self.port = port
|
||||
self.gpu_layers = gpu_layers
|
||||
self.context_size = context_size
|
||||
self.chat_format = chat_format
|
||||
self.extra_args = extra_args
|
||||
self._process: subprocess.Popen | None = None
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
return f"http://127.0.0.1:{self.port}/v1"
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the llama.cpp server subprocess."""
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"llama_cpp.server",
|
||||
"--model",
|
||||
str(self.model_path),
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(self.port),
|
||||
"--n_gpu_layers",
|
||||
str(self.gpu_layers),
|
||||
"--n_ctx",
|
||||
str(self.context_size),
|
||||
"--flash_attn",
|
||||
"true",
|
||||
"--n_batch",
|
||||
"2048",
|
||||
# Prompt cache: reuse KV cache for repeated system prompts
|
||||
"--cache",
|
||||
"true",
|
||||
]
|
||||
# Only pass chat_format if explicitly set (most GGUF models have it embedded)
|
||||
if self.chat_format:
|
||||
cmd.extend(["--chat_format", self.chat_format])
|
||||
# User-provided extra args (e.g. "--type_k 1 --type_v 1 --n_threads 8")
|
||||
if self.extra_args:
|
||||
cmd.extend(self.extra_args.split())
|
||||
|
||||
logger.info(f"Starting llama.cpp server: {' '.join(cmd)}")
|
||||
|
||||
# Write stderr to a log file to avoid pipe buffer deadlock
|
||||
# (llama.cpp outputs a lot of model metadata on stderr during loading)
|
||||
self._log_path = MODELS_DIR / "llamacpp_server.log"
|
||||
self._log_file = open(self._log_path, "w")
|
||||
|
||||
self._process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=self._log_file,
|
||||
# Ensure the subprocess is killed when the parent exits
|
||||
preexec_fn=os.setsid if hasattr(os, "setsid") else None,
|
||||
)
|
||||
|
||||
# Wait for the server to be ready
|
||||
await self._wait_for_ready()
|
||||
|
||||
async def _wait_for_ready(self, timeout: float = 120.0) -> None:
|
||||
"""Wait for the llama.cpp server to accept connections."""
|
||||
import httpx
|
||||
|
||||
start = time.monotonic()
|
||||
url = f"http://127.0.0.1:{self.port}/v1/models"
|
||||
last_log = start
|
||||
|
||||
while time.monotonic() - start < timeout:
|
||||
# Check if process died
|
||||
if self._process and self._process.poll() is not None:
|
||||
stderr = ""
|
||||
try:
|
||||
stderr = self._log_path.read_text()[-2000:]
|
||||
except Exception:
|
||||
pass
|
||||
raise RuntimeError(f"llama.cpp server exited with code {self._process.returncode}.\nstderr: {stderr}")
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
resp = await client.get(url, timeout=5.0)
|
||||
if resp.status_code == 200:
|
||||
logger.info(f"llama.cpp server ready on port {self.port}")
|
||||
return
|
||||
except (httpx.ConnectError, httpx.TimeoutException, httpx.ConnectTimeout):
|
||||
pass
|
||||
|
||||
# Log progress every 15s
|
||||
now = time.monotonic()
|
||||
if now - last_log > 15:
|
||||
elapsed = int(now - start)
|
||||
logger.info(f"Waiting for llama.cpp server to load model... ({elapsed}s)")
|
||||
last_log = now
|
||||
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
# Timeout — read the log to help debug
|
||||
stderr = ""
|
||||
try:
|
||||
stderr = self._log_path.read_text()[-2000:]
|
||||
except Exception:
|
||||
pass
|
||||
raise TimeoutError(
|
||||
f"llama.cpp server did not become ready within {timeout}s.\n"
|
||||
f"Check model compatibility and available memory.\n"
|
||||
f"Server log: {stderr}"
|
||||
)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the llama.cpp server subprocess."""
|
||||
if self._process is None:
|
||||
return
|
||||
|
||||
logger.info("Stopping llama.cpp server...")
|
||||
try:
|
||||
# Send SIGTERM to the process group
|
||||
if hasattr(os, "killpg"):
|
||||
os.killpg(os.getpgid(self._process.pid), signal.SIGTERM)
|
||||
else:
|
||||
self._process.terminate()
|
||||
|
||||
# Wait up to 10s for graceful shutdown
|
||||
try:
|
||||
self._process.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
if hasattr(os, "killpg"):
|
||||
os.killpg(os.getpgid(self._process.pid), signal.SIGKILL)
|
||||
else:
|
||||
self._process.kill()
|
||||
self._process.wait(timeout=5)
|
||||
except (ProcessLookupError, OSError):
|
||||
pass # Process already exited
|
||||
finally:
|
||||
self._process = None
|
||||
if hasattr(self, "_log_file") and self._log_file:
|
||||
self._log_file.close()
|
||||
self._log_file = None
|
||||
logger.info("llama.cpp server stopped")
|
||||
|
||||
|
||||
class LlamaCppLLM(LLMInterface):
|
||||
"""
|
||||
Built-in llama.cpp provider.
|
||||
|
||||
Manages a llama-cpp-python server subprocess and delegates to OpenAICompatibleLLM
|
||||
for actual inference calls. Handles model downloading and server lifecycle.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
model_path: str | None = None,
|
||||
gpu_layers: int = -1,
|
||||
context_size: int = 8192,
|
||||
chat_format: str | None = None,
|
||||
no_grammar: bool = False,
|
||||
extra_args: str | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(
|
||||
provider=provider,
|
||||
api_key=api_key or "llamacpp",
|
||||
base_url=base_url or "",
|
||||
model=model or DEFAULT_LLAMACPP_MODEL_ALIAS,
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
self._model_path_str = model_path
|
||||
self._gpu_layers = gpu_layers
|
||||
self._context_size = context_size
|
||||
self._chat_format = chat_format
|
||||
self._no_grammar = no_grammar
|
||||
self._extra_args = extra_args
|
||||
self._server: LlamaCppServer | None = None
|
||||
self._delegate: Any = None # OpenAICompatibleLLM, created after server starts
|
||||
self._initialized = False
|
||||
|
||||
async def _ensure_initialized(self) -> None:
|
||||
"""Lazy initialization: download model + start shared server on first use."""
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
global _shared_server
|
||||
|
||||
from .openai_compatible_llm import OpenAICompatibleLLM
|
||||
|
||||
async with _shared_server_lock:
|
||||
if _shared_server is None:
|
||||
# Resolve and potentially download the model
|
||||
model_path = _resolve_model_path(self._model_path_str)
|
||||
logger.info(f"Using GGUF model: {model_path}")
|
||||
|
||||
# Start the shared llama.cpp server
|
||||
port = _find_free_port()
|
||||
_shared_server = LlamaCppServer(
|
||||
model_path=model_path,
|
||||
port=port,
|
||||
gpu_layers=self._gpu_layers,
|
||||
context_size=self._context_size,
|
||||
chat_format=self._chat_format,
|
||||
extra_args=self._extra_args,
|
||||
)
|
||||
await _shared_server.start()
|
||||
|
||||
self._server = _shared_server
|
||||
|
||||
# Create the delegate that talks to the shared server's OpenAI-compatible API
|
||||
if self._no_grammar:
|
||||
logger.info("Grammar enforcement disabled (HINDSIGHT_API_LLAMACPP_NO_GRAMMAR=true)")
|
||||
self._delegate = OpenAICompatibleLLM(
|
||||
provider="llamacpp",
|
||||
api_key="llamacpp",
|
||||
base_url=self._server.base_url,
|
||||
model=self.model,
|
||||
reasoning_effort=self.reasoning_effort,
|
||||
)
|
||||
|
||||
self._initialized = True
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""Verify the llama.cpp server is running and can generate text."""
|
||||
await self._ensure_initialized()
|
||||
# Make a simple test call to verify the model can actually generate
|
||||
await self._delegate.call(
|
||||
messages=[{"role": "user", "content": "Say 'ok'"}],
|
||||
max_completion_tokens=10,
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
scope="verification",
|
||||
)
|
||||
logger.info("llama.cpp LLM verification passed")
|
||||
|
||||
async def call(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
response_format: Any | None = None,
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "memory",
|
||||
max_retries: int = 10,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 60.0,
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
"""Delegate call to the OpenAI-compatible API."""
|
||||
await self._ensure_initialized()
|
||||
return await self._delegate.call(
|
||||
messages=messages,
|
||||
response_format=response_format,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
temperature=temperature,
|
||||
scope=scope,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
skip_validation=skip_validation,
|
||||
strict_schema=strict_schema,
|
||||
return_usage=return_usage,
|
||||
)
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]],
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "tools",
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
) -> LLMToolCallResult:
|
||||
"""Delegate tool calls to the OpenAI-compatible API."""
|
||||
await self._ensure_initialized()
|
||||
return await self._delegate.call_with_tools(
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
temperature=temperature,
|
||||
scope=scope,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
tool_choice=tool_choice,
|
||||
)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Stop the shared llama.cpp server."""
|
||||
global _shared_server
|
||||
|
||||
if self._delegate:
|
||||
await self._delegate.cleanup()
|
||||
self._delegate = None
|
||||
|
||||
# Stop the shared server (only the first cleanup call actually stops it)
|
||||
async with _shared_server_lock:
|
||||
if _shared_server is not None:
|
||||
await _shared_server.stop()
|
||||
_shared_server = None
|
||||
|
||||
self._server = None
|
||||
self._initialized = False
|
||||
@@ -0,0 +1,78 @@
|
||||
"""
|
||||
No-op LLM provider for chunk-only storage mode.
|
||||
|
||||
When the LLM provider is set to "none", the system operates without any LLM dependency.
|
||||
Retain uses chunks mode (no fact extraction), and reflect/consolidation are disabled.
|
||||
This provider acts as a safety net — if any code path unexpectedly tries to call the LLM,
|
||||
it raises a clear error instead of a confusing connection failure.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from ..llm_interface import LLMInterface
|
||||
from ..response_models import LLMToolCallResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMNotAvailableError(Exception):
|
||||
"""Raised when an operation requires an LLM but the provider is set to 'none'."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class NoneLLM(LLMInterface):
|
||||
"""
|
||||
No-op LLM provider that rejects all LLM calls.
|
||||
|
||||
Used when HINDSIGHT_API_LLM_PROVIDER=none to run Hindsight as a chunk store
|
||||
with semantic search but without LLM-based features (fact extraction, reflect,
|
||||
consolidation).
|
||||
"""
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""No-op — no LLM connection to verify."""
|
||||
logger.debug("NoneLLM: no LLM connection to verify (provider=none)")
|
||||
|
||||
async def call(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
response_format: Any | None = None,
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "memory",
|
||||
max_retries: int = 10,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 60.0,
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
"""Raise LLMNotAvailableError — no LLM is configured."""
|
||||
raise LLMNotAvailableError(
|
||||
"LLM provider is set to 'none'. This operation requires an LLM. "
|
||||
"Set HINDSIGHT_API_LLM_PROVIDER to a real provider (e.g., openai, anthropic, gemini)."
|
||||
)
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]],
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "tools",
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
) -> LLMToolCallResult:
|
||||
"""Raise LLMNotAvailableError — no LLM is configured."""
|
||||
raise LLMNotAvailableError(
|
||||
"LLM provider is set to 'none'. This operation requires an LLM. "
|
||||
"Set HINDSIGHT_API_LLM_PROVIDER to a real provider (e.g., openai, anthropic, gemini)."
|
||||
)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""No-op — nothing to clean up."""
|
||||
pass
|
||||
@@ -24,6 +24,7 @@ import os
|
||||
import re
|
||||
import time
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
from openai import APIConnectionError, APIStatusError, AsyncOpenAI, LengthFinishReasonError
|
||||
@@ -32,6 +33,7 @@ from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -39,6 +41,51 @@ logger = logging.getLogger(__name__)
|
||||
DEFAULT_LLM_SEED = 4242
|
||||
|
||||
|
||||
def _strip_code_fences(content: str) -> str:
|
||||
"""Strip markdown code fences from LLM response if present.
|
||||
|
||||
Many LLM providers (MiniMax, some Ollama models, Claude via proxies)
|
||||
wrap JSON responses in ```json ... ``` fences even when json_object
|
||||
response format is requested. This strips the fences while preserving
|
||||
the JSON content inside. Returns the original content unchanged if
|
||||
no fences are detected.
|
||||
"""
|
||||
if "```" not in content:
|
||||
return content
|
||||
try:
|
||||
if "```json" in content:
|
||||
return content.split("```json")[1].split("```")[0].strip()
|
||||
return content.split("```")[1].split("```")[0].strip()
|
||||
except (IndexError, ValueError):
|
||||
return content
|
||||
|
||||
|
||||
def _summarize_status_error(e: APIStatusError, body_max: int = 400) -> str:
|
||||
"""Render an APIStatusError with status code + truncated response body.
|
||||
|
||||
Without this, retry loops only log "API error after N attempts" with the
|
||||
bare exception message — losing the provider's actual error payload, which
|
||||
is the only thing that explains *why* a request failed (rate limit reason,
|
||||
invalid tool schema, model overloaded, etc.).
|
||||
"""
|
||||
body: Any = getattr(e, "body", None)
|
||||
if body is None:
|
||||
try:
|
||||
body = e.response.text
|
||||
except Exception:
|
||||
body = None
|
||||
if isinstance(body, (dict, list)):
|
||||
try:
|
||||
body_str = json.dumps(body, default=str, ensure_ascii=False)
|
||||
except Exception:
|
||||
body_str = str(body)
|
||||
else:
|
||||
body_str = str(body or "").strip()
|
||||
if len(body_str) > body_max:
|
||||
body_str = body_str[:body_max] + "...TRUNCATED"
|
||||
return f"HTTP {e.status_code}: {body_str or '<no body>'}"
|
||||
|
||||
|
||||
class OpenAICompatibleLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for OpenAI-compatible APIs.
|
||||
@@ -60,6 +107,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float | None = None,
|
||||
groq_service_tier: str | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
@@ -73,12 +121,13 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
reasoning_effort: Reasoning effort level for supported models ("low", "medium", "high").
|
||||
timeout: Request timeout in seconds (uses env var or 300s default).
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto").
|
||||
extra_body: Extra body params merged into every API call.
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
"""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
# Validate provider
|
||||
valid_providers = ["openai", "groq", "ollama", "lmstudio", "minimax"]
|
||||
valid_providers = ["openai", "groq", "ollama", "lmstudio", "llamacpp", "minimax", "volcano", "openrouter"]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"OpenAICompatibleLLM only supports: {', '.join(valid_providers)}. Got: {self.provider}")
|
||||
|
||||
@@ -92,26 +141,38 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
self.base_url = "https://api.minimax.io/v1"
|
||||
elif self.provider == "openrouter":
|
||||
self.base_url = "https://openrouter.ai/api/v1"
|
||||
|
||||
# For ollama/lmstudio, use dummy key if not provided
|
||||
if self.provider in ("ollama", "lmstudio") and not self.api_key:
|
||||
self.api_key = "local"
|
||||
|
||||
# Validate API key for cloud providers
|
||||
if self.provider in ("openai", "groq", "minimax") and not self.api_key:
|
||||
if self.provider in ("openai", "groq", "minimax", "openrouter") and not self.api_key:
|
||||
raise ValueError(f"API key is required for {self.provider}")
|
||||
|
||||
# Service tier configuration (from config, not env vars)
|
||||
self.groq_service_tier = groq_service_tier
|
||||
self.openai_service_tier = kwargs.get("openai_service_tier")
|
||||
# User-configured extra body params (merged into every API call)
|
||||
self._config_extra_body = extra_body or {}
|
||||
|
||||
# Get timeout config
|
||||
self.timeout = timeout or float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT)))
|
||||
|
||||
# Create OpenAI client
|
||||
# Create OpenAI client — extract query params from base_url (e.g. Azure api-version)
|
||||
client_kwargs: dict[str, Any] = {"api_key": self.api_key, "max_retries": 0}
|
||||
if self.base_url:
|
||||
client_kwargs["base_url"] = self.base_url
|
||||
parsed = urlparse(self.base_url)
|
||||
if parsed.query:
|
||||
clean_url = urlunparse(parsed._replace(query=""))
|
||||
client_kwargs["base_url"] = clean_url
|
||||
default_query = {k: v[0] for k, v in parse_qs(parsed.query).items()}
|
||||
client_kwargs["default_query"] = default_query
|
||||
self.base_url = clean_url
|
||||
else:
|
||||
client_kwargs["base_url"] = self.base_url
|
||||
if self.timeout:
|
||||
client_kwargs["timeout"] = self.timeout
|
||||
|
||||
@@ -159,6 +220,36 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
return None
|
||||
|
||||
def _max_tokens_param_name(self) -> str:
|
||||
"""Return the correct parameter name for limiting response tokens.
|
||||
|
||||
Native OpenAI, Azure OpenAI, Groq, and llamacpp accept 'max_completion_tokens'.
|
||||
Mistral and other OpenAI-compatible endpoints that haven't adopted the newer
|
||||
parameter name require 'max_tokens', so when the openai provider is configured
|
||||
with a non-Azure custom base_url we fall back to the widely-supported
|
||||
'max_tokens'.
|
||||
|
||||
Reasoning models (GPT-5, o1, o3) only accept 'max_completion_tokens' and reject
|
||||
'max_tokens' outright, so they always use the new parameter name regardless of
|
||||
base_url.
|
||||
"""
|
||||
# Reasoning models (GPT-5, o1, o3, ...) only accept max_completion_tokens.
|
||||
# Azure OpenAI + GPT-5 is the canonical example: issue #978.
|
||||
if self._supports_reasoning_model():
|
||||
return "max_completion_tokens"
|
||||
# Native OpenAI (no custom base URL), Groq, and llamacpp use max_completion_tokens
|
||||
if self.provider in ("groq", "llamacpp"):
|
||||
return "max_completion_tokens"
|
||||
if self.provider == "openai" and not self.base_url:
|
||||
return "max_completion_tokens"
|
||||
# Azure OpenAI is fully OpenAI-API-compatible — detect it by hostname so users
|
||||
# can keep provider=openai + an Azure base_url (the documented setup).
|
||||
if self.provider == "openai" and self.base_url and ".openai.azure.com" in self.base_url:
|
||||
return "max_completion_tokens"
|
||||
# openai with custom base_url, ollama, lmstudio, minimax, volcano —
|
||||
# use the widely-supported max_tokens
|
||||
return "max_tokens"
|
||||
|
||||
async def call(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
@@ -231,9 +322,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# For reasoning models, enforce minimum to ensure space for reasoning + output
|
||||
if is_reasoning_model and max_completion_tokens < 16000:
|
||||
max_completion_tokens = 16000
|
||||
call_params["max_completion_tokens"] = max_completion_tokens
|
||||
|
||||
# Temperature - reasoning models don't support custom temperature
|
||||
call_params[self._max_tokens_param_name()] = max_completion_tokens
|
||||
if temperature is not None and not is_reasoning_model:
|
||||
# MiniMax requires temperature in (0.0, 1.0] — clamp accordingly
|
||||
if self.provider == "minimax":
|
||||
@@ -245,17 +334,17 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
call_params["reasoning_effort"] = self.reasoning_effort
|
||||
|
||||
# Provider-specific parameters
|
||||
extra_body: dict[str, Any] = {**self._config_extra_body}
|
||||
if self.provider == "groq":
|
||||
call_params["seed"] = DEFAULT_LLM_SEED
|
||||
extra_body: dict[str, Any] = {}
|
||||
# Add service_tier if configured
|
||||
if self.groq_service_tier:
|
||||
extra_body["service_tier"] = self.groq_service_tier
|
||||
# Add reasoning parameters for reasoning models
|
||||
if is_reasoning_model:
|
||||
extra_body["include_reasoning"] = False
|
||||
if extra_body:
|
||||
call_params["extra_body"] = extra_body
|
||||
if extra_body:
|
||||
call_params["extra_body"] = extra_body
|
||||
|
||||
# Prepare response format ONCE before retry loop
|
||||
if response_format is not None:
|
||||
@@ -276,9 +365,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
else:
|
||||
# Soft enforcement: add schema to prompt and use json_object mode
|
||||
if schema is not None:
|
||||
schema_msg = (
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
)
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
|
||||
if call_params["messages"] and call_params["messages"][0].get("role") == "system":
|
||||
first_msg = call_params["messages"][0]
|
||||
@@ -288,13 +375,23 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
first_msg = call_params["messages"][0]
|
||||
if isinstance(first_msg, dict) and isinstance(first_msg.get("content"), str):
|
||||
first_msg["content"] = schema_msg + "\n\n" + first_msg["content"]
|
||||
if self.provider not in ("lmstudio", "ollama"):
|
||||
# LM Studio and Ollama don't support json_object response format reliably
|
||||
# Providers that skip json_object grammar enforcement
|
||||
skip_grammar = self.provider in ("lmstudio", "ollama", "volcano")
|
||||
if self.provider == "llamacpp":
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
skip_grammar = get_config().llamacpp_no_grammar
|
||||
if not skip_grammar:
|
||||
call_params["response_format"] = {"type": "json_object"}
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
# Surface attempt count in worker stage so JSON-schema retry loops
|
||||
# are visible from logs (small models on strict structured output
|
||||
# often loop here). Cheap no-op outside worker context.
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
if response_format is not None:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
@@ -313,20 +410,14 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
if len(content) < original_len:
|
||||
logger.debug(f"Stripped {original_len - len(content)} chars of reasoning tokens")
|
||||
|
||||
# For local models, they may wrap JSON in markdown code blocks
|
||||
if self.provider in ("lmstudio", "ollama"):
|
||||
clean_content = content
|
||||
if "```json" in content:
|
||||
clean_content = content.split("```json")[1].split("```")[0].strip()
|
||||
elif "```" in content:
|
||||
clean_content = content.split("```")[1].split("```")[0].strip()
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
# Fallback to parsing raw content
|
||||
json_data = json.loads(content)
|
||||
else:
|
||||
# Log raw LLM response for debugging JSON parse issues
|
||||
# Strip markdown code fences if present — any provider may
|
||||
# produce these (confirmed with MiniMax, some Ollama models,
|
||||
# Claude via proxies). No-op when content is already bare JSON.
|
||||
clean_content = _strip_code_fences(content)
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
# Fallback to parsing raw content in case stripping was wrong
|
||||
try:
|
||||
json_data = json.loads(content)
|
||||
except json.JSONDecodeError as json_err:
|
||||
@@ -483,12 +574,19 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"APIStatusError ({self.provider}/{self.model}, scope={scope}, "
|
||||
f"attempt {attempt + 1}/{max_retries + 1}): {_summarize_status_error(e)}"
|
||||
)
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
|
||||
sleep_time = backoff + jitter
|
||||
await asyncio.sleep(sleep_time)
|
||||
else:
|
||||
logger.error(f"API error after {max_retries + 1} attempts: {str(e)}")
|
||||
logger.error(
|
||||
f"API error after {max_retries + 1} attempts ({self.provider}/{self.model}, "
|
||||
f"scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
@@ -551,7 +649,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
}
|
||||
|
||||
if max_completion_tokens is not None:
|
||||
call_params["max_completion_tokens"] = max_completion_tokens
|
||||
call_params[self._max_tokens_param_name()] = max_completion_tokens
|
||||
if temperature is not None:
|
||||
# MiniMax requires temperature in (0.0, 1.0] — clamp accordingly
|
||||
if self.provider == "minimax":
|
||||
@@ -559,12 +657,17 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
# Provider-specific parameters
|
||||
extra_body: dict[str, Any] = {**self._config_extra_body}
|
||||
if self.provider == "groq":
|
||||
call_params["seed"] = DEFAULT_LLM_SEED
|
||||
if extra_body:
|
||||
call_params["extra_body"] = extra_body
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self.provider}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
|
||||
@@ -634,18 +737,41 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
except APIConnectionError as e:
|
||||
last_exception = e
|
||||
status_code = getattr(e, "status_code", None) or getattr(
|
||||
getattr(e, "response", None), "status_code", None
|
||||
)
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"APIConnectionError in tool call ({self.provider}/{self.model}, scope={scope}, "
|
||||
f"attempt {attempt + 1}/{max_retries + 1}, HTTP {status_code}): {str(e)[:200]}"
|
||||
)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
logger.error(
|
||||
f"Connection error in tool call after {max_retries + 1} attempts "
|
||||
f"({self.provider}/{self.model}, scope={scope}): {str(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except APIStatusError as e:
|
||||
if e.status_code in (401, 403):
|
||||
logger.error(
|
||||
f"Auth error in tool call (HTTP {e.status_code}, {self.provider}/{self.model}), "
|
||||
f"not retrying: {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"APIStatusError in tool call ({self.provider}/{self.model}, scope={scope}, "
|
||||
f"attempt {attempt + 1}/{max_retries + 1}): {_summarize_status_error(e)}"
|
||||
)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
logger.error(
|
||||
f"API error in tool call after {max_retries + 1} attempts "
|
||||
f"({self.provider}/{self.model}, scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
@@ -693,6 +819,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
"think": False, # Disable thinking for reasoning models (qwen3.5, etc.)
|
||||
}
|
||||
|
||||
# Add schema as format parameter for structured output
|
||||
@@ -714,6 +841,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await client.post(native_url, json=payload)
|
||||
response.raise_for_status()
|
||||
@@ -721,26 +850,33 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
result = response.json()
|
||||
content = result.get("message", {}).get("content", "")
|
||||
|
||||
# Parse JSON response
|
||||
# Strip markdown code fences if present (safety net —
|
||||
# Ollama with schema enforcement usually returns bare JSON,
|
||||
# but some models may still wrap in fences)
|
||||
clean_content = _strip_code_fences(content)
|
||||
try:
|
||||
json_data = json.loads(content)
|
||||
except json.JSONDecodeError as json_err:
|
||||
content_preview = content[:500] if content else "<empty>"
|
||||
if content and len(content) > 700:
|
||||
content_preview = f"{content[:500]}...TRUNCATED...{content[-200:]}"
|
||||
logger.warning(
|
||||
f"Ollama JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {json_err}\n"
|
||||
f" Model: ollama/{self.model}\n"
|
||||
f" Content length: {len(content) if content else 0} chars\n"
|
||||
f" Content preview: {content_preview!r}"
|
||||
)
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
last_exception = json_err
|
||||
continue
|
||||
else:
|
||||
raise
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
# Fallback to raw content
|
||||
try:
|
||||
json_data = json.loads(content)
|
||||
except json.JSONDecodeError as json_err:
|
||||
content_preview = content[:500] if content else "<empty>"
|
||||
if content and len(content) > 700:
|
||||
content_preview = f"{content[:500]}...TRUNCATED...{content[-200:]}"
|
||||
logger.warning(
|
||||
f"Ollama JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {json_err}\n"
|
||||
f" Model: ollama/{self.model}\n"
|
||||
f" Content length: {len(content) if content else 0} chars\n"
|
||||
f" Content preview: {content_preview!r}"
|
||||
)
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
last_exception = json_err
|
||||
continue
|
||||
else:
|
||||
raise
|
||||
|
||||
# Extract token usage from Ollama response
|
||||
duration = time.time() - start_time
|
||||
@@ -838,7 +974,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
logger.info(f"Submitting batch with {len(requests)} requests to {self.provider}")
|
||||
|
||||
# Format requests as JSONL
|
||||
jsonl_content = "\n".join(json.dumps(req) for req in requests)
|
||||
jsonl_content = "\n".join(json.dumps(req, ensure_ascii=False) for req in requests)
|
||||
|
||||
# Upload file to provider (wrap in BytesIO with filename)
|
||||
file_bytes = io.BytesIO(jsonl_content.encode("utf-8"))
|
||||
|
||||
@@ -137,7 +137,21 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
"RETURN_AS_TIMEZONE_AWARE": False,
|
||||
}
|
||||
|
||||
results = self._search_dates(query, settings=settings)
|
||||
# Wrap dateparser in a defensive try/except. dateparser has been
|
||||
# observed to crash with internal errors (e.g., IndexError from
|
||||
# locale.translate_search) on certain query inputs. A parser bug
|
||||
# should not bring down the whole search/consolidation pipeline —
|
||||
# treat any failure as "no temporal constraint found" so the caller
|
||||
# can fall back to non-temporal retrieval.
|
||||
try:
|
||||
results = self._search_dates(query, settings=settings)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"dateparser raised %s on query (treating as no temporal constraint): %s",
|
||||
type(e).__name__,
|
||||
e,
|
||||
)
|
||||
return QueryAnalysis(temporal_constraint=None)
|
||||
|
||||
if not results:
|
||||
return QueryAnalysis(temporal_constraint=None)
|
||||
|
||||
@@ -17,7 +17,12 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
import tiktoken
|
||||
|
||||
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall
|
||||
from .prompts import FINAL_SYSTEM_PROMPT, _extract_directive_rules, build_final_prompt, build_system_prompt_for_tools
|
||||
from .prompts import (
|
||||
_extract_directive_rules,
|
||||
build_final_prompt,
|
||||
build_final_system_prompt,
|
||||
build_system_prompt_for_tools,
|
||||
)
|
||||
from .tools_schema import get_reflect_tools
|
||||
|
||||
|
||||
@@ -186,7 +191,7 @@ async def _generate_structured_output(
|
||||
DynamicModel = create_model("StructuredResponse", **fields)
|
||||
|
||||
# Include the full schema in the prompt for better LLM guidance
|
||||
schema_str = json.dumps(response_schema, indent=2)
|
||||
schema_str = json.dumps(response_schema, indent=2, ensure_ascii=False)
|
||||
|
||||
# Build field descriptions for the prompt
|
||||
field_descriptions = []
|
||||
@@ -316,6 +321,8 @@ async def run_reflect_agent(
|
||||
response_schema: dict | None = None,
|
||||
directives: list[dict[str, Any]] | None = None,
|
||||
has_mental_models: bool = False,
|
||||
include_observations: bool = True,
|
||||
include_recall: bool = True,
|
||||
budget: str | None = None,
|
||||
max_context_tokens: int = 100_000,
|
||||
) -> ReflectAgentResult:
|
||||
@@ -355,7 +362,14 @@ async def run_reflect_agent(
|
||||
directive_rules = _extract_directive_rules(directives) if directives else None
|
||||
|
||||
# Get tools for this agent (with directive compliance field if directives exist)
|
||||
tools = get_reflect_tools(directive_rules=directive_rules)
|
||||
tools = get_reflect_tools(
|
||||
directive_rules=directive_rules,
|
||||
include_mental_models=has_mental_models,
|
||||
include_observations=include_observations,
|
||||
include_recall=include_recall,
|
||||
)
|
||||
# Build set of enabled tool names to guard against LLM hallucinating disabled tool calls
|
||||
enabled_tools: frozenset[str] = frozenset(t["function"]["name"] for t in tools if t.get("type") == "function")
|
||||
|
||||
# Build initial messages (directives are injected into system prompt at START and END)
|
||||
system_prompt = build_system_prompt_for_tools(
|
||||
@@ -437,7 +451,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -494,7 +508,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -538,19 +552,18 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
|
||||
# Determine tool_choice for this iteration.
|
||||
# Force the full hierarchical retrieval path before allowing auto:
|
||||
# With mental models:
|
||||
# 0 → search_mental_models, 1 → search_observations, 2 → recall, 3+ → auto
|
||||
# Without mental models:
|
||||
# 0 → search_observations, 1 → recall, 2+ → auto
|
||||
if iteration == 0 and has_mental_models:
|
||||
iter_tool_choice: str | dict = {"type": "function", "function": {"name": "search_mental_models"}}
|
||||
elif iteration == 0:
|
||||
iter_tool_choice = {"type": "function", "function": {"name": "search_observations"}}
|
||||
elif iteration == 1 and has_mental_models:
|
||||
iter_tool_choice = {"type": "function", "function": {"name": "search_observations"}}
|
||||
elif iteration == 1 or (iteration == 2 and has_mental_models):
|
||||
iter_tool_choice = {"type": "function", "function": {"name": "recall"}}
|
||||
# Force the full hierarchical retrieval path (only for enabled tools) before allowing auto.
|
||||
# Build the forced sequence from the tools that are actually enabled.
|
||||
forced_sequence = []
|
||||
if has_mental_models:
|
||||
forced_sequence.append("search_mental_models")
|
||||
if include_observations:
|
||||
forced_sequence.append("search_observations")
|
||||
if include_recall:
|
||||
forced_sequence.append("recall")
|
||||
|
||||
if iteration < len(forced_sequence):
|
||||
iter_tool_choice: str | dict = {"type": "function", "function": {"name": forced_sequence[iteration]}}
|
||||
else:
|
||||
iter_tool_choice = "auto"
|
||||
|
||||
@@ -598,7 +611,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -641,9 +654,57 @@ async def run_reflect_agent(
|
||||
|
||||
# No tool calls - LLM wants to respond with text
|
||||
if not result.tool_calls:
|
||||
if result.content:
|
||||
# When directives are present but no evidence has been gathered,
|
||||
# the LLM tends to echo directive content verbatim as its answer.
|
||||
# Fall through to the final-prompt path which doesn't include
|
||||
# directives and handles "no data" gracefully.
|
||||
has_gathered_evidence = (
|
||||
bool(available_memory_ids) or bool(available_mental_model_ids) or bool(available_observation_ids)
|
||||
)
|
||||
directive_leak_risk = directives and not has_gathered_evidence
|
||||
if result.content and not directive_leak_risk:
|
||||
answer = _clean_answer_text(result.content.strip())
|
||||
|
||||
# The call_with_tools call above is intentionally uncapped so the
|
||||
# LLM has headroom to emit tool-call JSON plus any intermediate
|
||||
# reasoning. But when the LLM short-circuits and returns text
|
||||
# directly, that text becomes the user-visible final answer and
|
||||
# must respect max_tokens like the forced-final paths do. If it
|
||||
# overshoots, run one extra capped call to rewrite it within
|
||||
# the cap.
|
||||
if max_tokens is not None and len(_TIKTOKEN_ENCODING.encode(answer)) > max_tokens:
|
||||
rewrite_start = time.time()
|
||||
rewritten, rewrite_usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"Rewrite the user's text so it fits within the requested token "
|
||||
"budget. Preserve the key facts and structure; drop lower-priority "
|
||||
"detail. Respond with the rewritten text only, no preamble."
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}",
|
||||
},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
total_input_tokens += rewrite_usage.input_tokens
|
||||
total_output_tokens += rewrite_usage.output_tokens
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final_rewrite",
|
||||
"duration_ms": int((time.time() - rewrite_start) * 1000),
|
||||
"input_tokens": rewrite_usage.input_tokens,
|
||||
"output_tokens": rewrite_usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = _clean_answer_text(rewritten.strip())
|
||||
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
@@ -671,7 +732,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -735,7 +796,8 @@ async def run_reflect_agent(
|
||||
"content": json.dumps(
|
||||
{
|
||||
"error": "You must search for information first. Use search_mental_models(), search_observations(), or recall() before providing your final answer."
|
||||
}
|
||||
},
|
||||
ensure_ascii=False,
|
||||
),
|
||||
}
|
||||
)
|
||||
@@ -769,7 +831,17 @@ async def run_reflect_agent(
|
||||
# Execute other tools in parallel (exclude done tool in all its format variants)
|
||||
other_tools = [tc for tc in result.tool_calls if not _is_done_tool(tc.name)]
|
||||
if other_tools:
|
||||
# Add assistant message with tool calls
|
||||
# Partition into enabled vs hallucinated (not in enabled_tools set)
|
||||
allowed_tools = []
|
||||
hallucinated_tools = []
|
||||
for tc in other_tools:
|
||||
norm = _normalize_tool_name(tc.name)
|
||||
if enabled_tools is not None and norm not in enabled_tools and norm not in ("done", "expand"):
|
||||
hallucinated_tools.append(tc)
|
||||
else:
|
||||
allowed_tools.append(tc)
|
||||
|
||||
# Build assistant message with all tool calls (LLM requires them for history)
|
||||
messages.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
@@ -777,6 +849,24 @@ async def run_reflect_agent(
|
||||
}
|
||||
)
|
||||
|
||||
# Immediately reject hallucinated tool calls without adding to trace
|
||||
for tc in hallucinated_tools:
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.id,
|
||||
"name": tc.name,
|
||||
"content": json.dumps(
|
||||
{
|
||||
"error": f"Tool '{_normalize_tool_name(tc.name)}' is not available. Use only the tools provided to you."
|
||||
},
|
||||
ensure_ascii=False,
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
other_tools = allowed_tools
|
||||
|
||||
# Execute tools in parallel
|
||||
tool_tasks = [
|
||||
_execute_tool_with_timing(
|
||||
@@ -785,6 +875,7 @@ async def run_reflect_agent(
|
||||
search_observations_fn,
|
||||
recall_fn,
|
||||
expand_fn,
|
||||
enabled_tools=enabled_tools,
|
||||
)
|
||||
for tc in other_tools
|
||||
]
|
||||
@@ -840,7 +931,7 @@ async def run_reflect_agent(
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.id,
|
||||
"name": tc.name, # Required by Gemini
|
||||
"content": json.dumps(output, default=str),
|
||||
"content": json.dumps(output, default=str, ensure_ascii=False),
|
||||
}
|
||||
)
|
||||
|
||||
@@ -863,7 +954,7 @@ async def run_reflect_agent(
|
||||
)
|
||||
|
||||
try:
|
||||
output_chars = len(json.dumps(output))
|
||||
output_chars = len(json.dumps(output, ensure_ascii=False))
|
||||
except (TypeError, ValueError):
|
||||
output_chars = len(str(output))
|
||||
|
||||
@@ -900,7 +991,7 @@ def _tool_call_to_dict(tc: "LLMToolCall") -> dict[str, Any]:
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.name,
|
||||
"arguments": json.dumps(tc.arguments),
|
||||
"arguments": json.dumps(tc.arguments, ensure_ascii=False),
|
||||
},
|
||||
}
|
||||
if tc.thought_signature is not None:
|
||||
@@ -974,6 +1065,7 @@ async def _execute_tool_with_timing(
|
||||
search_observations_fn: Callable[[str, int], Awaitable[dict[str, Any]]],
|
||||
recall_fn: Callable[[str, int, int], Awaitable[dict[str, Any]]],
|
||||
expand_fn: Callable[[list[str], str], Awaitable[dict[str, Any]]],
|
||||
enabled_tools: frozenset[str] | None = None,
|
||||
) -> tuple[dict[str, Any], int]:
|
||||
"""Execute a tool call and return result with timing."""
|
||||
from hindsight_api.tracing import get_tracer
|
||||
@@ -997,7 +1089,7 @@ async def _execute_tool_with_timing(
|
||||
# Set attributes
|
||||
span.set_attribute("hindsight.tool.name", normalized_name)
|
||||
span.set_attribute("hindsight.tool.id", tc.id)
|
||||
span.set_attribute("hindsight.tool.arguments", json.dumps(tc.arguments))
|
||||
span.set_attribute("hindsight.tool.arguments", json.dumps(tc.arguments, ensure_ascii=False))
|
||||
|
||||
try:
|
||||
result = await _execute_tool(
|
||||
@@ -1007,6 +1099,7 @@ async def _execute_tool_with_timing(
|
||||
search_observations_fn,
|
||||
recall_fn,
|
||||
expand_fn,
|
||||
enabled_tools=enabled_tools,
|
||||
)
|
||||
|
||||
# Set success attributes
|
||||
@@ -1046,11 +1139,16 @@ async def _execute_tool(
|
||||
search_observations_fn: Callable[[str, int], Awaitable[dict[str, Any]]],
|
||||
recall_fn: Callable[[str, int, int], Awaitable[dict[str, Any]]],
|
||||
expand_fn: Callable[[list[str], str], Awaitable[dict[str, Any]]],
|
||||
enabled_tools: frozenset[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Execute a single tool by name."""
|
||||
# Normalize tool name for various LLM output formats
|
||||
tool_name = _normalize_tool_name(tool_name)
|
||||
|
||||
# Guard against LLMs hallucinating calls to tools that were not provided
|
||||
if enabled_tools is not None and tool_name not in enabled_tools and tool_name not in ("done", "expand"):
|
||||
return {"error": f"Tool '{tool_name}' is not available. Use only the tools provided to you."}
|
||||
|
||||
if tool_name == "search_mental_models":
|
||||
query = args.get("query")
|
||||
if not query:
|
||||
|
||||
@@ -0,0 +1,307 @@
|
||||
"""Delta operations for structured mental models.
|
||||
|
||||
The LLM's job during a delta refresh is to emit a list of these operations,
|
||||
each targeting an existing section (by id) or referencing a position relative
|
||||
to one. ``apply_operations`` validates and applies each op in turn against a
|
||||
copy of the document; invalid ops (unknown ``section_id``, out-of-range
|
||||
``block_index``, malformed payloads) are dropped with a debug-friendly reason.
|
||||
|
||||
Sections and blocks not mentioned by any op are physically copied through
|
||||
unchanged — there is no LLM-mediated re-emission of unchanged text, so prose
|
||||
drift is structurally impossible.
|
||||
|
||||
Why operations and not "output the new structured doc":
|
||||
- "Output the new doc" still asks the LLM to *generate* every section's
|
||||
blocks, including ones it didn't intend to modify, which gives it the same
|
||||
opportunity to drift.
|
||||
- Operations make the no-change case mechanical: zero ops → identical doc.
|
||||
- Operations are auditable: each refresh produces a log of exactly what
|
||||
changed, useful for debugging the LLM's behaviour and explaining diffs.
|
||||
|
||||
Failure modes are by design conservative: an operation list that fails to
|
||||
parse against the Pydantic schema, or an LLM that returns invalid ops, results
|
||||
in zero changes — the document stays as-is. The structure can only get better
|
||||
or stay the same per refresh, never get worse.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Annotated, Any, Literal, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from .structured_doc import (
|
||||
Block,
|
||||
Section,
|
||||
StructuredDocument,
|
||||
make_unique_id,
|
||||
slugify_heading,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Op payloads ---------------------------------------------------------------
|
||||
|
||||
|
||||
class _OpBase(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
class AppendBlockOp(_OpBase):
|
||||
"""Add a new block at the end of an existing section."""
|
||||
|
||||
op: Literal["append_block"] = "append_block"
|
||||
section_id: str
|
||||
block: Block
|
||||
|
||||
|
||||
class InsertBlockOp(_OpBase):
|
||||
"""Insert a new block at ``index`` in an existing section.
|
||||
|
||||
``index`` may equal ``len(section.blocks)`` (append) but not be greater.
|
||||
"""
|
||||
|
||||
op: Literal["insert_block"] = "insert_block"
|
||||
section_id: str
|
||||
index: int = Field(ge=0)
|
||||
block: Block
|
||||
|
||||
|
||||
class ReplaceBlockOp(_OpBase):
|
||||
"""Replace the block at ``index`` of an existing section."""
|
||||
|
||||
op: Literal["replace_block"] = "replace_block"
|
||||
section_id: str
|
||||
index: int = Field(ge=0)
|
||||
block: Block
|
||||
|
||||
|
||||
class RemoveBlockOp(_OpBase):
|
||||
"""Remove the block at ``index`` of an existing section."""
|
||||
|
||||
op: Literal["remove_block"] = "remove_block"
|
||||
section_id: str
|
||||
index: int = Field(ge=0)
|
||||
|
||||
|
||||
class AddSectionOp(_OpBase):
|
||||
"""Add a brand-new section.
|
||||
|
||||
``after_section_id`` is optional; when omitted the new section is appended
|
||||
at the end. ``new_id`` is optional; when omitted we slugify the heading
|
||||
and disambiguate against existing IDs.
|
||||
"""
|
||||
|
||||
op: Literal["add_section"] = "add_section"
|
||||
heading: str
|
||||
level: int = Field(default=2, ge=1, le=6)
|
||||
blocks: list[Block] = Field(default_factory=list)
|
||||
after_section_id: str | None = None
|
||||
new_id: str | None = None
|
||||
|
||||
|
||||
class RemoveSectionOp(_OpBase):
|
||||
"""Remove an entire section by id."""
|
||||
|
||||
op: Literal["remove_section"] = "remove_section"
|
||||
section_id: str
|
||||
|
||||
|
||||
class ReplaceSectionBlocksOp(_OpBase):
|
||||
"""Replace all blocks of a section in one go.
|
||||
|
||||
Used when most of a section's contents are stale and rebuilding it as a
|
||||
unit is clearer than emitting many block-level ops. The section's heading
|
||||
and id are preserved.
|
||||
"""
|
||||
|
||||
op: Literal["replace_section_blocks"] = "replace_section_blocks"
|
||||
section_id: str
|
||||
blocks: list[Block] = Field(default_factory=list)
|
||||
|
||||
|
||||
class RenameSectionOp(_OpBase):
|
||||
"""Rename a section's heading. The id is unchanged so future ops still resolve."""
|
||||
|
||||
op: Literal["rename_section"] = "rename_section"
|
||||
section_id: str
|
||||
new_heading: str
|
||||
|
||||
|
||||
Operation = Annotated[
|
||||
Union[
|
||||
AppendBlockOp,
|
||||
InsertBlockOp,
|
||||
ReplaceBlockOp,
|
||||
RemoveBlockOp,
|
||||
AddSectionOp,
|
||||
RemoveSectionOp,
|
||||
ReplaceSectionBlocksOp,
|
||||
RenameSectionOp,
|
||||
],
|
||||
Field(discriminator="op"),
|
||||
]
|
||||
|
||||
|
||||
class DeltaOperationList(BaseModel):
|
||||
"""Container for the operations produced by an LLM delta call."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
operations: list[Operation] = Field(default_factory=list)
|
||||
|
||||
|
||||
# Application ---------------------------------------------------------------
|
||||
|
||||
|
||||
class AppliedDelta(BaseModel):
|
||||
"""Outcome of applying a list of operations to a document."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
document: StructuredDocument
|
||||
applied: list[dict[str, Any]] = Field(default_factory=list)
|
||||
skipped: list[dict[str, Any]] = Field(default_factory=list)
|
||||
|
||||
@property
|
||||
def changed(self) -> bool:
|
||||
return len(self.applied) > 0
|
||||
|
||||
|
||||
def _op_summary(op: Operation) -> dict[str, Any]:
|
||||
"""Compact dict suitable for the audit trail."""
|
||||
data = op.model_dump()
|
||||
return {k: v for k, v in data.items() if k != "block" and k != "blocks"} | {
|
||||
"op": data["op"],
|
||||
}
|
||||
|
||||
|
||||
def apply_operations(
|
||||
doc: StructuredDocument,
|
||||
operations: list[Operation],
|
||||
) -> AppliedDelta:
|
||||
"""Apply a list of operations to a document, returning a new document.
|
||||
|
||||
The original document is never mutated. Invalid operations (unknown
|
||||
section, out-of-range index, name collision when adding a section) are
|
||||
skipped and recorded in ``skipped`` with a ``reason`` string.
|
||||
"""
|
||||
new_doc = doc.model_copy(deep=True)
|
||||
applied: list[dict[str, Any]] = []
|
||||
skipped: list[dict[str, Any]] = []
|
||||
|
||||
def skip(op: Operation, reason: str) -> None:
|
||||
entry = _op_summary(op)
|
||||
entry["reason"] = reason
|
||||
skipped.append(entry)
|
||||
logger.debug(f"[STRUCTURED_DELTA] skipping op {entry}")
|
||||
|
||||
for op in operations:
|
||||
if isinstance(op, AppendBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
section.blocks.append(op.block)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, InsertBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
if op.index > len(section.blocks):
|
||||
skip(
|
||||
op,
|
||||
f"index out of range: {op.index} > {len(section.blocks)}",
|
||||
)
|
||||
continue
|
||||
section.blocks.insert(op.index, op.block)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, ReplaceBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
if op.index >= len(section.blocks):
|
||||
skip(
|
||||
op,
|
||||
f"index out of range: {op.index} >= {len(section.blocks)}",
|
||||
)
|
||||
continue
|
||||
section.blocks[op.index] = op.block
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, RemoveBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
if op.index >= len(section.blocks):
|
||||
skip(
|
||||
op,
|
||||
f"index out of range: {op.index} >= {len(section.blocks)}",
|
||||
)
|
||||
continue
|
||||
section.blocks.pop(op.index)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, AddSectionOp):
|
||||
existing_ids = {s.id for s in new_doc.sections}
|
||||
base_id = op.new_id or slugify_heading(op.heading)
|
||||
section_id = make_unique_id(base_id, existing_ids)
|
||||
new_section = Section(
|
||||
id=section_id,
|
||||
heading=op.heading,
|
||||
level=op.level,
|
||||
blocks=list(op.blocks),
|
||||
)
|
||||
if op.after_section_id is None:
|
||||
new_doc.sections.append(new_section)
|
||||
else:
|
||||
idx = new_doc.section_index(op.after_section_id)
|
||||
if idx is None:
|
||||
skip(op, f"unknown after_section_id: {op.after_section_id}")
|
||||
continue
|
||||
new_doc.sections.insert(idx + 1, new_section)
|
||||
entry = _op_summary(op)
|
||||
entry["assigned_id"] = section_id
|
||||
applied.append(entry)
|
||||
continue
|
||||
|
||||
if isinstance(op, RemoveSectionOp):
|
||||
idx = new_doc.section_index(op.section_id)
|
||||
if idx is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
new_doc.sections.pop(idx)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, ReplaceSectionBlocksOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
section.blocks = list(op.blocks)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, RenameSectionOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
section.heading = op.new_heading
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
skip(op, f"unhandled op type: {type(op).__name__}") # pragma: no cover
|
||||
|
||||
return AppliedDelta(document=new_doc, applied=applied, skipped=skipped)
|
||||
@@ -18,6 +18,9 @@ _TIKTOKEN_ENCODING = tiktoken.get_encoding("cl100k_base")
|
||||
# The remainder covers the system prompt, question, bank context, and output tokens.
|
||||
_FINAL_PROMPT_CONTEXT_FRACTION = 0.8
|
||||
|
||||
_DEFAULT_ROLE = "You are a reflection agent that answers questions by reasoning over retrieved memories."
|
||||
_DEFAULT_FINAL_ROLE = "You are a thoughtful assistant that synthesizes answers from retrieved memories."
|
||||
|
||||
|
||||
def _extract_directive_rules(directives: list[dict[str, Any]]) -> list[str]:
|
||||
"""Extract directive rules as a list of strings."""
|
||||
@@ -133,7 +136,9 @@ def build_system_prompt_for_tools(
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
"You are a reflection agent that answers questions by reasoning over retrieved memories.",
|
||||
mission.strip() if mission else _DEFAULT_ROLE,
|
||||
"",
|
||||
"Answer the user's question by reasoning over retrieved memories.",
|
||||
"",
|
||||
]
|
||||
)
|
||||
@@ -369,7 +374,7 @@ def build_agent_prompt(
|
||||
output = entry["output"]
|
||||
# Format as proper JSON for LLM readability
|
||||
try:
|
||||
output_str = json.dumps(output, indent=2, default=str)
|
||||
output_str = json.dumps(output, indent=2, default=str, ensure_ascii=False)
|
||||
except (TypeError, ValueError):
|
||||
output_str = str(output)
|
||||
parts.append(f"\n### Call {i}: {tool}\n```json\n{output_str}\n```")
|
||||
@@ -444,7 +449,7 @@ def build_final_prompt(
|
||||
tool = entry["tool"]
|
||||
output = entry["output"]
|
||||
try:
|
||||
output_str = json.dumps(output, indent=2, default=str)
|
||||
output_str = json.dumps(output, indent=2, default=str, ensure_ascii=False)
|
||||
except (TypeError, ValueError):
|
||||
output_str = str(output)
|
||||
block = f"\n### From {tool}:\n```json\n{output_str}\n```"
|
||||
@@ -479,9 +484,9 @@ def build_final_prompt(
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
FINAL_SYSTEM_PROMPT = """CRITICAL: You MUST ONLY use information from retrieved tool results. NEVER make up names, people, events, or entities.
|
||||
_FINAL_SYSTEM_PROMPT_BASE = """CRITICAL: You MUST ONLY use information from retrieved tool results. NEVER make up names, people, events, or entities.
|
||||
|
||||
You are a thoughtful assistant that synthesizes answers from retrieved memories.
|
||||
{role_section}
|
||||
|
||||
Your approach:
|
||||
- Reason over the retrieved memories to answer the question
|
||||
@@ -508,3 +513,213 @@ CRITICAL: Output ONLY the final synthesized answer. Do NOT include:
|
||||
Just provide the direct answer with proper markdown formatting.
|
||||
|
||||
CRITICAL: This is a NON-CONVERSATIONAL system. NEVER ask follow-up questions, offer to search again, suggest alternatives, or end with anything like "Would you like me to..." or "Let me know if...". The user cannot reply. Your answer must be complete and self-contained."""
|
||||
|
||||
|
||||
def build_final_system_prompt(mission: str | None = None) -> str:
|
||||
"""Build the final synthesis system prompt, using mission as role when set."""
|
||||
role_section = mission.strip() if mission else _DEFAULT_FINAL_ROLE
|
||||
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section)
|
||||
|
||||
|
||||
# Backward-compatible constant for non-identity missions
|
||||
FINAL_SYSTEM_PROMPT = build_final_system_prompt()
|
||||
|
||||
|
||||
STRUCTURED_DELTA_SYSTEM_PROMPT = """You are integrating *new information* into an existing structured document.
|
||||
|
||||
You will be given:
|
||||
1. TOPIC — the question this document answers. Content that does not help
|
||||
answer this question is OFF-TOPIC and should be removed.
|
||||
2. CURRENT DOCUMENT (JSON) — the existing structured mental model. Each section
|
||||
has a stable ``id``, a ``heading``, a ``level`` (1..6), and an ordered list
|
||||
of ``blocks``. Blocks are typed: ``paragraph``, ``bullet_list``,
|
||||
``ordered_list``, or ``code``.
|
||||
3. NEW INFORMATION SYNTHESIS (markdown) — a synthesis showing how the new facts
|
||||
relate to the document's topic. Use it to understand context and relevance,
|
||||
but do NOT copy its formatting or wording wholesale.
|
||||
4. SUPPORTING FACTS — observations and facts created since the last refresh.
|
||||
These are genuinely new — they were NOT available when the current document
|
||||
was written.
|
||||
|
||||
Your task: output a JSON object ``{"operations": [...]}``. Applied to CURRENT
|
||||
DOCUMENT, the operations must produce a document that best answers the TOPIC
|
||||
by integrating the new facts.
|
||||
|
||||
RULES
|
||||
- These facts are NEW since the last refresh. The existing document already
|
||||
captures all prior information from earlier refreshes. Your job is to
|
||||
integrate the new facts into the existing document.
|
||||
- **Preserve existing content**: The current document was built from prior facts
|
||||
that you cannot see. Do NOT remove or replace existing sections just because
|
||||
the new facts do not reference them. Only remove content when the new facts
|
||||
explicitly contradict or supersede it.
|
||||
- **Merge overlapping topics**: When new facts cover topics that overlap with
|
||||
existing sections, merge the new information INTO the existing section
|
||||
rather than creating duplicates. When new facts provide more specific or
|
||||
authoritative guidance on a topic already covered generically, update the
|
||||
existing content to reflect the more specific guidance.
|
||||
- **Preserve examples**: Concrete examples, before/after pairs, sample sentences,
|
||||
and illustrative ✅/❌ comparisons are MORE valuable than abstract rules.
|
||||
When facts contain examples, include them. Never drop an example to make
|
||||
room for an abstract restatement of the same point.
|
||||
- Operations target sections by ``section_id`` (use the ``id`` field of the
|
||||
section in CURRENT DOCUMENT, NOT the heading). Block operations target
|
||||
blocks by ``index`` (0-based, against the section's current block list).
|
||||
- **Add** new content with ``append_block``, ``insert_block``, or ``add_section``
|
||||
when facts introduce information not yet covered. Prefer extending an
|
||||
existing section over creating a new one.
|
||||
- **Update** existing content with ``replace_block`` or ``replace_section_blocks``
|
||||
when new facts provide corrections, updates, or more specific information
|
||||
about topics already in the document.
|
||||
- **Remove** content with ``remove_block`` or ``remove_section`` ONLY when
|
||||
the new facts explicitly contradict or supersede it.
|
||||
- NEVER emit operations whose only effect is to reword unchanged content.
|
||||
- NEVER emit operations to "normalize" formatting (numbered → bulleted, casing
|
||||
changes, paragraph → list, etc).
|
||||
- Every operation MUST be justifiable by a specific fact in SUPPORTING FACTS.
|
||||
- Output ``{"operations": []}`` only if the new facts are already reflected
|
||||
in the document (e.g., from a concurrent update).
|
||||
|
||||
ALLOWED OPERATIONS (each line shows the JSON shape)
|
||||
- ``{"op": "append_block", "section_id": "...", "block": {...}}``
|
||||
- ``{"op": "insert_block", "section_id": "...", "index": N, "block": {...}}``
|
||||
- ``{"op": "replace_block", "section_id": "...", "index": N, "block": {...}}``
|
||||
- ``{"op": "remove_block", "section_id": "...", "index": N}``
|
||||
- ``{"op": "add_section", "heading": "...", "level": 2, "blocks": [...], "after_section_id": "..."}``
|
||||
- ``{"op": "remove_section", "section_id": "..."}``
|
||||
- ``{"op": "replace_section_blocks", "section_id": "...", "blocks": [...]}``
|
||||
- ``{"op": "rename_section", "section_id": "...", "new_heading": "..."}``
|
||||
|
||||
Block shapes
|
||||
- ``{"type": "paragraph", "text": "..."}``
|
||||
- ``{"type": "bullet_list", "items": ["...", "..."]}``
|
||||
- ``{"type": "ordered_list", "items": ["...", "..."]}``
|
||||
- ``{"type": "code", "language": "json", "text": "..."}``
|
||||
|
||||
OUTPUT FORMAT
|
||||
Return ONLY a single JSON object on its own, with no prose before or after,
|
||||
no markdown code fences, no commentary. The object must have exactly one
|
||||
top-level key, ``operations``, whose value is an array of operation objects
|
||||
(empty array when nothing changes).
|
||||
|
||||
Examples
|
||||
- No changes needed → ``{"operations": []}``
|
||||
- Add one bullet to an existing "Members" section →
|
||||
``{"operations": [{"op": "append_block", "section_id": "members",
|
||||
"block": {"type": "bullet_list", "items": ["Carol — junior engineer"]}}]}``
|
||||
- Replace a paragraph that has been corrected by new facts →
|
||||
``{"operations": [{"op": "replace_block", "section_id": "overview",
|
||||
"index": 0, "block": {"type": "paragraph", "text": "Updated summary."}}]}``
|
||||
- Remove an obsolete block →
|
||||
``{"operations": [{"op": "remove_block", "section_id": "status", "index": 2}]}``"""
|
||||
|
||||
|
||||
def build_structured_delta_prompt(
|
||||
*,
|
||||
current_document_json: str,
|
||||
candidate_markdown: str,
|
||||
supporting_facts: list[dict[str, Any]],
|
||||
source_query: str,
|
||||
max_output_tokens: int | None = None,
|
||||
) -> str:
|
||||
"""Build the user prompt for a structured-delta mental model refresh.
|
||||
|
||||
The LLM's job is to emit operations against ``current_document_json``;
|
||||
the surrounding ``candidate_markdown`` and ``supporting_facts`` are
|
||||
references for *what new information exists*, not templates to mimic.
|
||||
|
||||
``max_output_tokens`` is surfaced in the prompt so the model can keep its
|
||||
op list within the provider's response cap. The actual cap is enforced by
|
||||
the caller; this is just an advisory anchor — without it the model often
|
||||
returns op lists whose JSON gets truncated mid-string.
|
||||
"""
|
||||
fact_lines: list[str] = []
|
||||
for f in supporting_facts:
|
||||
fid = f.get("id", "")
|
||||
text = (f.get("text") or "").strip().replace("\n", " ")
|
||||
ftype = f.get("type", "")
|
||||
fact_lines.append(f"- [{ftype}:{fid}] {text}")
|
||||
facts_block = "\n".join(fact_lines) if fact_lines else "(no supporting facts retrieved)"
|
||||
|
||||
budget_hint = ""
|
||||
if max_output_tokens is not None:
|
||||
budget_hint = (
|
||||
f"\n\n## Output budget\n"
|
||||
f"Your JSON response must fit within ~{max_output_tokens} tokens. If you "
|
||||
"would need more than this to express every change, prefer the highest-"
|
||||
"leverage edits first (a few ``replace_section_blocks`` ops over many "
|
||||
"block-level ops) so the response always parses as valid JSON."
|
||||
)
|
||||
|
||||
return (
|
||||
f"## Topic\n{source_query}\n\n"
|
||||
f"## CURRENT DOCUMENT (apply ops to this; reference section ids as listed)\n"
|
||||
f"```json\n{current_document_json}\n```\n\n"
|
||||
f"## NEW INFORMATION SYNTHESIS (context for how new facts relate to the topic)\n"
|
||||
f"```markdown\n{candidate_markdown}\n```\n\n"
|
||||
f"## SUPPORTING FACTS (new since last refresh — integrate these)\n{facts_block}"
|
||||
f"{budget_hint}\n\n"
|
||||
"## Task\n"
|
||||
"Output a JSON object matching the operations schema. Integrate the new "
|
||||
"supporting facts into CURRENT DOCUMENT. Add, update, or remove content "
|
||||
"as needed. Preserve unchanged sections and blocks by not mentioning them."
|
||||
)
|
||||
|
||||
|
||||
DELTA_SYSTEM_PROMPT = """You are performing a surgical delta update to an existing mental model document.
|
||||
|
||||
You will be given:
|
||||
1. CURRENT DOCUMENT: the existing mental model content (markdown).
|
||||
2. CANDIDATE UPDATE: a freshly generated synthesis based on the latest retrieved memories.
|
||||
3. SUPPORTING FACTS: the observations and facts that support the CANDIDATE UPDATE.
|
||||
|
||||
Your task: produce an updated version of the CURRENT DOCUMENT that reflects the new reality, with the MINIMUM possible changes.
|
||||
|
||||
ABSOLUTE RULES:
|
||||
- Preserve unchanged content BYTE-FOR-BYTE. If a sentence, heading, bullet, code block, or section is still accurate according to the CANDIDATE UPDATE and SUPPORTING FACTS, copy it verbatim — same wording, same punctuation, same whitespace, same markdown structure.
|
||||
- Do NOT reformat, rephrase, or re-style content that is still accurate. No "light edits for clarity", no reordering for flow, no synonym swaps.
|
||||
- Remove content that is contradicted by the CANDIDATE UPDATE or SUPPORTING FACTS (stale content).
|
||||
- Add new content ONLY when the SUPPORTING FACTS contain information not already in the CURRENT DOCUMENT.
|
||||
- When adding new content, prefer appending to an existing relevant section. Creating a new section is acceptable when the new information does not fit any existing section.
|
||||
- When creating a new section, match the heading style, tone, and formatting conventions used in the CURRENT DOCUMENT.
|
||||
- Every assertion in your output MUST be grounded in either (a) the CURRENT DOCUMENT (preserved) or (b) the SUPPORTING FACTS. Never introduce outside knowledge.
|
||||
- If nothing in the SUPPORTING FACTS contradicts or extends the CURRENT DOCUMENT, return the CURRENT DOCUMENT UNCHANGED, character for character.
|
||||
|
||||
OUTPUT FORMAT:
|
||||
- Output ONLY the updated markdown document. No preamble, no explanation, no diff markers, no commentary.
|
||||
- Do not wrap the output in code fences unless the CURRENT DOCUMENT itself was entirely a code fence."""
|
||||
|
||||
|
||||
def build_delta_prompt(
|
||||
*,
|
||||
current_content: str,
|
||||
candidate_content: str,
|
||||
supporting_facts: list[dict[str, Any]],
|
||||
source_query: str,
|
||||
) -> str:
|
||||
"""Build the user prompt for a delta-mode mental model refresh.
|
||||
|
||||
Args:
|
||||
current_content: The existing mental model content (to preserve as much as possible).
|
||||
candidate_content: Fresh synthesis from the reflect agent reflecting new reality.
|
||||
supporting_facts: Flat list of fact dicts (id, text, type) supporting the candidate.
|
||||
source_query: The mental model's source query, for topical framing.
|
||||
"""
|
||||
fact_lines: list[str] = []
|
||||
for f in supporting_facts:
|
||||
fid = f.get("id", "")
|
||||
text = (f.get("text") or "").strip().replace("\n", " ")
|
||||
ftype = f.get("type", "")
|
||||
fact_lines.append(f"- [{ftype}:{fid}] {text}")
|
||||
facts_block = "\n".join(fact_lines) if fact_lines else "(no supporting facts retrieved)"
|
||||
|
||||
return (
|
||||
f"## Topic\n{source_query}\n\n"
|
||||
f"## CURRENT DOCUMENT\n```markdown\n{current_content}\n```\n\n"
|
||||
f"## CANDIDATE UPDATE\n```markdown\n{candidate_content}\n```\n\n"
|
||||
f"## SUPPORTING FACTS\n{facts_block}\n\n"
|
||||
"## Task\n"
|
||||
"Produce the updated mental model document by applying the minimum necessary changes "
|
||||
"to CURRENT DOCUMENT so that it reflects CANDIDATE UPDATE and SUPPORTING FACTS. "
|
||||
"Preserve unchanged content byte-for-byte. Output only the final markdown."
|
||||
)
|
||||
|
||||
@@ -0,0 +1,301 @@
|
||||
"""Structured representation of a mental model document.
|
||||
|
||||
Why this exists
|
||||
---------------
|
||||
Storing mental models as raw markdown forces every refresh to round-trip prose
|
||||
through an LLM, which then drifts on stylistic details (numbered vs bulleted
|
||||
lists, casing, separator lines, paraphrasing) even when instructed to preserve
|
||||
content byte-for-byte. The intrinsic mechanism of an LLM is to *generate* the
|
||||
next token from a gestalt of the input — not to copy tokens verbatim — so any
|
||||
"preserve unchanged content" instruction is fundamentally a soft constraint.
|
||||
|
||||
The fix is to give the LLM no opportunity to drift on unchanged content. We
|
||||
keep an authoritative structured representation of the document; the markdown
|
||||
shown to users is a deterministic render of that structure. Delta refreshes
|
||||
emit *operations* against the structure (see ``delta_ops.py``); sections and
|
||||
blocks not mentioned by any operation are physically untouched.
|
||||
|
||||
Schema (v1)
|
||||
-----------
|
||||
A document is an ordered list of ``Section``s. Each section has:
|
||||
- ``id`` : stable slug derived from ``heading`` (used as the operation
|
||||
target across refreshes; surviving renames is a separate
|
||||
concern handled by an explicit ``rename`` op).
|
||||
- ``heading``: the markdown heading text (without the ``#`` prefix).
|
||||
- ``level`` : 1 (``#``) … 6 (``######``). Default 2.
|
||||
- ``blocks``: ordered list of typed blocks — paragraph, bullet_list,
|
||||
ordered_list, code.
|
||||
|
||||
The schema is intentionally narrow: it covers what real mental-model documents
|
||||
actually contain (the kind a coding agent writes for itself or a user writes as
|
||||
a "skill" doc). Tables, images, and raw HTML are out of scope until needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Annotated, Literal, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
# Blocks ---------------------------------------------------------------------
|
||||
|
||||
|
||||
class ParagraphBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["paragraph"] = "paragraph"
|
||||
text: str
|
||||
|
||||
|
||||
class BulletListBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["bullet_list"] = "bullet_list"
|
||||
items: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class OrderedListBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["ordered_list"] = "ordered_list"
|
||||
items: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class CodeBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["code"] = "code"
|
||||
language: str = ""
|
||||
text: str
|
||||
|
||||
|
||||
Block = Annotated[
|
||||
Union[ParagraphBlock, BulletListBlock, OrderedListBlock, CodeBlock],
|
||||
Field(discriminator="type"),
|
||||
]
|
||||
|
||||
|
||||
# Section / Document ---------------------------------------------------------
|
||||
|
||||
|
||||
class Section(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
id: str
|
||||
heading: str
|
||||
level: int = Field(default=2, ge=1, le=6)
|
||||
blocks: list[Block] = Field(default_factory=list)
|
||||
|
||||
|
||||
class StructuredDocument(BaseModel):
|
||||
"""Top-level structured representation of a mental model."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
version: Literal[1] = 1
|
||||
sections: list[Section] = Field(default_factory=list)
|
||||
|
||||
def section_by_id(self, section_id: str) -> Section | None:
|
||||
for s in self.sections:
|
||||
if s.id == section_id:
|
||||
return s
|
||||
return None
|
||||
|
||||
def section_index(self, section_id: str) -> int | None:
|
||||
for i, s in enumerate(self.sections):
|
||||
if s.id == section_id:
|
||||
return i
|
||||
return None
|
||||
|
||||
|
||||
# Slug helpers ---------------------------------------------------------------
|
||||
|
||||
_SLUG_RX = re.compile(r"[^a-z0-9]+")
|
||||
|
||||
|
||||
def slugify_heading(heading: str) -> str:
|
||||
"""Stable, deterministic slug from a heading.
|
||||
|
||||
"Stop Conditions" -> "stop-conditions"
|
||||
"Inputs and Context" -> "inputs-and-context"
|
||||
"""
|
||||
slug = _SLUG_RX.sub("-", heading.strip().lower()).strip("-")
|
||||
return slug or "section"
|
||||
|
||||
|
||||
def make_unique_id(base: str, existing: set[str]) -> str:
|
||||
"""Disambiguate by appending -2, -3, … if the slug is already in use."""
|
||||
if base not in existing:
|
||||
return base
|
||||
i = 2
|
||||
while f"{base}-{i}" in existing:
|
||||
i += 1
|
||||
return f"{base}-{i}"
|
||||
|
||||
|
||||
# Renderer -------------------------------------------------------------------
|
||||
|
||||
|
||||
def render_block(block: Block) -> str:
|
||||
"""Render a single block to markdown. No trailing newline."""
|
||||
if isinstance(block, ParagraphBlock):
|
||||
return block.text.rstrip()
|
||||
if isinstance(block, BulletListBlock):
|
||||
return "\n".join(f"- {item.rstrip()}" for item in block.items)
|
||||
if isinstance(block, OrderedListBlock):
|
||||
return "\n".join(f"{i + 1}. {item.rstrip()}" for i, item in enumerate(block.items))
|
||||
if isinstance(block, CodeBlock):
|
||||
fence_lang = block.language or ""
|
||||
return f"```{fence_lang}\n{block.text}\n```"
|
||||
raise TypeError(f"Unknown block type: {type(block)!r}")
|
||||
|
||||
|
||||
def render_section(section: Section) -> str:
|
||||
"""Render a section: heading + blank line + blocks separated by blank lines."""
|
||||
parts = ["#" * section.level + " " + section.heading.strip()]
|
||||
for block in section.blocks:
|
||||
parts.append("") # blank line before each block
|
||||
parts.append(render_block(block))
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def render_document(doc: StructuredDocument) -> str:
|
||||
"""Render the whole document. Sections separated by a single blank line.
|
||||
|
||||
The output is byte-stable: same structured input always produces the same
|
||||
markdown, modulo the inherent ordering of sections/blocks/items.
|
||||
"""
|
||||
if not doc.sections:
|
||||
return ""
|
||||
return "\n\n".join(render_section(s) for s in doc.sections) + "\n"
|
||||
|
||||
|
||||
# Parser ---------------------------------------------------------------------
|
||||
#
|
||||
# The parser is intentionally lenient: it accepts the markdown produced by
|
||||
# our own renderer (round-trip-safe) and the markdown an LLM tends to produce
|
||||
# for mental-model documents. It is *not* a general CommonMark parser — it
|
||||
# does not need to be. When it cannot classify a block it falls back to a
|
||||
# paragraph so that no content is silently dropped.
|
||||
|
||||
_HEADING_RX = re.compile(r"^(#{1,6})\s+(.+?)\s*$")
|
||||
_BULLET_RX = re.compile(r"^\s*[-*+]\s+(.*)$")
|
||||
_ORDERED_RX = re.compile(r"^\s*\d+[.)]\s+(.*)$")
|
||||
_FENCE_RX = re.compile(r"^```([A-Za-z0-9_+-]*)\s*$")
|
||||
|
||||
|
||||
def _strip_separators(lines: list[str]) -> list[str]:
|
||||
"""Drop horizontal-rule lines (`---`, `***`) used as section separators.
|
||||
|
||||
Our renderer never emits these, but LLM output frequently includes them
|
||||
between sections; treating them as blank lines avoids parsing them as
|
||||
paragraphs.
|
||||
"""
|
||||
return ["" if re.fullmatch(r"\s*([-*_])\1{2,}\s*", line) else line for line in lines]
|
||||
|
||||
|
||||
def _split_blocks(lines: list[str]) -> list[list[str]]:
|
||||
"""Group consecutive non-blank lines into block chunks."""
|
||||
chunks: list[list[str]] = []
|
||||
current: list[str] = []
|
||||
in_fence = False
|
||||
for line in lines:
|
||||
if _FENCE_RX.match(line):
|
||||
current.append(line)
|
||||
in_fence = not in_fence
|
||||
continue
|
||||
if in_fence:
|
||||
current.append(line)
|
||||
continue
|
||||
if line.strip() == "":
|
||||
if current:
|
||||
chunks.append(current)
|
||||
current = []
|
||||
else:
|
||||
current.append(line)
|
||||
if current:
|
||||
chunks.append(current)
|
||||
return chunks
|
||||
|
||||
|
||||
def _parse_block(chunk: list[str]) -> Block:
|
||||
"""Parse a single non-empty chunk into a block."""
|
||||
if chunk and _FENCE_RX.match(chunk[0]):
|
||||
m = _FENCE_RX.match(chunk[0])
|
||||
lang = m.group(1) if m else ""
|
||||
body_lines = chunk[1:]
|
||||
if body_lines and _FENCE_RX.match(body_lines[-1]):
|
||||
body_lines = body_lines[:-1]
|
||||
return CodeBlock(language=lang, text="\n".join(body_lines))
|
||||
|
||||
if all(_BULLET_RX.match(line) for line in chunk):
|
||||
items = []
|
||||
for line in chunk:
|
||||
m = _BULLET_RX.match(line)
|
||||
assert m is not None
|
||||
items.append(m.group(1).strip())
|
||||
return BulletListBlock(items=items)
|
||||
|
||||
if all(_ORDERED_RX.match(line) for line in chunk):
|
||||
items = []
|
||||
for line in chunk:
|
||||
m = _ORDERED_RX.match(line)
|
||||
assert m is not None
|
||||
items.append(m.group(1).strip())
|
||||
return OrderedListBlock(items=items)
|
||||
|
||||
return ParagraphBlock(text=" ".join(line.strip() for line in chunk).strip())
|
||||
|
||||
|
||||
def parse_markdown(markdown: str) -> StructuredDocument:
|
||||
"""Best-effort parse of a markdown document into the structured schema.
|
||||
|
||||
Sections are introduced by ATX headings (``#``..``######``). Anything
|
||||
before the first heading is wrapped into an implicit "Overview" section
|
||||
so we never silently drop user content. Section IDs are unique slugs of
|
||||
their headings.
|
||||
"""
|
||||
raw_lines = (markdown or "").splitlines()
|
||||
lines = _strip_separators(raw_lines)
|
||||
|
||||
sections: list[Section] = []
|
||||
used_ids: set[str] = set()
|
||||
pending: list[str] = []
|
||||
current: Section | None = None
|
||||
|
||||
def flush_pending_into(section: Section) -> None:
|
||||
if not pending:
|
||||
return
|
||||
for chunk in _split_blocks(pending):
|
||||
section.blocks.append(_parse_block(chunk))
|
||||
pending.clear()
|
||||
|
||||
for line in lines:
|
||||
m = _HEADING_RX.match(line)
|
||||
if m:
|
||||
if current is not None:
|
||||
flush_pending_into(current)
|
||||
sections.append(current)
|
||||
elif pending:
|
||||
# Content before the first heading: wrap in implicit section.
|
||||
base = "overview"
|
||||
section_id = make_unique_id(base, used_ids)
|
||||
used_ids.add(section_id)
|
||||
implicit = Section(id=section_id, heading="Overview", level=2)
|
||||
flush_pending_into(implicit)
|
||||
sections.append(implicit)
|
||||
level = len(m.group(1))
|
||||
heading = m.group(2).strip()
|
||||
section_id = make_unique_id(slugify_heading(heading), used_ids)
|
||||
used_ids.add(section_id)
|
||||
current = Section(id=section_id, heading=heading, level=level)
|
||||
else:
|
||||
pending.append(line)
|
||||
|
||||
if current is not None:
|
||||
flush_pending_into(current)
|
||||
sections.append(current)
|
||||
elif pending:
|
||||
base = "overview"
|
||||
section_id = make_unique_id(base, used_ids)
|
||||
used_ids.add(section_id)
|
||||
implicit = Section(id=section_id, heading="Overview", level=2)
|
||||
flush_pending_into(implicit)
|
||||
sections.append(implicit)
|
||||
|
||||
return StructuredDocument(sections=sections)
|
||||
@@ -9,6 +9,7 @@ Implements hierarchical retrieval:
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
from dataclasses import replace
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
@@ -22,6 +23,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def tool_search_mental_models(
|
||||
memory_engine: "MemoryEngine",
|
||||
conn: "Connection",
|
||||
bank_id: str,
|
||||
query: str,
|
||||
@@ -31,7 +33,6 @@ async def tool_search_mental_models(
|
||||
tags_match: str = "any",
|
||||
tag_groups: "list | None" = None,
|
||||
exclude_ids: list[str] | None = None,
|
||||
pending_consolidation: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Search user-curated mental models by semantic similarity.
|
||||
@@ -81,7 +82,7 @@ async def tool_search_mental_models(
|
||||
f"""
|
||||
SELECT
|
||||
id, name, content,
|
||||
tags, created_at, last_refreshed_at,
|
||||
tags, created_at, last_refreshed_at, trigger,
|
||||
1 - (embedding <=> $2::vector) as relevance
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1 AND embedding IS NOT NULL {filters}
|
||||
@@ -98,10 +99,9 @@ async def tool_search_mental_models(
|
||||
if last_refreshed_at and last_refreshed_at.tzinfo is None:
|
||||
last_refreshed_at = last_refreshed_at.replace(tzinfo=timezone.utc)
|
||||
|
||||
# A mental model is stale when there are memories that haven't been consolidated yet —
|
||||
# the same signal used for observations staleness.
|
||||
is_stale = pending_consolidation > 0
|
||||
staleness_reason = f"{pending_consolidation} memories pending consolidation" if is_stale else None
|
||||
# Per-MM staleness: new in-scope memories since last refresh (includes pending).
|
||||
is_stale = await memory_engine.compute_mental_model_is_stale(conn, bank_id, row)
|
||||
staleness_reason = "new in-scope memories ingested since last refresh" if is_stale else None
|
||||
|
||||
mental_models.append(
|
||||
{
|
||||
@@ -134,9 +134,12 @@ async def tool_search_observations(
|
||||
tag_groups: "list | None" = None,
|
||||
last_consolidated_at: datetime | None = None,
|
||||
pending_consolidation: int = 0,
|
||||
source_facts_max_tokens: int = -1,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Search consolidated observations using recall with include_source_facts.
|
||||
Search consolidated observations using recall.
|
||||
|
||||
Observations are auto-generated from memories. Returns freshness info
|
||||
so the agent knows if it should also verify with recall().
|
||||
@@ -151,24 +154,37 @@ async def tool_search_observations(
|
||||
tags_match: How to match tags - "any" (OR), "all" (AND)
|
||||
last_consolidated_at: When consolidation last ran (for staleness check)
|
||||
pending_consolidation: Number of memories waiting to be consolidated
|
||||
source_facts_max_tokens: Token budget for source facts (-1 = disabled, 0+ = enabled with limit)
|
||||
|
||||
Returns:
|
||||
Dict with matching observations including freshness info and source memories
|
||||
"""
|
||||
include_source_facts = source_facts_max_tokens != -1
|
||||
recall_kwargs: dict[str, Any] = {}
|
||||
if include_source_facts and source_facts_max_tokens > 0:
|
||||
recall_kwargs["max_source_facts_tokens"] = source_facts_max_tokens
|
||||
|
||||
# Use an internal request context so this recall is not billed as a
|
||||
# user-facing operation. The reflect caller is already billed for the
|
||||
# overall reflect operation; double-billing the sub-recalls would
|
||||
# overcharge the customer.
|
||||
internal_ctx = replace(request_context, internal=True)
|
||||
result = await memory_engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
fact_type=["observation"],
|
||||
max_tokens=max_tokens,
|
||||
enable_trace=False,
|
||||
request_context=request_context,
|
||||
request_context=internal_ctx,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
include_source_facts=True,
|
||||
max_source_facts_tokens=-1, # No token limit — include all source facts
|
||||
include_source_facts=include_source_facts,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
_connection_budget=1,
|
||||
_quiet=True,
|
||||
**recall_kwargs,
|
||||
)
|
||||
|
||||
is_stale = pending_consolidation > 0
|
||||
@@ -200,6 +216,10 @@ async def tool_recall(
|
||||
tag_groups: "list | None" = None,
|
||||
connection_budget: int = 1,
|
||||
max_chunk_tokens: int = 1000,
|
||||
fact_types: list[str] | None = None,
|
||||
include_chunks: bool = True,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Search memories using TEMPR retrieval.
|
||||
@@ -216,22 +236,28 @@ async def tool_recall(
|
||||
tags: Filter by tags (includes untagged memories)
|
||||
tags_match: How to match tags - "any" (OR), "all" (AND), or "exact"
|
||||
connection_budget: Max DB connections for this recall (default 1 for internal ops)
|
||||
max_chunk_tokens: Maximum tokens for raw source chunk text (default 1000, always included)
|
||||
max_chunk_tokens: Maximum tokens for raw source chunk text (default 1000)
|
||||
fact_types: Optional filter for fact types to retrieve. Defaults to ["experience", "world"].
|
||||
include_chunks: Whether to fetch raw chunk text alongside facts (default True).
|
||||
|
||||
Returns:
|
||||
Dict with list of matching memories including raw chunk text
|
||||
Dict with list of matching memories including raw chunk text (when include_chunks)
|
||||
"""
|
||||
include_chunks = True
|
||||
# Only world/experience are valid for raw recall (observation is handled by search_observations)
|
||||
recall_fact_type = [ft for ft in (fact_types or ["experience", "world"]) if ft in ("world", "experience")]
|
||||
internal_ctx = replace(request_context, internal=True)
|
||||
result = await memory_engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
fact_type=["experience", "world"],
|
||||
fact_type=recall_fact_type,
|
||||
max_tokens=max_tokens,
|
||||
enable_trace=False,
|
||||
request_context=request_context,
|
||||
request_context=internal_ctx,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
_connection_budget=connection_budget,
|
||||
_quiet=True, # Suppress logging for internal operations
|
||||
include_chunks=include_chunks,
|
||||
|
||||
@@ -227,7 +227,12 @@ def _build_done_tool_with_directives(directive_rules: list[str]) -> dict:
|
||||
}
|
||||
|
||||
|
||||
def get_reflect_tools(directive_rules: list[str] | None = None) -> list[dict]:
|
||||
def get_reflect_tools(
|
||||
directive_rules: list[str] | None = None,
|
||||
include_mental_models: bool = True,
|
||||
include_observations: bool = True,
|
||||
include_recall: bool = True,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Get the list of tools for the reflect agent.
|
||||
|
||||
@@ -239,16 +244,23 @@ def get_reflect_tools(directive_rules: list[str] | None = None) -> list[dict]:
|
||||
Args:
|
||||
directive_rules: Optional list of directive rule strings. If provided,
|
||||
the done() tool will require directive compliance confirmation.
|
||||
include_mental_models: Whether to include the search_mental_models tool.
|
||||
include_observations: Whether to include the search_observations tool.
|
||||
include_recall: Whether to include the recall tool.
|
||||
|
||||
Returns:
|
||||
List of tool definitions in OpenAI format
|
||||
"""
|
||||
tools = [
|
||||
TOOL_SEARCH_MENTAL_MODELS,
|
||||
TOOL_SEARCH_OBSERVATIONS,
|
||||
TOOL_RECALL,
|
||||
TOOL_EXPAND,
|
||||
]
|
||||
tools = []
|
||||
|
||||
if include_mental_models:
|
||||
tools.append(TOOL_SEARCH_MENTAL_MODELS)
|
||||
if include_observations:
|
||||
tools.append(TOOL_SEARCH_OBSERVATIONS)
|
||||
if include_recall:
|
||||
tools.append(TOOL_RECALL)
|
||||
|
||||
tools.append(TOOL_EXPAND)
|
||||
|
||||
# Use directive-aware done tool if directives are present
|
||||
if directive_rules:
|
||||
|
||||
@@ -8,9 +8,8 @@ API stability even if internal models change.
|
||||
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
# Valid fact types for recall operations (excludes 'opinion' which is deprecated)
|
||||
VALID_RECALL_FACT_TYPES = frozenset(["world", "experience", "observation"])
|
||||
|
||||
|
||||
@@ -159,6 +158,19 @@ class MemoryFact(BaseModel):
|
||||
mentioned_at: str | None = Field(None, description="ISO format date when the fact was mentioned/learned")
|
||||
document_id: str | None = Field(None, description="ID of the document this memory belongs to")
|
||||
metadata: dict[str, str] | None = Field(None, description="User-defined metadata")
|
||||
|
||||
@field_validator("metadata", mode="before")
|
||||
@classmethod
|
||||
def parse_metadata(cls, v: Any) -> dict[str, str] | None:
|
||||
"""Parse metadata from JSON string if needed (asyncpg may return JSONB as str)."""
|
||||
if v is None:
|
||||
return None
|
||||
if isinstance(v, str):
|
||||
import json
|
||||
|
||||
return json.loads(v)
|
||||
return v
|
||||
|
||||
chunk_id: str | None = Field(
|
||||
None, description="ID of the chunk this fact was extracted from (format: bank_id_document_id_chunk_index)"
|
||||
)
|
||||
|
||||
@@ -10,32 +10,47 @@ from typing import TypedDict
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table, get_current_schema
|
||||
from ..response_models import DispositionTraits
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Fact types that get per-bank partial HNSW indexes, mapped to their 4-char index suffix.
|
||||
_HNSW_FACT_TYPES: dict[str, str] = {
|
||||
# Fact types that get per-bank partial vector indexes, mapped to their 4-char index suffix.
|
||||
_BANK_INDEX_FACT_TYPES: dict[str, str] = {
|
||||
"world": "worl",
|
||||
"experience": "expr",
|
||||
"observation": "obsv",
|
||||
}
|
||||
|
||||
|
||||
def _hnsw_index_name(ft: str, internal_id: str) -> str:
|
||||
"""Deterministic, schema-safe HNSW index name for a (bank, fact_type) pair.
|
||||
def _bank_index_name(ft: str, internal_id: str) -> str:
|
||||
"""Deterministic, schema-safe vector index name for a (bank, fact_type) pair.
|
||||
|
||||
Uses the first 16 hex chars of internal_id (8 bytes of entropy) — unique
|
||||
enough in practice, fits comfortably within PostgreSQL's 63-char identifier limit.
|
||||
"""
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
return f"idx_mu_emb_{_HNSW_FACT_TYPES[ft]}_{uid}"
|
||||
return f"idx_mu_emb_{_BANK_INDEX_FACT_TYPES[ft]}_{uid}"
|
||||
|
||||
|
||||
async def create_bank_hnsw_indexes(conn, bank_id: str, internal_id: str) -> None:
|
||||
"""Create per-(bank, fact_type) partial HNSW indexes for a newly created bank.
|
||||
def _vector_index_clause() -> str:
|
||||
"""Return the USING clause for vector index creation based on the configured extension."""
|
||||
ext = get_config().vector_extension
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
elif ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
else: # pgvector (default)
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
|
||||
async def create_bank_vector_indexes(conn, bank_id: str, internal_id: str) -> None:
|
||||
"""Create per-(bank, fact_type) partial vector indexes for a newly created bank.
|
||||
|
||||
Respects the HINDSIGHT_API_VECTOR_EXTENSION config to use the appropriate
|
||||
index type (HNSW for pgvector, DiskANN for pgvectorscale, vchordrq for vchord).
|
||||
|
||||
Called immediately after the bank row is first inserted. Safe on empty banks
|
||||
(index build is instant). Idempotent via CREATE INDEX IF NOT EXISTS.
|
||||
@@ -43,24 +58,25 @@ async def create_bank_hnsw_indexes(conn, bank_id: str, internal_id: str) -> None
|
||||
"""
|
||||
table = fq_table("memory_units")
|
||||
escaped = bank_id.replace("'", "''")
|
||||
for ft in _HNSW_FACT_TYPES:
|
||||
idx = _hnsw_index_name(ft, internal_id)
|
||||
using_clause = _vector_index_clause()
|
||||
for ft in _BANK_INDEX_FACT_TYPES:
|
||||
idx = _bank_index_name(ft, internal_id)
|
||||
await conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx} "
|
||||
f"ON {table} USING hnsw (embedding vector_cosine_ops) "
|
||||
f"ON {table} {using_clause} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped}'"
|
||||
)
|
||||
|
||||
|
||||
async def drop_bank_hnsw_indexes(conn, internal_id: str) -> None:
|
||||
"""Drop per-(bank, fact_type) partial HNSW indexes for a bank being deleted.
|
||||
async def drop_bank_vector_indexes(conn, internal_id: str) -> None:
|
||||
"""Drop per-(bank, fact_type) partial vector indexes for a bank being deleted.
|
||||
|
||||
Called before the bank row is deleted so internal_id is still known.
|
||||
Idempotent via DROP INDEX IF EXISTS.
|
||||
"""
|
||||
schema = get_current_schema()
|
||||
for ft in _HNSW_FACT_TYPES:
|
||||
idx = _hnsw_index_name(ft, internal_id)
|
||||
for ft in _BANK_INDEX_FACT_TYPES:
|
||||
idx = _bank_index_name(ft, internal_id)
|
||||
await conn.execute(f"DROP INDEX IF EXISTS {schema}.{idx}")
|
||||
|
||||
|
||||
@@ -97,6 +113,22 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
Returns:
|
||||
BankProfile with name, typed DispositionTraits, and mission
|
||||
"""
|
||||
profile, _ = await get_or_create_bank_profile(pool, bank_id)
|
||||
return profile
|
||||
|
||||
|
||||
async def get_or_create_bank_profile(pool, bank_id: str) -> tuple[BankProfile, bool]:
|
||||
"""
|
||||
Get bank profile, auto-creating with defaults if it doesn't exist.
|
||||
|
||||
Same as get_bank_profile, but also returns a flag indicating whether the
|
||||
bank was freshly created on this call. Used by the memory engine to apply
|
||||
the HINDSIGHT_API_DEFAULT_BANK_TEMPLATE hook on first bank creation.
|
||||
|
||||
Returns:
|
||||
Tuple of (BankProfile, created) where created is True if the bank
|
||||
did not exist before this call.
|
||||
"""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Try to get existing bank
|
||||
row = await conn.fetchrow(
|
||||
@@ -113,15 +145,18 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
if isinstance(disposition_data, str):
|
||||
disposition_data = json.loads(disposition_data)
|
||||
|
||||
return BankProfile(
|
||||
name=row["name"],
|
||||
disposition=DispositionTraits(**disposition_data),
|
||||
mission=row["mission"] or "",
|
||||
return (
|
||||
BankProfile(
|
||||
name=row["name"],
|
||||
disposition=DispositionTraits(**disposition_data),
|
||||
mission=row["mission"] or "",
|
||||
),
|
||||
False,
|
||||
)
|
||||
|
||||
# Bank doesn't exist, create with defaults.
|
||||
# Generate internal_id here so we control the value and can use it
|
||||
# immediately for HNSW index creation without a RETURNING round-trip.
|
||||
# immediately for vector index creation without a RETURNING round-trip.
|
||||
internal_id = uuid.uuid4()
|
||||
inserted = await conn.fetchval(
|
||||
f"""
|
||||
@@ -137,11 +172,15 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
internal_id,
|
||||
)
|
||||
|
||||
if inserted:
|
||||
# Fresh insert — create per-bank HNSW indexes (instant on empty bank)
|
||||
await create_bank_hnsw_indexes(conn, bank_id, str(internal_id))
|
||||
created = inserted is not None
|
||||
if created:
|
||||
# Fresh insert — create per-bank vector indexes (instant on empty bank)
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id))
|
||||
|
||||
return BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission="")
|
||||
return (
|
||||
BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission=""),
|
||||
created,
|
||||
)
|
||||
|
||||
|
||||
async def update_bank_disposition(pool, bank_id: str, disposition: dict[str, int]) -> None:
|
||||
|
||||
@@ -4,7 +4,9 @@ Chunk storage for retain pipeline.
|
||||
Handles storage of document chunks in the database.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ..memory_engine import fq_table
|
||||
from .types import ChunkMetadata
|
||||
@@ -12,6 +14,61 @@ from .types import ChunkMetadata
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def compute_chunk_hash(chunk_text: str) -> str:
|
||||
"""Compute SHA256 hash of chunk text for delta comparison."""
|
||||
return hashlib.sha256(chunk_text.encode()).hexdigest()
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExistingChunk:
|
||||
"""Represents a chunk already stored in the database."""
|
||||
|
||||
chunk_id: str
|
||||
chunk_index: int
|
||||
content_hash: str | None
|
||||
|
||||
|
||||
async def load_existing_chunks(conn, bank_id: str, document_id: str) -> list[ExistingChunk]:
|
||||
"""
|
||||
Load existing chunk metadata for a document.
|
||||
|
||||
Returns list of ExistingChunk with chunk_id, chunk_index, and content_hash.
|
||||
"""
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT chunk_id, chunk_index, content_hash
|
||||
FROM {fq_table("chunks")}
|
||||
WHERE document_id = $1 AND bank_id = $2
|
||||
ORDER BY chunk_index
|
||||
""",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
return [
|
||||
ExistingChunk(
|
||||
chunk_id=row["chunk_id"],
|
||||
chunk_index=row["chunk_index"],
|
||||
content_hash=row["content_hash"],
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
|
||||
async def delete_chunks_by_ids(conn, chunk_ids: list[str]) -> None:
|
||||
"""
|
||||
Delete specific chunks by their IDs.
|
||||
|
||||
This cascades to memory_units (via FK with CASCADE delete)
|
||||
and their links.
|
||||
"""
|
||||
if not chunk_ids:
|
||||
return
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('chunks')} WHERE chunk_id = ANY($1::text[])",
|
||||
chunk_ids,
|
||||
)
|
||||
|
||||
|
||||
async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[ChunkMetadata]) -> dict[int, str]:
|
||||
"""
|
||||
Store document chunks in the database.
|
||||
@@ -32,6 +89,7 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[
|
||||
chunk_ids = []
|
||||
chunk_texts = []
|
||||
chunk_indices = []
|
||||
content_hashes = []
|
||||
chunk_id_map = {}
|
||||
|
||||
for chunk in chunks:
|
||||
@@ -39,19 +97,30 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[
|
||||
chunk_ids.append(chunk_id)
|
||||
chunk_texts.append(chunk.chunk_text)
|
||||
chunk_indices.append(chunk.chunk_index)
|
||||
content_hashes.append(compute_chunk_hash(chunk.chunk_text))
|
||||
chunk_id_map[chunk.chunk_index] = chunk_id
|
||||
|
||||
# Batch insert all chunks
|
||||
# Batch upsert all chunks. ON CONFLICT makes this idempotent: re-submitting
|
||||
# a retain under the same document_id (the pattern in vectorize-io/hindsight#977)
|
||||
# may produce chunk_ids that already exist when upstream cascade-delete or
|
||||
# delta-retain paths don't run (or race with a concurrent task). Overwriting
|
||||
# is the correct behavior per the document_id grouping semantics — the caller
|
||||
# intends this chunk to hold the latest content at that (document_id, index).
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("chunks")} (chunk_id, document_id, bank_id, chunk_text, chunk_index)
|
||||
SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::integer[])
|
||||
INSERT INTO {fq_table("chunks")} (chunk_id, document_id, bank_id, chunk_text, chunk_index, content_hash)
|
||||
SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::integer[], $6::text[])
|
||||
ON CONFLICT (chunk_id) DO UPDATE SET
|
||||
chunk_text = EXCLUDED.chunk_text,
|
||||
chunk_index = EXCLUDED.chunk_index,
|
||||
content_hash = EXCLUDED.content_hash
|
||||
""",
|
||||
chunk_ids,
|
||||
[document_id] * len(chunk_texts),
|
||||
[bank_id] * len(chunk_texts),
|
||||
chunk_texts,
|
||||
chunk_indices,
|
||||
content_hashes,
|
||||
)
|
||||
|
||||
return chunk_id_map
|
||||
|
||||
@@ -47,6 +47,16 @@ async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> lis
|
||||
embeddings_backend.encode,
|
||||
texts,
|
||||
)
|
||||
return embeddings
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate batch embeddings: {str(e)}")
|
||||
|
||||
# Guarantee 1:1 alignment with input texts. A silent length mismatch here
|
||||
# propagates downstream as zip() drops items, eventually surfacing as an
|
||||
# IndexError in retain mapping (see issue #1037).
|
||||
if len(embeddings) != len(texts):
|
||||
raise RuntimeError(
|
||||
f"Embeddings backend returned {len(embeddings)} vectors for {len(texts)} input texts; "
|
||||
"expected exact 1:1 alignment"
|
||||
)
|
||||
|
||||
return embeddings
|
||||
|
||||
@@ -12,61 +12,27 @@ from .types import EntityLink, ProcessedFact
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def process_entities_batch(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
def _prepare_facts_for_entity_processing(
|
||||
facts: list[ProcessedFact],
|
||||
log_buffer: list[str] = None,
|
||||
user_entities_per_content: dict[int, list[dict]] = None,
|
||||
entity_labels: list | None = None,
|
||||
) -> list[EntityLink]:
|
||||
user_entities_per_content: dict[int, list[dict]] | None = None,
|
||||
) -> tuple[list[str], list, list[list[dict]]]:
|
||||
"""
|
||||
Process entities for all facts and create entity links.
|
||||
|
||||
This function:
|
||||
1. Extracts entity mentions from fact texts
|
||||
2. Merges user-provided entities with LLM-extracted entities
|
||||
3. Resolves entity names to canonical entities
|
||||
4. Creates entity records in the database
|
||||
5. Returns entity links ready for insertion
|
||||
|
||||
Args:
|
||||
entity_resolver: EntityResolver instance for entity resolution
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
unit_ids: List of unit IDs (same length as facts)
|
||||
facts: List of ProcessedFact objects
|
||||
log_buffer: Optional buffer for detailed logging
|
||||
user_entities_per_content: Dict mapping content_index to list of user-provided entities
|
||||
Extract fact texts, dates, and merged entity lists from ProcessedFact objects.
|
||||
|
||||
Returns:
|
||||
List of EntityLink objects for batch insertion
|
||||
Tuple of (fact_texts, fact_dates, entities_per_fact)
|
||||
"""
|
||||
if not unit_ids or not facts:
|
||||
return []
|
||||
|
||||
if len(unit_ids) != len(facts):
|
||||
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
|
||||
|
||||
user_entities_per_content = user_entities_per_content or {}
|
||||
|
||||
# Extract data for link_utils function
|
||||
fact_texts = [fact.fact_text for fact in facts]
|
||||
# Use occurred_start if available, otherwise use mentioned_at for entity timestamps
|
||||
fact_dates = [fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at for fact in facts]
|
||||
|
||||
# Convert EntityRef objects to dict format and merge with user-provided entities
|
||||
entities_per_fact = []
|
||||
for fact in facts:
|
||||
# Start with LLM-extracted entities
|
||||
llm_entities = [{"text": entity.name, "type": "CONCEPT"} for entity in (fact.entities or [])]
|
||||
|
||||
# Get user entities for this content (use content_index from fact)
|
||||
user_entities = user_entities_per_content.get(fact.content_index, [])
|
||||
|
||||
# Merge with case-insensitive deduplication
|
||||
seen_texts = {e["text"].lower() for e in llm_entities}
|
||||
for user_entity in user_entities:
|
||||
if user_entity["text"].lower() not in seen_texts:
|
||||
@@ -80,8 +46,48 @@ async def process_entities_batch(
|
||||
|
||||
entities_per_fact.append(llm_entities)
|
||||
|
||||
# Use existing link_utils function for entity processing
|
||||
entity_links = await link_utils.extract_entities_batch_optimized(
|
||||
return fact_texts, fact_dates, entities_per_fact
|
||||
|
||||
|
||||
async def resolve_entities(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
facts: list[ProcessedFact],
|
||||
log_buffer: list[str] = None,
|
||||
user_entities_per_content: dict[int, list[dict]] = None,
|
||||
entity_labels: list | None = None,
|
||||
) -> tuple[list[str], list[tuple], dict[str, list[str]]]:
|
||||
"""
|
||||
Phase 1: Resolve entity names to canonical IDs (read-heavy).
|
||||
|
||||
Should be called on a SEPARATE connection OUTSIDE the main write transaction
|
||||
to avoid holding the transaction open during expensive trigram scans.
|
||||
|
||||
Args:
|
||||
entity_resolver: EntityResolver instance
|
||||
conn: Database connection (separate from the main write transaction)
|
||||
bank_id: Bank identifier
|
||||
unit_ids: Placeholder unit IDs (used only for grouping)
|
||||
facts: List of ProcessedFact objects
|
||||
log_buffer: Optional buffer for detailed logging
|
||||
user_entities_per_content: Dict mapping content_index to user-provided entities
|
||||
entity_labels: Optional entity label taxonomy
|
||||
|
||||
Returns:
|
||||
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids)
|
||||
to pass to build_entity_links().
|
||||
"""
|
||||
if not unit_ids or not facts:
|
||||
return [], [], {}
|
||||
|
||||
if len(unit_ids) != len(facts):
|
||||
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
|
||||
|
||||
fact_texts, fact_dates, entities_per_fact = _prepare_facts_for_entity_processing(facts, user_entities_per_content)
|
||||
|
||||
return await link_utils.resolve_entities_only(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id,
|
||||
@@ -90,22 +96,67 @@ async def process_entities_batch(
|
||||
"", # context (not used in current implementation)
|
||||
fact_dates,
|
||||
entities_per_fact,
|
||||
log_buffer, # Pass log_buffer for detailed logging
|
||||
log_buffer,
|
||||
entity_labels=entity_labels,
|
||||
)
|
||||
|
||||
return entity_links
|
||||
|
||||
async def build_entity_links(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
resolved_entity_ids: list[str],
|
||||
entity_to_unit: list[tuple],
|
||||
unit_to_entity_ids: dict[str, list[str]],
|
||||
log_buffer: list[str] = None,
|
||||
skip_unit_entities_insert: bool = False,
|
||||
) -> list[EntityLink]:
|
||||
"""
|
||||
Build entity links for UI graph visualization.
|
||||
|
||||
Queries unit_entities to find shared entities between new and existing units,
|
||||
then generates EntityLink objects. When called from Phase 3 (post-transaction),
|
||||
set skip_unit_entities_insert=True since unit_entities were already inserted
|
||||
in Phase 2.
|
||||
|
||||
Args:
|
||||
entity_resolver: EntityResolver instance
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
unit_ids: Actual unit IDs (must already be inserted in the DB)
|
||||
resolved_entity_ids: From resolve_entities()
|
||||
entity_to_unit: From resolve_entities()
|
||||
unit_to_entity_ids: From resolve_entities()
|
||||
log_buffer: Optional buffer for detailed logging
|
||||
skip_unit_entities_insert: Skip unit_entities INSERT (already done in Phase 2)
|
||||
|
||||
Returns:
|
||||
List of EntityLink objects for batch insertion
|
||||
"""
|
||||
return await link_utils.build_entity_links_from_resolved(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
resolved_entity_ids,
|
||||
entity_to_unit,
|
||||
unit_to_entity_ids,
|
||||
log_buffer,
|
||||
skip_unit_entities_insert=skip_unit_entities_insert,
|
||||
)
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, entity_links: list[EntityLink]) -> None:
|
||||
async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str) -> None:
|
||||
"""
|
||||
Insert entity links in batch.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
entity_links: List of EntityLink objects
|
||||
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
|
||||
"""
|
||||
if not entity_links:
|
||||
return
|
||||
|
||||
await link_utils.insert_entity_links_batch(conn, entity_links)
|
||||
await link_utils.insert_entity_links_batch(conn, entity_links, bank_id)
|
||||
|
||||
@@ -87,7 +87,7 @@ class Fact(BaseModel):
|
||||
|
||||
# Required fields
|
||||
fact: str = Field(description="Combined fact text: what | when | where | who | why")
|
||||
fact_type: Literal["world", "experience", "opinion"] = Field(description="Perspective: world/experience/opinion")
|
||||
fact_type: Literal["world", "experience"] = Field(description="Perspective: world/experience")
|
||||
|
||||
# Optional temporal fields
|
||||
occurred_start: str | None = None
|
||||
@@ -159,7 +159,9 @@ class ExtractedFact(BaseModel):
|
||||
fact_kind: str = Field(default="conversation", description="'event' or 'conversation'")
|
||||
occurred_start: str | None = Field(default=None, description="ISO timestamp for events")
|
||||
occurred_end: str | None = Field(default=None, description="ISO timestamp for event end")
|
||||
fact_type: Literal["world", "assistant"] = Field(description="'world' or 'assistant'")
|
||||
fact_type: Literal["world", "assistant"] = Field(
|
||||
description="'world' = objective/external facts. 'assistant' = first-person actions, experiences, or observations by the speaker."
|
||||
)
|
||||
entities: list[Entity] | None = Field(default=None, description="People, places, concepts")
|
||||
causal_relations: list[FactCausalRelation] | None = Field(
|
||||
default=None, description="Links to previous facts (target_index < this fact's index)"
|
||||
@@ -261,7 +263,7 @@ class ExtractedFactVerbose(BaseModel):
|
||||
)
|
||||
|
||||
fact_type: Literal["world", "assistant"] = Field(
|
||||
description="'world' = about the user/others (background, experiences). 'assistant' = experience with the assistant."
|
||||
description="'world' = objective/external facts about other people, events, general knowledge. 'assistant' = first-person actions, experiences, or observations by the speaker (e.g., 'I changed X', 'I discovered Y')."
|
||||
)
|
||||
|
||||
entities: list[Entity] | None = Field(
|
||||
@@ -352,7 +354,9 @@ class VerbatimExtractedFact(BaseModel):
|
||||
fact_kind: str = Field(default="conversation", description="'event' or 'conversation'")
|
||||
occurred_start: str | None = Field(default=None, description="ISO timestamp for events")
|
||||
occurred_end: str | None = Field(default=None, description="ISO timestamp for event end")
|
||||
fact_type: Literal["world", "assistant"] = Field(description="'world' or 'assistant'")
|
||||
fact_type: Literal["world", "assistant"] = Field(
|
||||
description="'world' = objective/external facts. 'assistant' = first-person actions, experiences, or observations by the speaker."
|
||||
)
|
||||
entities: list[Entity] | None = Field(default=None, description="People, places, concepts")
|
||||
|
||||
@field_validator("entities", mode="before")
|
||||
@@ -499,8 +503,8 @@ fact_kind:
|
||||
- "conversation": Ongoing state, preference, trait (no dates)
|
||||
|
||||
fact_type:
|
||||
- "world": About user's life, other people, external events
|
||||
- "assistant": Interactions with assistant (requests, recommendations)
|
||||
- "world": About other people, external events, general knowledge, objective facts
|
||||
- "assistant": First-person actions, experiences, or observations by the speaker/author (e.g., "I changed X", "I discovered Y", "I debugged Z"). Also includes interactions with the user (requests, recommendations). If the narrator describes something they did, tried, learned, or decided — use "assistant".
|
||||
|
||||
══════════════════════════════════════════════════════════════════════════
|
||||
TEMPORAL HANDLING
|
||||
@@ -616,7 +620,7 @@ VERBOSE_FACT_EXTRACTION_PROMPT = """Extract facts from text into structured form
|
||||
|
||||
LANGUAGE: MANDATORY — Detect the language of the input text and produce ALL output in that EXACT same language. You are STRICTLY FORBIDDEN from translating or switching to any other language. Every single word of your output must be in the same language as the input. Do NOT output in a different language under any circumstance.
|
||||
|
||||
══════════════════════════════════════════════════════════════════════════
|
||||
{retain_mission_section}══════════════════════════════════════════════════════════════════════════
|
||||
FACT FORMAT - ALL FIVE DIMENSIONS REQUIRED - MAXIMUM VERBOSITY
|
||||
══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
@@ -827,7 +831,9 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
custom_instructions=config.retain_custom_instructions,
|
||||
)
|
||||
elif extraction_mode == "verbose":
|
||||
prompt = VERBOSE_FACT_EXTRACTION_PROMPT
|
||||
prompt = VERBOSE_FACT_EXTRACTION_PROMPT.format(
|
||||
retain_mission_section=retain_mission_section,
|
||||
)
|
||||
elif extraction_mode == "verbatim":
|
||||
prompt = VERBATIM_FACT_EXTRACTION_PROMPT.format(
|
||||
retain_mission_section=retain_mission_section,
|
||||
@@ -903,6 +909,7 @@ def _build_user_message(
|
||||
event_date: datetime | None,
|
||||
context: str,
|
||||
metadata: dict[str, str] | None = None,
|
||||
agent_name: str | None = None,
|
||||
) -> str:
|
||||
"""Build user message for fact extraction."""
|
||||
from .orchestrator import parse_datetime_flexible
|
||||
@@ -921,11 +928,15 @@ def _build_user_message(
|
||||
metadata_lines = "\n".join(f" {k}: {v}" for k, v in metadata.items())
|
||||
metadata_section = f"\nMetadata:\n{metadata_lines}"
|
||||
|
||||
narrator_section = ""
|
||||
if agent_name:
|
||||
narrator_section = f'\nNarrator: {agent_name} (AI agent — first-person statements like "I did X" are the agent\'s own actions; classify as "assistant")'
|
||||
|
||||
return f"""Extract facts from the following text chunk.
|
||||
|
||||
Chunk: {chunk_index + 1}/{total_chunks}
|
||||
Event Date: {event_date_str}
|
||||
Context: {sanitized_context}{metadata_section}
|
||||
Context: {sanitized_context}{metadata_section}{narrator_section}
|
||||
|
||||
Text:
|
||||
{sanitized_chunk}"""
|
||||
@@ -989,7 +1000,7 @@ async def _extract_facts_from_chunk(
|
||||
extract_causal_links = config.retain_extract_causal_links
|
||||
|
||||
# Build user message using helper function
|
||||
user_message = _build_user_message(chunk, chunk_index, total_chunks, event_date, context, metadata)
|
||||
user_message = _build_user_message(chunk, chunk_index, total_chunks, event_date, context, metadata, agent_name)
|
||||
|
||||
# Retry logic for JSON validation errors
|
||||
# Use retain-specific overrides if set, otherwise fall back to global LLM config
|
||||
@@ -1049,7 +1060,7 @@ async def _extract_facts_from_chunk(
|
||||
f"LLM response missing 'facts' field or returned empty list. "
|
||||
f"Response: {extraction_response_json}. "
|
||||
f"Input: "
|
||||
f"date: {event_date.isoformat()}, "
|
||||
f"date: {event_date.isoformat() if event_date else 'unset'}, "
|
||||
f"context: {context if context else 'none'}, "
|
||||
f"text: {chunk}"
|
||||
)
|
||||
@@ -1083,30 +1094,16 @@ async def _extract_facts_from_chunk(
|
||||
logger.warning(f"Skipping fact {i}: missing 'what' field")
|
||||
continue
|
||||
|
||||
# Critical field: fact_type
|
||||
# LLM uses "assistant" but we convert to "experience" for storage
|
||||
original_fact_type = llm_fact.get("fact_type")
|
||||
fact_type = original_fact_type
|
||||
|
||||
# Convert "assistant" → "experience" for storage
|
||||
if fact_type == "assistant":
|
||||
# Critical field: fact_type — "assistant" maps to "experience", everything else is "world".
|
||||
# If fact_type is unexpected, fall back to fact_kind before defaulting to "world".
|
||||
raw_fact_type = llm_fact.get("fact_type")
|
||||
if raw_fact_type == "assistant":
|
||||
fact_type = "experience"
|
||||
|
||||
# Validate fact_type (after conversion)
|
||||
if fact_type not in ["world", "experience", "opinion"]:
|
||||
# Try to fix common mistakes - check if they swapped fact_type and fact_kind
|
||||
fact_kind = llm_fact.get("fact_kind")
|
||||
if fact_kind == "assistant":
|
||||
fact_type = "experience"
|
||||
elif fact_kind in ["world", "experience", "opinion"]:
|
||||
fact_type = fact_kind
|
||||
else:
|
||||
# Default to 'world' if we can't determine
|
||||
fact_type = "world"
|
||||
logger.warning(
|
||||
f"Fact {i}: defaulting to fact_type='world' "
|
||||
f"(original fact_type={original_fact_type!r}, fact_kind={fact_kind!r})"
|
||||
)
|
||||
elif raw_fact_type == "world":
|
||||
fact_type = "world"
|
||||
else:
|
||||
raw_fact_kind = llm_fact.get("fact_kind")
|
||||
fact_type = "experience" if raw_fact_kind == "assistant" else "world"
|
||||
|
||||
# Get fact_kind for temporal handling (but don't store it)
|
||||
fact_kind = llm_fact.get("fact_kind", "conversation")
|
||||
@@ -1472,28 +1469,76 @@ async def extract_facts_from_text(
|
||||
f"chunk_size={config.retain_chunk_size:,}) - starting parallel LLM extraction"
|
||||
)
|
||||
|
||||
tasks = [
|
||||
_extract_facts_with_auto_split(
|
||||
chunk=chunk,
|
||||
chunk_index=i,
|
||||
total_chunks=len(chunks),
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name=agent_name,
|
||||
metadata=metadata,
|
||||
)
|
||||
for i, chunk in enumerate(chunks)
|
||||
]
|
||||
chunk_results = await asyncio.gather(*tasks)
|
||||
# Per-chunk retry wrapper: each chunk gets up to MAX_CHUNK_RETRIES attempts.
|
||||
# This handles transient LLM failures (timeouts, rate limits, malformed responses)
|
||||
# without discarding the entire batch. If a chunk still fails after all retries,
|
||||
# the ENTIRE retain fails — we do not accept partial extraction.
|
||||
MAX_CHUNK_RETRIES = 3
|
||||
CHUNK_RETRY_BASE_DELAY = 2.0 # seconds, doubles each retry
|
||||
|
||||
async def _extract_chunk_with_retry(chunk: str, chunk_index: int) -> tuple:
|
||||
"""Extract facts from a single chunk with retries on failure."""
|
||||
last_exception = None
|
||||
for attempt in range(MAX_CHUNK_RETRIES):
|
||||
try:
|
||||
return await _extract_facts_with_auto_split(
|
||||
chunk=chunk,
|
||||
chunk_index=chunk_index,
|
||||
total_chunks=len(chunks),
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name=agent_name,
|
||||
metadata=metadata,
|
||||
)
|
||||
except Exception as e:
|
||||
last_exception = e
|
||||
if attempt < MAX_CHUNK_RETRIES - 1:
|
||||
delay = CHUNK_RETRY_BASE_DELAY * (2**attempt)
|
||||
logger.warning(
|
||||
f"Chunk {chunk_index}/{len(chunks)} extraction failed "
|
||||
f"(attempt {attempt + 1}/{MAX_CHUNK_RETRIES}): "
|
||||
f"{type(e).__name__}. Retrying in {delay:.0f}s..."
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
logger.error(
|
||||
f"Chunk {chunk_index}/{len(chunks)} extraction failed after "
|
||||
f"{MAX_CHUNK_RETRIES} attempts: {type(e).__name__}: {e}"
|
||||
)
|
||||
raise last_exception
|
||||
|
||||
tasks = [_extract_chunk_with_retry(chunk, i) for i, chunk in enumerate(chunks)]
|
||||
|
||||
# return_exceptions=True so we can collect all results even if some chunks
|
||||
# exhausted their retries. We check for failures below and fail the retain
|
||||
# if ANY chunk could not be extracted — partial extraction is not acceptable.
|
||||
chunk_results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
all_facts = []
|
||||
chunk_metadata = [] # [(chunk_text, fact_count), ...]
|
||||
total_usage = TokenUsage()
|
||||
for chunk, (chunk_facts, chunk_usage) in zip(chunks, chunk_results):
|
||||
failed_chunks = []
|
||||
for i, (chunk, result) in enumerate(zip(chunks, chunk_results)):
|
||||
if isinstance(result, Exception):
|
||||
failed_chunks.append((i, result))
|
||||
continue
|
||||
chunk_facts, chunk_usage = result
|
||||
all_facts.extend(chunk_facts)
|
||||
chunk_metadata.append((chunk, len(chunk_facts)))
|
||||
total_usage = total_usage + chunk_usage
|
||||
|
||||
if failed_chunks:
|
||||
# Fail the entire retain — partial extraction is not acceptable.
|
||||
# All successfully extracted facts are discarded because the transaction
|
||||
# hasn't committed yet. The worker poller will retry the entire task.
|
||||
failed_summary = ", ".join(f"chunk {idx}: {type(err).__name__}" for idx, err in failed_chunks[:5])
|
||||
raise RuntimeError(
|
||||
f"Fact extraction failed: {len(failed_chunks)}/{len(chunks)} chunks failed "
|
||||
f"after {MAX_CHUNK_RETRIES} retries each. First failures: {failed_summary}"
|
||||
)
|
||||
|
||||
return all_facts, chunk_metadata, total_usage
|
||||
|
||||
|
||||
@@ -1592,7 +1637,13 @@ async def extract_facts_from_contents_batch_api(
|
||||
|
||||
# Build user message using helper function
|
||||
user_message = _build_user_message(
|
||||
chunk, chunk_index_in_content, len(chunks), item.event_date, item.context, item.metadata or None
|
||||
chunk,
|
||||
chunk_index_in_content,
|
||||
len(chunks),
|
||||
item.event_date,
|
||||
item.context,
|
||||
item.metadata or None,
|
||||
agent_name,
|
||||
)
|
||||
|
||||
# Build request body using helper function
|
||||
@@ -1754,23 +1805,17 @@ async def extract_facts_from_contents_batch_api(
|
||||
who = get_value("who")
|
||||
why = get_value("why")
|
||||
|
||||
# Critical field: fact_type
|
||||
original_fact_type = llm_fact.get("fact_type")
|
||||
fact_type = original_fact_type
|
||||
|
||||
# Convert "assistant" → "experience"
|
||||
if fact_type == "assistant":
|
||||
# Critical field: fact_type — only "assistant" maps to "experience", everything else is "world"
|
||||
# Critical field: fact_type — "assistant" maps to "experience", everything else is "world".
|
||||
# If fact_type is unexpected, fall back to fact_kind before defaulting to "world".
|
||||
raw_fact_type = llm_fact.get("fact_type")
|
||||
if raw_fact_type == "assistant":
|
||||
fact_type = "experience"
|
||||
|
||||
# Validate fact_type
|
||||
if fact_type not in ["world", "experience", "opinion"]:
|
||||
fact_kind = llm_fact.get("fact_kind")
|
||||
if fact_kind == "assistant":
|
||||
fact_type = "experience"
|
||||
elif fact_kind in ["world", "experience", "opinion"]:
|
||||
fact_type = fact_kind
|
||||
else:
|
||||
fact_type = "world"
|
||||
elif raw_fact_type == "world":
|
||||
fact_type = "world"
|
||||
else:
|
||||
raw_fact_kind = llm_fact.get("fact_kind")
|
||||
fact_type = "experience" if raw_fact_kind == "assistant" else "world"
|
||||
|
||||
# Build combined fact text
|
||||
combined_parts = [what]
|
||||
@@ -2069,8 +2114,9 @@ async def extract_facts_from_contents(
|
||||
)
|
||||
fact_extraction_tasks.append(task)
|
||||
|
||||
# Step 2: Wait for all fact extractions to complete
|
||||
all_fact_results = await asyncio.gather(*fact_extraction_tasks)
|
||||
# Step 2: Wait for all fact extractions to complete.
|
||||
# Use return_exceptions=True so one content item failure doesn't discard the rest.
|
||||
all_fact_results = await asyncio.gather(*fact_extraction_tasks, return_exceptions=True)
|
||||
|
||||
# Step 3: Flatten and convert to typed objects
|
||||
extracted_facts: list[ExtractedFactType] = []
|
||||
@@ -2080,9 +2126,16 @@ async def extract_facts_from_contents(
|
||||
global_chunk_idx = 0
|
||||
global_fact_idx = 0
|
||||
|
||||
for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate(
|
||||
zip(contents, all_fact_results)
|
||||
):
|
||||
# Filter out failed content items
|
||||
valid_results = []
|
||||
for content, result in zip(contents, all_fact_results):
|
||||
if isinstance(result, Exception):
|
||||
logger.warning(f"Content extraction failed (skipping): {type(result).__name__}: {result}")
|
||||
valid_results.append((content, ([], [], TokenUsage())))
|
||||
else:
|
||||
valid_results.append((content, result))
|
||||
|
||||
for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate(valid_results):
|
||||
total_usage = total_usage + content_usage
|
||||
chunk_start_idx = global_chunk_idx
|
||||
|
||||
|
||||
@@ -7,16 +7,34 @@ Handles insertion of facts into the database.
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
from ...config import get_config
|
||||
from ..memory_engine import fq_table
|
||||
from .bank_utils import DEFAULT_DISPOSITION, create_bank_hnsw_indexes
|
||||
from .bank_utils import DEFAULT_DISPOSITION, create_bank_vector_indexes
|
||||
from .fact_extraction import _sanitize_text
|
||||
from .types import ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def get_document_content(
|
||||
conn,
|
||||
bank_id: str,
|
||||
document_id: str,
|
||||
) -> str | None:
|
||||
"""Fetch the original_text of an existing document.
|
||||
|
||||
Returns None if the document does not exist.
|
||||
"""
|
||||
row = await conn.fetchval(
|
||||
f"SELECT original_text FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
return row
|
||||
|
||||
|
||||
async def insert_facts_batch(
|
||||
conn, bank_id: str, facts: list[ProcessedFact], document_id: str | None = None
|
||||
) -> list[str]:
|
||||
@@ -44,7 +62,6 @@ async def insert_facts_batch(
|
||||
mentioned_ats = []
|
||||
contexts = []
|
||||
fact_types = []
|
||||
confidence_scores = []
|
||||
metadata_jsons = []
|
||||
chunk_ids = []
|
||||
document_ids = []
|
||||
@@ -64,8 +81,6 @@ async def insert_facts_batch(
|
||||
mentioned_ats.append(fact.mentioned_at)
|
||||
contexts.append(_sanitize_text(fact.context))
|
||||
fact_types.append(fact.fact_type)
|
||||
# confidence_score is only for opinion facts
|
||||
confidence_scores.append(1.0 if fact.fact_type == "opinion" else None)
|
||||
metadata_jsons.append(json.dumps(fact.metadata))
|
||||
chunk_ids.append(fact.chunk_id)
|
||||
# Use per-fact document_id if available, otherwise fallback to batch-level document_id
|
||||
@@ -81,9 +96,15 @@ async def insert_facts_batch(
|
||||
if fact.entities:
|
||||
signal_parts.extend(e.name for e in fact.entities)
|
||||
if fact.occurred_start:
|
||||
signal_parts.append(fact.occurred_start.strftime("%B %-d %Y"))
|
||||
try:
|
||||
signal_parts.append(fact.occurred_start.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
|
||||
except (ValueError, AttributeError):
|
||||
pass
|
||||
if fact.occurred_end and fact.occurred_end != fact.occurred_start:
|
||||
signal_parts.append(fact.occurred_end.strftime("%B %-d %Y"))
|
||||
try:
|
||||
signal_parts.append(fact.occurred_end.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
|
||||
except (ValueError, AttributeError):
|
||||
pass
|
||||
text_signals_list.append(" ".join(signal_parts) if signal_parts else None)
|
||||
|
||||
# Batch insert all facts
|
||||
@@ -97,18 +118,18 @@ async def insert_facts_batch(
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[], $15::jsonb[], $16::text[]
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
@@ -129,18 +150,18 @@ async def insert_facts_batch(
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[], $15::jsonb[], $16::text[]
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
@@ -162,7 +183,6 @@ async def insert_facts_batch(
|
||||
mentioned_ats,
|
||||
contexts,
|
||||
fact_types,
|
||||
confidence_scores,
|
||||
metadata_jsons,
|
||||
chunk_ids,
|
||||
document_ids,
|
||||
@@ -201,8 +221,87 @@ async def ensure_bank_exists(conn, bank_id: str) -> None:
|
||||
internal_id,
|
||||
)
|
||||
if inserted:
|
||||
# Fresh insert — create per-bank HNSW indexes
|
||||
await create_bank_hnsw_indexes(conn, bank_id, str(internal_id))
|
||||
# Fresh insert — create per-bank vector indexes
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id))
|
||||
|
||||
|
||||
async def delete_stale_observations_for_memories(
|
||||
conn,
|
||||
bank_id: str,
|
||||
fact_ids: "list[str | uuid.UUID]",
|
||||
) -> int:
|
||||
"""Delete observations whose source memories are about to be removed.
|
||||
|
||||
Mirrors the cleanup performed by ``MemoryEngine.delete_document`` so that
|
||||
every code path that removes ``memory_units`` also removes the
|
||||
observations derived from them. Without this, ingesting a fresh version
|
||||
of a document via the retain pipeline (which does a full-replace
|
||||
``DELETE FROM documents`` cascade) used to leave orphan observations
|
||||
pointing at memory IDs that no longer existed.
|
||||
|
||||
For each observation referencing any of ``fact_ids``:
|
||||
1. Delete the observation row (its text is stale once even one source
|
||||
memory disappears).
|
||||
2. Reset ``consolidated_at = NULL`` on the surviving source memories so
|
||||
they get re-consolidated under fresh observations on the next run.
|
||||
|
||||
Must be called within an active transaction, before the source memories
|
||||
are deleted.
|
||||
|
||||
Returns the number of observations deleted.
|
||||
"""
|
||||
if not fact_ids:
|
||||
return 0
|
||||
|
||||
fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in fact_ids]
|
||||
|
||||
affected_obs = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, source_memory_ids
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND fact_type = 'observation'
|
||||
AND source_memory_ids && $2::uuid[]
|
||||
""",
|
||||
bank_id,
|
||||
fact_uuids,
|
||||
)
|
||||
|
||||
if not affected_obs:
|
||||
return 0
|
||||
|
||||
deleted_set = {str(uid) for uid in fact_uuids}
|
||||
obs_ids = [obs["id"] for obs in affected_obs]
|
||||
seen_remaining: set[str] = set()
|
||||
remaining_source_ids: list[uuid.UUID] = []
|
||||
for obs in affected_obs:
|
||||
for src_id in obs["source_memory_ids"] or []:
|
||||
src_str = str(src_id)
|
||||
if src_str not in deleted_set and src_str not in seen_remaining:
|
||||
remaining_source_ids.append(src_id)
|
||||
seen_remaining.add(src_str)
|
||||
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE id = ANY($1::uuid[])",
|
||||
obs_ids,
|
||||
)
|
||||
|
||||
if remaining_source_ids:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET consolidated_at = NULL
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
remaining_source_ids,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[OBSERVATIONS] Deleted {len(obs_ids)} observations, reset {len(remaining_source_ids)} "
|
||||
f"source memories for re-consolidation in bank {bank_id}"
|
||||
)
|
||||
return len(obs_ids)
|
||||
|
||||
|
||||
async def handle_document_tracking(
|
||||
@@ -215,7 +314,10 @@ async def handle_document_tracking(
|
||||
document_tags: list[str] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Handle document tracking in the database.
|
||||
Handle document tracking in the database (full-replace mode).
|
||||
|
||||
Deletes the existing document (cascading to all units and links) on the
|
||||
first batch, then inserts the new document record.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
@@ -232,22 +334,106 @@ async def handle_document_tracking(
|
||||
combined_content = _sanitize_text(combined_content) or ""
|
||||
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
|
||||
|
||||
# Always delete old document first if it exists (cascades to units and links)
|
||||
# Only delete on the first batch to avoid deleting data we just inserted
|
||||
# Delete old document first (cascades to units and links).
|
||||
# Only delete on the first batch to avoid deleting data we just inserted.
|
||||
# Before the cascade, fan out to delete observations derived from the
|
||||
# outgoing memory_units — otherwise the FK ON DELETE CASCADE removes the
|
||||
# source memory_units but leaves observation rows pointing at IDs that
|
||||
# no longer exist (consolidated_at on co-source memories also stays
|
||||
# frozen). Same cleanup the explicit ``delete_document`` API performs.
|
||||
preserved_created_at = None
|
||||
if is_first_batch:
|
||||
await conn.fetchval(
|
||||
f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING id", document_id, bank_id
|
||||
existing_unit_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id FROM {fq_table("memory_units")}
|
||||
WHERE document_id = $1 AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
document_id,
|
||||
)
|
||||
existing_unit_ids = [row["id"] for row in existing_unit_rows]
|
||||
if existing_unit_ids:
|
||||
invalidated = await delete_stale_observations_for_memories(conn, bank_id, existing_unit_ids)
|
||||
if invalidated:
|
||||
logger.info(
|
||||
f"[RETAIN] Document {document_id} re-ingested: invalidated "
|
||||
f"{invalidated} observation(s) derived from {len(existing_unit_ids)} outgoing memory_units"
|
||||
)
|
||||
# Explicitly delete memory_units by document_id BEFORE deleting the
|
||||
# document row. The CASCADE from documents→chunks→memory_units only
|
||||
# catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL
|
||||
# (e.g. from partial writes or edge cases) would survive the cascade.
|
||||
# This explicit delete ensures complete cleanup.
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE document_id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
# Capture created_at before deletion so re-ingestion preserves it.
|
||||
preserved_created_at = await conn.fetchval(
|
||||
f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING created_at",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
# Insert document (or update if exists from concurrent operations)
|
||||
await _upsert_document_row(
|
||||
conn,
|
||||
bank_id,
|
||||
document_id,
|
||||
combined_content,
|
||||
content_hash,
|
||||
retain_params,
|
||||
document_tags,
|
||||
preserved_created_at=preserved_created_at,
|
||||
)
|
||||
|
||||
|
||||
async def upsert_document_metadata(
|
||||
conn,
|
||||
bank_id: str,
|
||||
document_id: str,
|
||||
combined_content: str,
|
||||
retain_params: dict | None = None,
|
||||
document_tags: list[str] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Update document metadata without deleting existing facts/chunks.
|
||||
|
||||
Used by delta retain: the document row is upserted but chunks and
|
||||
memory_units are managed separately at the chunk level.
|
||||
"""
|
||||
import hashlib
|
||||
|
||||
combined_content = _sanitize_text(combined_content) or ""
|
||||
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
|
||||
|
||||
await _upsert_document_row(conn, bank_id, document_id, combined_content, content_hash, retain_params, document_tags)
|
||||
|
||||
|
||||
async def _upsert_document_row(
|
||||
conn,
|
||||
bank_id: str,
|
||||
document_id: str,
|
||||
combined_content: str,
|
||||
content_hash: str,
|
||||
retain_params: dict | None = None,
|
||||
document_tags: list[str] | None = None,
|
||||
preserved_created_at: datetime | None = None,
|
||||
) -> None:
|
||||
"""Insert or update a document row.
|
||||
|
||||
When ``preserved_created_at`` is provided, it is used for ``created_at`` on
|
||||
INSERT so that re-ingesting a document (which deletes + inserts the row)
|
||||
keeps the original creation timestamp. ``updated_at`` is always set to
|
||||
``NOW()`` on both INSERT and the ON CONFLICT UPDATE branch.
|
||||
"""
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, metadata, retain_params, tags)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, retain_params, tags, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, COALESCE($7, NOW()), NOW())
|
||||
ON CONFLICT (id, bank_id) DO UPDATE
|
||||
SET original_text = EXCLUDED.original_text,
|
||||
content_hash = EXCLUDED.content_hash,
|
||||
metadata = EXCLUDED.metadata,
|
||||
retain_params = EXCLUDED.retain_params,
|
||||
tags = EXCLUDED.tags,
|
||||
updated_at = NOW()
|
||||
@@ -256,7 +442,38 @@ async def handle_document_tracking(
|
||||
bank_id,
|
||||
combined_content,
|
||||
content_hash,
|
||||
json.dumps({}), # Empty metadata dict
|
||||
json.dumps(retain_params) if retain_params else None,
|
||||
document_tags or [],
|
||||
preserved_created_at,
|
||||
)
|
||||
|
||||
|
||||
async def update_memory_units_tags(
|
||||
conn,
|
||||
bank_id: str,
|
||||
document_id: str,
|
||||
tags: list[str],
|
||||
) -> int:
|
||||
"""
|
||||
Update tags on all memory_units belonging to a document.
|
||||
|
||||
Used during delta retain to propagate tag changes to unchanged facts.
|
||||
|
||||
Returns:
|
||||
Number of memory units updated.
|
||||
"""
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET tags = $3, updated_at = NOW()
|
||||
WHERE bank_id = $1 AND document_id = $2
|
||||
""",
|
||||
bank_id,
|
||||
document_id,
|
||||
tags or [],
|
||||
)
|
||||
# result is a status string like "UPDATE 5"
|
||||
try:
|
||||
return int(result.split()[-1])
|
||||
except (ValueError, IndexError):
|
||||
return 0
|
||||
|
||||
@@ -32,17 +32,26 @@ async def create_temporal_links_batch(conn, bank_id: str, unit_ids: list[str]) -
|
||||
return await link_utils.create_temporal_links_batch_per_fact(conn, bank_id, unit_ids, log_buffer=[])
|
||||
|
||||
|
||||
async def create_semantic_links_batch(conn, bank_id: str, unit_ids: list[str], embeddings: list[list[float]]) -> int:
|
||||
async def create_semantic_links_batch(
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
embeddings: list[list[float]],
|
||||
pre_computed_ann_links: list[tuple] | None = None,
|
||||
) -> int:
|
||||
"""
|
||||
Create semantic links between facts.
|
||||
|
||||
Links facts that are semantically similar based on embeddings.
|
||||
When pre_computed_ann_links are provided (from Phase 1), they are used
|
||||
instead of running ANN queries inside the transaction.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
unit_ids: List of unit IDs to create links for
|
||||
embeddings: List of embedding vectors (same length as unit_ids)
|
||||
pre_computed_ann_links: Pre-computed ANN results from Phase 1
|
||||
|
||||
Returns:
|
||||
Number of semantic links created
|
||||
@@ -53,10 +62,12 @@ async def create_semantic_links_batch(conn, bank_id: str, unit_ids: list[str], e
|
||||
if len(unit_ids) != len(embeddings):
|
||||
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and embeddings ({len(embeddings)})")
|
||||
|
||||
return await link_utils.create_semantic_links_batch(conn, bank_id, unit_ids, embeddings, log_buffer=[])
|
||||
return await link_utils.create_semantic_links_batch(
|
||||
conn, bank_id, unit_ids, embeddings, log_buffer=[], pre_computed_ann_links=pre_computed_ann_links
|
||||
)
|
||||
|
||||
|
||||
async def create_causal_links_batch(conn, unit_ids: list[str], facts: list[ProcessedFact]) -> int:
|
||||
async def create_causal_links_batch(conn, bank_id: str, unit_ids: list[str], facts: list[ProcessedFact]) -> int:
|
||||
"""
|
||||
Create causal links between facts.
|
||||
|
||||
@@ -94,6 +105,6 @@ async def create_causal_links_batch(conn, unit_ids: list[str], facts: list[Proce
|
||||
else:
|
||||
causal_relations_per_fact.append([])
|
||||
|
||||
link_count = await link_utils.create_causal_links_batch(conn, unit_ids, causal_relations_per_fact)
|
||||
link_count = await link_utils.create_causal_links_batch(conn, bank_id, unit_ids, causal_relations_per_fact)
|
||||
|
||||
return link_count
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -25,6 +25,9 @@ class RetainContentDict(TypedDict, total=False):
|
||||
observation_scopes: How to scope observations for consolidation (optional).
|
||||
"per_tag" runs one pass per individual tag; "combined" (default) runs a
|
||||
single pass with all tags; a list[list[str]] specifies exact passes.
|
||||
update_mode: How to handle existing documents with the same document_id (optional).
|
||||
"replace" (default) deletes old data and reprocesses. "append" concatenates
|
||||
new content to the existing document and reprocesses.
|
||||
"""
|
||||
|
||||
content: str # Required
|
||||
@@ -37,6 +40,7 @@ class RetainContentDict(TypedDict, total=False):
|
||||
observation_scopes: (
|
||||
Literal["per_tag", "combined", "all_combinations"] | list[list[str]]
|
||||
) # Observation scopes for consolidation
|
||||
update_mode: Literal["replace", "append"]
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -107,7 +111,7 @@ class ExtractedFact:
|
||||
"""
|
||||
|
||||
fact_text: str
|
||||
fact_type: str # "world", "experience", "opinion", "observation"
|
||||
fact_type: str # "world", "experience", "observation"
|
||||
entities: list[str] = field(default_factory=list)
|
||||
occurred_start: datetime | None = None
|
||||
occurred_end: datetime | None = None
|
||||
@@ -221,6 +225,45 @@ class ProcessedFact:
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Phase3Context:
|
||||
"""
|
||||
Data passed from Phase 2 to Phase 3 for entity link building.
|
||||
|
||||
Contains the unit IDs and entity resolution data needed to build
|
||||
entity links for UI graph visualization after the write transaction commits.
|
||||
"""
|
||||
|
||||
unit_ids: list[str] = field(default_factory=list)
|
||||
resolved_entity_ids: list[str] = field(default_factory=list)
|
||||
entity_to_unit: list[tuple] = field(default_factory=list)
|
||||
unit_to_entity_ids: dict[str, list[str]] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntityResolutionResult:
|
||||
"""
|
||||
Result of Phase 1 entity resolution.
|
||||
|
||||
Contains resolved entity IDs and the mapping data needed to remap
|
||||
placeholder unit IDs to real IDs after fact insertion in Phase 2.
|
||||
"""
|
||||
|
||||
resolved_entity_ids: list[str]
|
||||
entity_to_unit: list[tuple]
|
||||
unit_to_entity_ids: dict[str, list[str]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Phase1Result:
|
||||
"""
|
||||
Full result of Phase 1 (entity resolution + optional semantic ANN).
|
||||
"""
|
||||
|
||||
entities: EntityResolutionResult
|
||||
semantic_ann_links: list[tuple]
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntityLink:
|
||||
"""
|
||||
@@ -248,7 +291,6 @@ class RetainBatch:
|
||||
contents: list[RetainContent]
|
||||
document_id: str | None = None
|
||||
fact_type_override: str | None = None
|
||||
confidence_score: float | None = None
|
||||
document_tags: list[str] = field(default_factory=list) # Tags applied to all items
|
||||
|
||||
# Extracted data (populated during processing)
|
||||
|
||||
@@ -3,12 +3,11 @@ Search module for memory retrieval.
|
||||
|
||||
Provides modular search architecture:
|
||||
- Retrieval: 4-way parallel (semantic + BM25 + graph + temporal)
|
||||
- Graph retrieval: Pluggable strategies (BFS, PPR)
|
||||
- Graph retrieval: Link expansion strategy
|
||||
- Reranking: Pluggable strategies (heuristic, cross-encoder)
|
||||
"""
|
||||
|
||||
from .graph_retrieval import BFSGraphRetriever, GraphRetriever
|
||||
from .mpfp_retrieval import MPFPGraphRetriever
|
||||
from .graph_retrieval import GraphRetriever
|
||||
from .reranking import CrossEncoderReranker
|
||||
from .retrieval import (
|
||||
ParallelRetrievalResult,
|
||||
@@ -21,7 +20,5 @@ __all__ = [
|
||||
"set_default_graph_retriever",
|
||||
"ParallelRetrievalResult",
|
||||
"GraphRetriever",
|
||||
"BFSGraphRetriever",
|
||||
"MPFPGraphRetriever",
|
||||
"CrossEncoderReranker",
|
||||
]
|
||||
|
||||
@@ -2,17 +2,16 @@
|
||||
Graph retrieval strategies for memory recall.
|
||||
|
||||
This module provides an abstraction for graph-based memory retrieval,
|
||||
allowing different algorithms (BFS spreading activation, PPR, etc.) to be
|
||||
swapped without changing the rest of the recall pipeline.
|
||||
allowing different algorithms to be swapped without changing the rest
|
||||
of the recall pipeline.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .tags import TagGroup, TagsMatch, filter_results_by_tag_groups, filter_results_by_tags
|
||||
from .types import MPFPTimings, RetrievalResult
|
||||
from .tags import TagGroup, TagsMatch
|
||||
from .types import GraphRetrievalTimings, RetrievalResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -29,7 +28,7 @@ class GraphRetriever(ABC):
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""Return identifier for this retrieval strategy (e.g., 'bfs', 'mpfp')."""
|
||||
"""Return identifier for this retrieval strategy (e.g., 'link_expansion')."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
@@ -47,7 +46,9 @@ class GraphRetriever(ABC):
|
||||
tags: list[str] | None = None, # Visibility scope tags for filtering
|
||||
tags_match: TagsMatch = "any", # How to match tags: 'any' (OR) or 'all' (AND)
|
||||
tag_groups: list[TagGroup] | None = None, # Compound boolean tag filter groups
|
||||
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
|
||||
created_after: datetime | None = None, # Only include memory_units created after this time
|
||||
created_before: datetime | None = None, # Only include memory_units created before this time
|
||||
) -> tuple[list[RetrievalResult], GraphRetrievalTimings | None]:
|
||||
"""
|
||||
Retrieve relevant facts via graph traversal.
|
||||
|
||||
@@ -55,228 +56,15 @@ class GraphRetriever(ABC):
|
||||
pool: Database connection pool
|
||||
query_embedding_str: Query embedding as string (for finding entry points)
|
||||
bank_id: Memory bank identifier
|
||||
fact_type: Fact type to filter ('world', 'experience', 'opinion', 'observation')
|
||||
fact_type: Fact type to filter ('world', 'experience', 'observation')
|
||||
budget: Maximum number of nodes to explore/return
|
||||
query_text: Original query text (optional, for some strategies)
|
||||
semantic_seeds: Pre-computed semantic entry points (from semantic retrieval)
|
||||
temporal_seeds: Pre-computed temporal entry points (from temporal retrieval)
|
||||
adjacency: Pre-loaded typed adjacency graph (optional, for MPFP)
|
||||
adjacency: Pre-loaded typed adjacency graph (optional)
|
||||
tags: Optional list of tags for visibility filtering (OR matching)
|
||||
|
||||
Returns:
|
||||
Tuple of (List of RetrievalResult with activation scores, optional timing info)
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class BFSGraphRetriever(GraphRetriever):
|
||||
"""
|
||||
Graph retrieval using BFS-style spreading activation.
|
||||
|
||||
Starting from semantic entry points, spreads activation through
|
||||
the memory graph (entity, temporal, causal links) using breadth-first
|
||||
traversal with decaying activation.
|
||||
|
||||
This is the original Hindsight graph retrieval algorithm.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
entry_point_limit: int = 5,
|
||||
entry_point_threshold: float = 0.5,
|
||||
activation_decay: float = 0.8,
|
||||
min_activation: float = 0.1,
|
||||
batch_size: int = 20,
|
||||
):
|
||||
"""
|
||||
Initialize BFS graph retriever.
|
||||
|
||||
Args:
|
||||
entry_point_limit: Maximum number of entry points to start from
|
||||
entry_point_threshold: Minimum semantic similarity for entry points
|
||||
activation_decay: Decay factor per hop (activation *= decay)
|
||||
min_activation: Minimum activation to continue spreading
|
||||
batch_size: Number of nodes to process per batch (for neighbor fetching)
|
||||
"""
|
||||
self.entry_point_limit = entry_point_limit
|
||||
self.entry_point_threshold = entry_point_threshold
|
||||
self.activation_decay = activation_decay
|
||||
self.min_activation = min_activation
|
||||
self.batch_size = batch_size
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "bfs"
|
||||
|
||||
async def retrieve(
|
||||
self,
|
||||
pool,
|
||||
query_embedding_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
budget: int,
|
||||
query_text: str | None = None,
|
||||
semantic_seeds: list[RetrievalResult] | None = None,
|
||||
temporal_seeds: list[RetrievalResult] | None = None,
|
||||
adjacency=None, # Not used by BFS
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
|
||||
"""
|
||||
Retrieve facts using BFS spreading activation.
|
||||
|
||||
Algorithm:
|
||||
1. Find entry points (top semantic matches above threshold)
|
||||
2. BFS traversal: visit neighbors, propagate decaying activation
|
||||
3. Boost causal links (causes, enables, prevents)
|
||||
4. Return visited nodes up to budget
|
||||
|
||||
Note: BFS finds its own entry points via embedding search.
|
||||
The semantic_seeds, temporal_seeds, and adjacency parameters are accepted
|
||||
for interface compatibility but not used.
|
||||
"""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
results = await self._retrieve_with_conn(
|
||||
conn,
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_type,
|
||||
budget,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
)
|
||||
return results, None
|
||||
|
||||
async def _retrieve_with_conn(
|
||||
self,
|
||||
conn,
|
||||
query_embedding_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
budget: int,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
) -> list[RetrievalResult]:
|
||||
"""Internal implementation with connection."""
|
||||
from .tags import build_tag_groups_where_clause, build_tags_where_clause_simple
|
||||
|
||||
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
|
||||
tag_groups_param_start = 6 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
params = [query_embedding_str, bank_id, fact_type, self.entry_point_threshold, self.entry_point_limit]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
|
||||
# Step 1: Find entry points
|
||||
entry_points = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = $3
|
||||
AND (1 - (embedding <=> $1::vector)) >= $4
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $5
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
if not entry_points:
|
||||
logger.debug(
|
||||
f"[BFS] No entry points found for fact_type={fact_type} (tags={tags}, tags_match={tags_match})"
|
||||
)
|
||||
return []
|
||||
|
||||
logger.debug(
|
||||
f"[BFS] Found {len(entry_points)} entry points for fact_type={fact_type} "
|
||||
f"(tags={tags}, tags_match={tags_match})"
|
||||
)
|
||||
|
||||
# Step 2: BFS spreading activation
|
||||
visited = set()
|
||||
results = []
|
||||
queue = [(RetrievalResult.from_db_row(dict(r)), r["similarity"]) for r in entry_points]
|
||||
budget_remaining = budget
|
||||
|
||||
while queue and budget_remaining > 0:
|
||||
# Collect a batch of nodes to process
|
||||
batch_nodes = []
|
||||
batch_activations = {}
|
||||
|
||||
while queue and len(batch_nodes) < self.batch_size and budget_remaining > 0:
|
||||
current, activation = queue.pop(0)
|
||||
unit_id = current.id
|
||||
|
||||
if unit_id not in visited:
|
||||
visited.add(unit_id)
|
||||
budget_remaining -= 1
|
||||
current.activation = activation
|
||||
results.append(current)
|
||||
batch_nodes.append(current.id)
|
||||
batch_activations[unit_id] = activation
|
||||
|
||||
# Batch fetch neighbors
|
||||
if batch_nodes and budget_remaining > 0:
|
||||
max_neighbors = len(batch_nodes) * 20
|
||||
neighbors = await conn.fetch(
|
||||
f"""
|
||||
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.occurred_end,
|
||||
mu.mentioned_at, mu.fact_type,
|
||||
mu.document_id, mu.chunk_id, mu.tags,
|
||||
ml.weight, ml.link_type, ml.from_unit_id
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.weight >= $2
|
||||
AND mu.fact_type = $3
|
||||
ORDER BY ml.weight DESC
|
||||
LIMIT $4
|
||||
""",
|
||||
batch_nodes,
|
||||
self.min_activation,
|
||||
fact_type,
|
||||
max_neighbors,
|
||||
)
|
||||
|
||||
for n in neighbors:
|
||||
neighbor_id = str(n["id"])
|
||||
if neighbor_id not in visited:
|
||||
parent_id = str(n["from_unit_id"])
|
||||
parent_activation = batch_activations.get(parent_id, 0.5)
|
||||
|
||||
# Boost causal links
|
||||
link_type = n["link_type"]
|
||||
base_weight = n["weight"]
|
||||
|
||||
if link_type in ("causes", "caused_by"):
|
||||
causal_boost = 2.0
|
||||
elif link_type in ("enables", "prevents"):
|
||||
causal_boost = 1.5
|
||||
else:
|
||||
causal_boost = 1.0
|
||||
|
||||
effective_weight = base_weight * causal_boost
|
||||
new_activation = parent_activation * effective_weight * self.activation_decay
|
||||
|
||||
if new_activation > self.min_activation:
|
||||
neighbor_result = RetrievalResult.from_db_row(dict(n))
|
||||
queue.append((neighbor_result, new_activation))
|
||||
|
||||
# Apply tags filtering (BFS may traverse into memories that don't match tags criteria)
|
||||
if tags:
|
||||
results = filter_results_by_tags(results, tags, match=tags_match)
|
||||
|
||||
# Apply compound tag group filtering (post-traversal)
|
||||
if tag_groups:
|
||||
results = filter_results_by_tag_groups(results, tag_groups)
|
||||
|
||||
return results
|
||||
|
||||
@@ -4,32 +4,39 @@ Link Expansion graph retrieval.
|
||||
Expands from semantic/temporal seeds through three parallel, first-class signals
|
||||
stored in memory_links:
|
||||
|
||||
1. Entity links — precomputed co-occurrence graph (created at retain time, bounded to
|
||||
MAX_LINKS_PER_ENTITY per entity). Score = number of distinct shared
|
||||
entities between the seed set and each candidate.
|
||||
1. Entity links — query-time self-join through unit_entities. Score = number of distinct
|
||||
shared entities between the seed set and each candidate, computed via
|
||||
COUNT(DISTINCT entity_id). Uses a LATERAL per-entity cap
|
||||
(graph_per_entity_limit, default 200) to prevent high-fanout entities
|
||||
from exploding the self-join intermediate rows.
|
||||
2. Semantic links — precomputed kNN graph (each new fact linked to its top-5 most
|
||||
similar existing facts at insert time, similarity >= 0.7). Checked
|
||||
in both directions since the graph is not symmetric. Score = weight.
|
||||
3. Causal links — explicit causal chains (causes/caused_by/enables/prevents).
|
||||
Score = weight + 1.0 (boosted as highest-quality signal).
|
||||
|
||||
All three signals are bounded at retain time, so no LATERAL fan-out caps are needed
|
||||
at query time. Each expansion is a simple aggregation over a small result set.
|
||||
Entity expansion is bounded by graph_per_entity_limit (LATERAL cap per entity).
|
||||
A timeout fallback (graph_expansion_timeout) drops entity expansion entirely if the
|
||||
query still exceeds the budget.
|
||||
|
||||
For non-observation fact types the three expansions are issued as a single CTE query
|
||||
(one roundtrip, one connection) with a `source` discriminator column so the Python
|
||||
merge step can apply per-signal score transformations.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .graph_retrieval import GraphRetriever
|
||||
from .tags import TagGroup, TagsMatch, filter_results_by_tag_groups, filter_results_by_tags
|
||||
from .types import MPFPTimings, RetrievalResult
|
||||
from .types import GraphRetrievalTimings, RetrievalResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -44,6 +51,8 @@ async def _find_semantic_seeds(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> list[RetrievalResult]:
|
||||
"""Find semantic seeds via embedding search."""
|
||||
from .tags import build_tag_groups_where_clause, build_tags_where_clause_simple
|
||||
@@ -51,15 +60,29 @@ async def _find_semantic_seeds(
|
||||
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
|
||||
tag_groups_param_start = 6 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
params = [query_embedding_str, bank_id, fact_type, threshold, limit]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags, proof_count,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
@@ -68,6 +91,7 @@ async def _find_semantic_seeds(
|
||||
AND (1 - (embedding <=> $1::vector)) >= $4
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{created_range_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $5
|
||||
""",
|
||||
@@ -116,7 +140,9 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
|
||||
created_after: "datetime | None" = None,
|
||||
created_before: "datetime | None" = None,
|
||||
) -> tuple[list[RetrievalResult], GraphRetrievalTimings | None]:
|
||||
"""
|
||||
Retrieve facts by expanding links from seeds.
|
||||
|
||||
@@ -136,7 +162,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
Tuple of (results, timings)
|
||||
"""
|
||||
start_time = time.time()
|
||||
timings = MPFPTimings(fact_type=fact_type)
|
||||
timings = GraphRetrievalTimings(fact_type=fact_type)
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Find seeds if not provided
|
||||
@@ -154,6 +180,8 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
timings.seeds_time = time.time() - seeds_start
|
||||
logger.debug(
|
||||
@@ -262,31 +290,48 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
idx_memory_links_to_type_weight (to_unit_id, link_type, weight DESC)
|
||||
→ replaces costly BitmapAnd of two separate scans
|
||||
"""
|
||||
config = get_config()
|
||||
ml = fq_table("memory_links")
|
||||
mu = fq_table("memory_units")
|
||||
all_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH entity_expanded AS (
|
||||
-- Entity co-occurrence: seeds → their precomputed entity-link neighbors.
|
||||
-- Score = distinct shared entities (bounded at retain time to
|
||||
-- MAX_LINKS_PER_ENTITY=50). GROUP BY mu.id is sufficient because mu.id
|
||||
-- is the primary key and functionally determines all other mu columns.
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
COUNT(DISTINCT ml.entity_id)::float AS score,
|
||||
'entity'::text AS source
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON mu.id = ml.to_unit_id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'entity'
|
||||
AND mu.fact_type = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
ue = fq_table("unit_entities")
|
||||
|
||||
per_entity_limit = config.link_expansion_per_entity_limit
|
||||
|
||||
# Entity CTE with LATERAL fanout cap.
|
||||
# Every seed entity (including high-frequency ones) is kept, but each
|
||||
# entity's expansion is capped to per_entity_limit target units. The
|
||||
# LATERAL subquery orders by unit_id DESC so the most recently inserted
|
||||
# units are preferred (a recency proxy that is free — it rides the PK
|
||||
# index with no extra sort).
|
||||
entity_cte = f"""
|
||||
seed_entities AS (
|
||||
SELECT DISTINCT ue.entity_id
|
||||
FROM {ue} ue
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
),
|
||||
entity_expanded AS (
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
COUNT(DISTINCT se.entity_id)::float AS score,
|
||||
'entity'::text AS source
|
||||
FROM seed_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM {ue} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
AND ue_target.unit_id != ALL($1::uuid[])
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
LIMIT {per_entity_limit}
|
||||
) t
|
||||
JOIN {mu} mu ON mu.id = t.unit_id
|
||||
WHERE mu.fact_type = $2
|
||||
GROUP BY mu.id
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
),
|
||||
)"""
|
||||
|
||||
semantic_causal_cte = f"""
|
||||
semantic_expanded AS (
|
||||
-- Semantic kNN: both outgoing (seeds → their kNN at insert time) and
|
||||
-- incoming (facts inserted after seeds that found seeds as kNN).
|
||||
@@ -294,14 +339,14 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags,
|
||||
fact_type, document_id, chunk_id, tags, proof_count,
|
||||
MAX(weight) AS score,
|
||||
'semantic'::text AS source
|
||||
FROM (
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
ml.weight
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON mu.id = ml.to_unit_id
|
||||
@@ -313,7 +358,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
ml.weight
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON mu.id = ml.from_unit_id
|
||||
@@ -324,7 +369,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
) sem_raw
|
||||
GROUP BY id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags
|
||||
fact_type, document_id, chunk_id, tags, proof_count
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
),
|
||||
@@ -335,7 +380,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
SELECT DISTINCT ON (mu.id)
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
ml.weight AS score,
|
||||
'causal'::text AS source
|
||||
FROM {ml} ml
|
||||
@@ -346,18 +391,37 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
AND mu.fact_type = $2
|
||||
ORDER BY mu.id, ml.weight DESC
|
||||
LIMIT $3
|
||||
)
|
||||
)"""
|
||||
|
||||
full_query = f"""
|
||||
WITH {entity_cte},
|
||||
{semantic_causal_cte}
|
||||
SELECT * FROM entity_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
""",
|
||||
seed_ids,
|
||||
fact_type,
|
||||
budget,
|
||||
self.causal_weight_threshold,
|
||||
)
|
||||
"""
|
||||
|
||||
params = [seed_ids, fact_type, budget, self.causal_weight_threshold]
|
||||
|
||||
try:
|
||||
all_rows = await asyncio.wait_for(
|
||||
conn.fetch(full_query, *params),
|
||||
timeout=config.link_expansion_timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
f"[LinkExpansion] Entity expansion timed out after {config.link_expansion_timeout}s "
|
||||
f"for fact_type={fact_type}, falling back to semantic+causal only"
|
||||
)
|
||||
fallback_query = f"""
|
||||
WITH {semantic_causal_cte}
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
"""
|
||||
all_rows = await conn.fetch(fallback_query, *params)
|
||||
|
||||
entity_rows = [r for r in all_rows if r["source"] == "entity"]
|
||||
semantic_rows = [r for r in all_rows if r["source"] == "semantic"]
|
||||
@@ -397,6 +461,33 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
f"{len(source_ids_found)} source_memory_ids found"
|
||||
)
|
||||
|
||||
config = get_config()
|
||||
ue = fq_table("unit_entities")
|
||||
per_entity_limit = config.link_expansion_per_entity_limit
|
||||
|
||||
connected_sources_cte = f"""
|
||||
source_entities AS (
|
||||
SELECT DISTINCT ue_seed.entity_id
|
||||
FROM seed_sources ss
|
||||
JOIN {ue} ue_seed ON ue_seed.unit_id = ss.source_id
|
||||
),
|
||||
connected_sources AS (
|
||||
-- Find sources sharing entities with seed observation sources
|
||||
-- via LATERAL-capped self-join (prevents hub entity fanout).
|
||||
SELECT DISTINCT t.unit_id AS source_id
|
||||
FROM source_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM {ue} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
LIMIT {per_entity_limit}
|
||||
) t
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
|
||||
)
|
||||
)"""
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH seed_sources AS (
|
||||
@@ -405,22 +496,14 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND source_memory_ids IS NOT NULL
|
||||
),
|
||||
connected_sources AS (
|
||||
-- Mirror the non-observation entity expansion: follow pre-bounded entity
|
||||
-- links in memory_links (capped to MAX_LINKS_PER_ENTITY=50 at retain time).
|
||||
-- Score = number of distinct shared entities, same as the non-obs path.
|
||||
SELECT DISTINCT ml.to_unit_id AS source_id
|
||||
FROM seed_sources ss
|
||||
JOIN {fq_table("memory_links")} ml ON ml.from_unit_id = ss.source_id
|
||||
WHERE ml.link_type = 'entity'
|
||||
),
|
||||
{connected_sources_cte},
|
||||
connected_array AS (
|
||||
SELECT array_agg(source_id) AS source_ids FROM connected_sources
|
||||
)
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
(SELECT COUNT(DISTINCT s) FROM unnest(mu.source_memory_ids) s WHERE s = ANY(ca.source_ids))::float AS score
|
||||
FROM {fq_table("memory_units")} mu, connected_array ca
|
||||
WHERE mu.fact_type = 'observation'
|
||||
@@ -444,13 +527,13 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags,
|
||||
fact_type, document_id, chunk_id, tags, proof_count,
|
||||
MAX(weight) AS score,
|
||||
'semantic'::text AS source
|
||||
FROM (
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight
|
||||
mu.chunk_id, mu.tags, mu.proof_count, ml.weight
|
||||
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.to_unit_id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
@@ -458,21 +541,21 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
UNION ALL
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight
|
||||
mu.chunk_id, mu.tags, mu.proof_count, ml.weight
|
||||
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.from_unit_id
|
||||
WHERE ml.to_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
) sem_raw
|
||||
GROUP BY id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags, proof_count
|
||||
ORDER BY score DESC LIMIT $2
|
||||
),
|
||||
causal_expanded AS (
|
||||
SELECT DISTINCT ON (mu.id)
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight AS score, 'causal'::text AS source
|
||||
mu.chunk_id, mu.tags, mu.proof_count, ml.weight AS score, 'causal'::text AS source
|
||||
FROM {ml} ml JOIN {mu} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||||
|
||||
@@ -1,702 +0,0 @@
|
||||
"""
|
||||
Meta-Path Forward Push (MPFP) graph retrieval.
|
||||
|
||||
A sublinear graph traversal algorithm for memory retrieval over heterogeneous
|
||||
graphs with multiple edge types (semantic, temporal, causal, entity).
|
||||
|
||||
Combines meta-path patterns from HIN literature with Forward Push local
|
||||
propagation from Approximate PPR.
|
||||
|
||||
Key properties:
|
||||
- Sublinear in graph size (threshold pruning bounds active nodes)
|
||||
- Lazy edge loading: only loads edges for frontier nodes, not entire graph
|
||||
- Predefined patterns capture different retrieval intents
|
||||
- All patterns run in parallel, results fused via RRF
|
||||
- No LLM in the loop during traversal
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .graph_retrieval import GraphRetriever
|
||||
from .tags import TagGroup, TagsMatch
|
||||
from .types import MPFPTimings, RetrievalResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Data Classes
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class EdgeTarget:
|
||||
"""A neighbor node with its edge weight."""
|
||||
|
||||
node_id: str
|
||||
weight: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class EdgeCache:
|
||||
"""
|
||||
Cache for lazily-loaded edges.
|
||||
|
||||
Grows per-hop as edges are loaded for frontier nodes.
|
||||
Shared across patterns to avoid redundant loads.
|
||||
Loads ALL edge types at once to minimize DB queries.
|
||||
Thread-safe via asyncio lock to prevent redundant concurrent loads.
|
||||
"""
|
||||
|
||||
# edge_type -> from_node_id -> list of EdgeTarget
|
||||
graphs: dict[str, dict[str, list[EdgeTarget]]] = field(default_factory=dict)
|
||||
# Track which nodes have been fully loaded (all edge types)
|
||||
_fully_loaded: set[str] = field(default_factory=set)
|
||||
# Timing stats
|
||||
db_queries: int = 0
|
||||
edge_load_time: float = 0.0
|
||||
# Detailed hop timing for debugging
|
||||
hop_details: list[dict] = field(default_factory=list)
|
||||
# Lock to prevent redundant concurrent loads
|
||||
_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
||||
|
||||
def get_neighbors(self, edge_type: str, node_id: str) -> list[EdgeTarget]:
|
||||
"""Get neighbors for a node via a specific edge type."""
|
||||
return self.graphs.get(edge_type, {}).get(node_id, [])
|
||||
|
||||
def get_normalized_neighbors(self, edge_type: str, node_id: str, top_k: int) -> list[EdgeTarget]:
|
||||
"""Get top-k neighbors with weights normalized to sum to 1."""
|
||||
neighbors = self.get_neighbors(edge_type, node_id)[:top_k]
|
||||
if not neighbors:
|
||||
return []
|
||||
|
||||
total = sum(n.weight for n in neighbors)
|
||||
if total == 0:
|
||||
return []
|
||||
|
||||
return [EdgeTarget(node_id=n.node_id, weight=n.weight / total) for n in neighbors]
|
||||
|
||||
def is_fully_loaded(self, node_id: str) -> bool:
|
||||
"""Check if all edges for this node have been loaded."""
|
||||
return node_id in self._fully_loaded
|
||||
|
||||
def get_uncached(self, node_ids: list[str]) -> list[str]:
|
||||
"""Get node IDs that haven't been fully loaded yet."""
|
||||
return [n for n in node_ids if not self.is_fully_loaded(n)]
|
||||
|
||||
def add_all_edges(self, edges_by_type: dict[str, dict[str, list[EdgeTarget]]], all_queried: list[str]):
|
||||
"""
|
||||
Add loaded edges to the cache (all edge types at once).
|
||||
|
||||
Args:
|
||||
edges_by_type: Dict mapping edge_type -> from_node_id -> list of EdgeTarget
|
||||
all_queried: All node IDs that were queried (marks them as fully loaded)
|
||||
"""
|
||||
for edge_type, edges in edges_by_type.items():
|
||||
if edge_type not in self.graphs:
|
||||
self.graphs[edge_type] = {}
|
||||
for node_id, neighbors in edges.items():
|
||||
self.graphs[edge_type][node_id] = neighbors
|
||||
|
||||
# Mark all queried nodes as fully loaded (even if they have no edges)
|
||||
self._fully_loaded.update(all_queried)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PatternResult:
|
||||
"""Result from a single pattern traversal."""
|
||||
|
||||
pattern: list[str]
|
||||
scores: dict[str, float] # node_id -> accumulated mass
|
||||
|
||||
|
||||
@dataclass
|
||||
class MPFPConfig:
|
||||
"""Configuration for MPFP algorithm."""
|
||||
|
||||
alpha: float = 0.15 # teleport/keep probability
|
||||
threshold: float = 1e-6 # mass pruning threshold (lower = explore more)
|
||||
top_k_neighbors: int = 20 # fan-out limit per node
|
||||
|
||||
# Patterns from semantic seeds
|
||||
patterns_semantic: list[list[str]] = field(
|
||||
default_factory=lambda: [
|
||||
["semantic", "semantic"], # topic expansion
|
||||
["entity", "temporal"], # entity timeline
|
||||
["semantic", "causes"], # reasoning chains (forward)
|
||||
["semantic", "caused_by"], # reasoning chains (backward)
|
||||
["entity", "semantic"], # entity context
|
||||
]
|
||||
)
|
||||
|
||||
# Patterns from temporal seeds
|
||||
patterns_temporal: list[list[str]] = field(
|
||||
default_factory=lambda: [
|
||||
["temporal", "semantic"], # what was happening then
|
||||
["temporal", "entity"], # who was involved then
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SeedNode:
|
||||
"""An entry point node with its initial score."""
|
||||
|
||||
node_id: str
|
||||
score: float # initial mass (e.g., similarity score)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Lazy Edge Loading
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def load_all_edges_for_frontier(
|
||||
pool,
|
||||
node_ids: list[str],
|
||||
top_k_per_type: int = 20,
|
||||
) -> dict[str, dict[str, list[EdgeTarget]]]:
|
||||
"""
|
||||
Load top-k edges per (node, edge_type) for frontier nodes.
|
||||
|
||||
Uses a LATERAL join to efficiently fetch only the top-k edges per type,
|
||||
avoiding loading hundreds of entity edges when only 20 are needed.
|
||||
|
||||
Requires composite index: (from_unit_id, link_type, weight DESC)
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
node_ids: Frontier node IDs to load edges for
|
||||
top_k_per_type: Max edges to load per (node, link_type) pair
|
||||
|
||||
Returns:
|
||||
Dict mapping edge_type -> from_node_id -> list of EdgeTarget
|
||||
"""
|
||||
if not node_ids:
|
||||
return {}
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Use LATERAL join to get top-k per (from_node, link_type)
|
||||
# This leverages the composite index for efficient early termination
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
WITH frontier(node_id) AS (SELECT unnest($1::uuid[]))
|
||||
SELECT f.node_id as from_unit_id, lt.link_type, edges.to_unit_id, edges.weight
|
||||
FROM frontier f
|
||||
CROSS JOIN (VALUES ('semantic'), ('temporal'), ('entity'), ('causes'), ('caused_by')) AS lt(link_type)
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ml.to_unit_id, ml.weight
|
||||
FROM {fq_table("memory_links")} ml
|
||||
WHERE ml.from_unit_id = f.node_id
|
||||
AND ml.link_type = lt.link_type
|
||||
AND ml.weight >= 0.1
|
||||
ORDER BY ml.weight DESC
|
||||
LIMIT $2
|
||||
) edges
|
||||
""",
|
||||
node_ids,
|
||||
top_k_per_type,
|
||||
)
|
||||
|
||||
# Group by edge_type -> from_node -> neighbors
|
||||
result: dict[str, dict[str, list[EdgeTarget]]] = defaultdict(lambda: defaultdict(list))
|
||||
for row in rows:
|
||||
edge_type = row["link_type"]
|
||||
from_id = str(row["from_unit_id"])
|
||||
to_id = str(row["to_unit_id"])
|
||||
weight = row["weight"]
|
||||
result[edge_type][from_id].append(EdgeTarget(node_id=to_id, weight=weight))
|
||||
|
||||
# Convert nested defaultdicts to regular dicts
|
||||
return {edge_type: dict(edges) for edge_type, edges in result.items()}
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Core Algorithm (Async with Lazy Loading)
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class PatternState:
|
||||
"""State for a pattern traversal between hops."""
|
||||
|
||||
pattern: list[str]
|
||||
hop_index: int
|
||||
scores: dict[str, float]
|
||||
frontier: dict[str, float]
|
||||
|
||||
|
||||
def _init_pattern_state(seeds: list[SeedNode], pattern: list[str]) -> PatternState:
|
||||
"""Initialize pattern state from seeds."""
|
||||
if not seeds:
|
||||
return PatternState(pattern=pattern, hop_index=0, scores={}, frontier={})
|
||||
|
||||
total_seed_score = sum(s.score for s in seeds)
|
||||
if total_seed_score == 0:
|
||||
total_seed_score = len(seeds)
|
||||
|
||||
frontier = {s.node_id: s.score / total_seed_score for s in seeds}
|
||||
return PatternState(pattern=pattern, hop_index=0, scores={}, frontier=frontier)
|
||||
|
||||
|
||||
def _execute_hop(state: PatternState, cache: EdgeCache, config: MPFPConfig) -> set[str]:
|
||||
"""
|
||||
Execute ONE hop of traversal, return frontier nodes for next hop.
|
||||
|
||||
This is a pure function that uses cached edges (no DB access).
|
||||
Returns set of uncached nodes needed for next hop.
|
||||
"""
|
||||
if state.hop_index >= len(state.pattern):
|
||||
return set()
|
||||
|
||||
edge_type = state.pattern[state.hop_index]
|
||||
|
||||
# Collect active nodes above threshold
|
||||
active_nodes = [node_id for node_id, mass in state.frontier.items() if mass >= config.threshold]
|
||||
if not active_nodes:
|
||||
state.frontier = {}
|
||||
return set()
|
||||
|
||||
# Propagate mass using cached edges
|
||||
next_frontier: dict[str, float] = {}
|
||||
uncached_for_next: set[str] = set()
|
||||
|
||||
for node_id, mass in state.frontier.items():
|
||||
if mass < config.threshold:
|
||||
continue
|
||||
|
||||
# Keep α portion for this node
|
||||
state.scores[node_id] = state.scores.get(node_id, 0) + config.alpha * mass
|
||||
|
||||
# Push (1-α) to neighbors
|
||||
push_mass = (1 - config.alpha) * mass
|
||||
neighbors = cache.get_normalized_neighbors(edge_type, node_id, config.top_k_neighbors)
|
||||
|
||||
for neighbor in neighbors:
|
||||
next_frontier[neighbor.node_id] = next_frontier.get(neighbor.node_id, 0) + push_mass * neighbor.weight
|
||||
# Track if we'll need edges for this node in the next hop
|
||||
if not cache.is_fully_loaded(neighbor.node_id):
|
||||
uncached_for_next.add(neighbor.node_id)
|
||||
|
||||
state.frontier = next_frontier
|
||||
state.hop_index += 1
|
||||
|
||||
return uncached_for_next
|
||||
|
||||
|
||||
def _finalize_pattern(state: PatternState, config: MPFPConfig) -> PatternResult:
|
||||
"""Finalize pattern by adding remaining frontier mass to scores."""
|
||||
for node_id, mass in state.frontier.items():
|
||||
if mass >= config.threshold:
|
||||
state.scores[node_id] = state.scores.get(node_id, 0) + mass
|
||||
|
||||
return PatternResult(pattern=state.pattern, scores=state.scores)
|
||||
|
||||
|
||||
async def mpfp_traverse_hop_synchronized(
|
||||
pool,
|
||||
pattern_jobs: list[tuple[list[SeedNode], list[str]]],
|
||||
config: MPFPConfig,
|
||||
cache: EdgeCache,
|
||||
) -> list[PatternResult]:
|
||||
"""
|
||||
Execute ALL patterns with hop-synchronized edge loading.
|
||||
|
||||
Instead of running each pattern independently (causing multiple DB queries),
|
||||
this function:
|
||||
1. Runs hop 1 for ALL patterns (using pre-warmed seed edges)
|
||||
2. Collects ALL unique hop-2 frontier nodes across patterns
|
||||
3. Pre-warms hop-2 edges in ONE query
|
||||
4. Runs hop 2 for ALL patterns
|
||||
|
||||
This reduces DB queries from O(patterns * hops) to O(hops).
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
pattern_jobs: List of (seeds, pattern) tuples
|
||||
config: Algorithm parameters
|
||||
cache: Shared edge cache (should be pre-warmed with seed edges)
|
||||
|
||||
Returns:
|
||||
List of PatternResult for each pattern
|
||||
"""
|
||||
import time
|
||||
|
||||
# Initialize all pattern states
|
||||
states = [_init_pattern_state(seeds, pattern) for seeds, pattern in pattern_jobs]
|
||||
|
||||
# Determine max hops (all patterns should be same length, but be safe)
|
||||
max_hops = max((len(p) for _, p in pattern_jobs), default=0)
|
||||
|
||||
# Detailed timing for debugging
|
||||
hop_times: list[dict] = []
|
||||
|
||||
# Execute hop-by-hop across ALL patterns
|
||||
for hop in range(max_hops):
|
||||
hop_start = time.time()
|
||||
hop_timing = {"hop": hop, "patterns_executed": 0, "uncached_count": 0, "load_time": 0.0}
|
||||
|
||||
# Execute this hop for all patterns, collect uncached nodes for next hop
|
||||
all_uncached: set[str] = set()
|
||||
exec_start = time.time()
|
||||
for state in states:
|
||||
if state.hop_index < len(state.pattern):
|
||||
uncached = _execute_hop(state, cache, config)
|
||||
all_uncached.update(uncached)
|
||||
hop_timing["patterns_executed"] += 1
|
||||
hop_timing["exec_time"] = time.time() - exec_start
|
||||
|
||||
# Pre-warm edges for ALL uncached nodes before next hop
|
||||
hop_timing["uncached_count"] = len(all_uncached)
|
||||
if all_uncached:
|
||||
uncached_list = list(all_uncached - cache._fully_loaded)
|
||||
hop_timing["uncached_after_filter"] = len(uncached_list)
|
||||
if uncached_list:
|
||||
load_start = time.time()
|
||||
edges_by_type = await load_all_edges_for_frontier(pool, uncached_list, config.top_k_neighbors)
|
||||
hop_timing["load_time"] = time.time() - load_start
|
||||
cache.edge_load_time += hop_timing["load_time"]
|
||||
cache.db_queries += 1
|
||||
cache.add_all_edges(edges_by_type, uncached_list)
|
||||
hop_timing["edges_loaded"] = sum(
|
||||
len(neighbors) for edges in edges_by_type.values() for neighbors in edges.values()
|
||||
)
|
||||
|
||||
hop_timing["total_time"] = time.time() - hop_start
|
||||
hop_times.append(hop_timing)
|
||||
|
||||
# Store hop timing details in cache for logging
|
||||
cache.hop_details = hop_times
|
||||
|
||||
# Finalize all patterns
|
||||
return [_finalize_pattern(state, config) for state in states]
|
||||
|
||||
|
||||
async def mpfp_traverse_async(
|
||||
pool,
|
||||
seeds: list[SeedNode],
|
||||
pattern: list[str],
|
||||
config: MPFPConfig,
|
||||
cache: EdgeCache,
|
||||
) -> PatternResult:
|
||||
"""
|
||||
Async Forward Push traversal with lazy edge loading.
|
||||
|
||||
NOTE: For better performance with multiple patterns, use mpfp_traverse_hop_synchronized().
|
||||
This function is kept for single-pattern use cases.
|
||||
"""
|
||||
if not seeds:
|
||||
return PatternResult(pattern=pattern, scores={})
|
||||
|
||||
results = await mpfp_traverse_hop_synchronized(pool, [(seeds, pattern)], config, cache)
|
||||
return results[0] if results else PatternResult(pattern=pattern, scores={})
|
||||
|
||||
|
||||
def rrf_fusion(
|
||||
results: list[PatternResult],
|
||||
k: int = 60,
|
||||
top_k: int = 50,
|
||||
) -> list[tuple[str, float]]:
|
||||
"""
|
||||
Reciprocal Rank Fusion to combine pattern results.
|
||||
|
||||
Args:
|
||||
results: List of pattern results
|
||||
k: RRF constant (higher = more uniform weighting)
|
||||
top_k: Number of results to return
|
||||
|
||||
Returns:
|
||||
List of (node_id, fused_score) tuples, sorted by score descending
|
||||
"""
|
||||
fused: dict[str, float] = {}
|
||||
|
||||
for result in results:
|
||||
if not result.scores:
|
||||
continue
|
||||
|
||||
# Rank nodes by their score in this pattern
|
||||
ranked = sorted(result.scores.keys(), key=lambda n: result.scores[n], reverse=True)
|
||||
|
||||
for rank, node_id in enumerate(ranked):
|
||||
fused[node_id] = fused.get(node_id, 0) + 1.0 / (k + rank + 1)
|
||||
|
||||
# Sort by fused score and return top-k
|
||||
sorted_results = sorted(fused.items(), key=lambda x: x[1], reverse=True)
|
||||
|
||||
return sorted_results[:top_k]
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Database Loading
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def fetch_memory_units_by_ids(
|
||||
pool,
|
||||
node_ids: list[str],
|
||||
fact_type: str,
|
||||
) -> list[RetrievalResult]:
|
||||
"""Fetch full memory unit details for a list of node IDs."""
|
||||
if not node_ids:
|
||||
return []
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND fact_type = $2
|
||||
""",
|
||||
node_ids,
|
||||
fact_type,
|
||||
)
|
||||
|
||||
return [RetrievalResult.from_db_row(dict(r)) for r in rows]
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Graph Retriever Implementation
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
class MPFPGraphRetriever(GraphRetriever):
|
||||
"""
|
||||
Graph retrieval using Meta-Path Forward Push with lazy edge loading.
|
||||
|
||||
Runs predefined patterns in parallel from semantic and temporal seeds,
|
||||
loading edges on-demand per hop instead of loading entire graph upfront.
|
||||
"""
|
||||
|
||||
def __init__(self, config: MPFPConfig | None = None):
|
||||
"""
|
||||
Initialize MPFP retriever.
|
||||
|
||||
Args:
|
||||
config: Algorithm configuration (uses defaults if None)
|
||||
"""
|
||||
if config is None:
|
||||
# Read top_k_neighbors from global config
|
||||
from ...config import get_config
|
||||
|
||||
global_config = get_config()
|
||||
config = MPFPConfig(top_k_neighbors=global_config.mpfp_top_k_neighbors)
|
||||
self.config = config
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "mpfp"
|
||||
|
||||
async def retrieve(
|
||||
self,
|
||||
pool,
|
||||
query_embedding_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
budget: int,
|
||||
query_text: str | None = None,
|
||||
semantic_seeds: list[RetrievalResult] | None = None,
|
||||
temporal_seeds: list[RetrievalResult] | None = None,
|
||||
adjacency=None, # Ignored - kept for interface compatibility
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
|
||||
"""
|
||||
Retrieve facts using MPFP algorithm with lazy edge loading.
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
query_embedding_str: Query embedding (used for fallback seed finding)
|
||||
bank_id: Memory bank ID
|
||||
fact_type: Fact type to filter
|
||||
budget: Maximum results to return
|
||||
query_text: Original query text (optional)
|
||||
semantic_seeds: Pre-computed semantic entry points
|
||||
temporal_seeds: Pre-computed temporal entry points
|
||||
adjacency: Ignored (kept for interface compatibility)
|
||||
tags: Optional list of tags for visibility filtering (OR matching)
|
||||
|
||||
Returns:
|
||||
Tuple of (List of RetrievalResult with activation scores, MPFPTimings)
|
||||
"""
|
||||
import time
|
||||
|
||||
timings = MPFPTimings(fact_type=fact_type)
|
||||
|
||||
# Convert seeds to SeedNode format
|
||||
semantic_seed_nodes = self._convert_seeds(semantic_seeds, "similarity")
|
||||
temporal_seed_nodes = self._convert_seeds(temporal_seeds, "temporal_score")
|
||||
|
||||
# If no semantic seeds provided, fall back to finding our own
|
||||
if not semantic_seed_nodes:
|
||||
seeds_start = time.time()
|
||||
semantic_seed_nodes = await self._find_semantic_seeds(
|
||||
pool,
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_type,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
)
|
||||
timings.seeds_time = time.time() - seeds_start
|
||||
logger.debug(
|
||||
f"[MPFP] Found {len(semantic_seed_nodes)} semantic seeds for fact_type={fact_type} (tags={tags}, tags_match={tags_match})"
|
||||
)
|
||||
|
||||
# Collect all pattern jobs
|
||||
pattern_jobs = []
|
||||
|
||||
# Patterns from semantic seeds
|
||||
for pattern in self.config.patterns_semantic:
|
||||
if semantic_seed_nodes:
|
||||
pattern_jobs.append((semantic_seed_nodes, pattern))
|
||||
|
||||
# Patterns from temporal seeds
|
||||
for pattern in self.config.patterns_temporal:
|
||||
if temporal_seed_nodes:
|
||||
pattern_jobs.append((temporal_seed_nodes, pattern))
|
||||
|
||||
if not pattern_jobs:
|
||||
logger.debug(
|
||||
f"[MPFP] No pattern jobs (semantic_seeds={len(semantic_seed_nodes)}, temporal_seeds={len(temporal_seed_nodes)})"
|
||||
)
|
||||
return [], timings
|
||||
|
||||
timings.pattern_count = len(pattern_jobs)
|
||||
|
||||
# Shared edge cache across all patterns
|
||||
cache = EdgeCache()
|
||||
|
||||
# Pre-warm cache with ALL seed node edges BEFORE running patterns
|
||||
# This prevents redundant DB queries at hop 1
|
||||
all_seed_ids = list({s.node_id for seeds, _ in pattern_jobs for s in seeds})
|
||||
if all_seed_ids:
|
||||
import time as time_module
|
||||
|
||||
prewarm_start = time_module.time()
|
||||
edges_by_type = await load_all_edges_for_frontier(pool, all_seed_ids, self.config.top_k_neighbors)
|
||||
cache.edge_load_time += time_module.time() - prewarm_start
|
||||
cache.db_queries += 1
|
||||
cache.add_all_edges(edges_by_type, all_seed_ids)
|
||||
|
||||
# Run all patterns with HOP-SYNCHRONIZED edge loading
|
||||
# This batches hop-2 edge loads across ALL patterns into ONE query
|
||||
# Reduces DB queries from O(patterns * hops) to O(hops)
|
||||
step_start = time.time()
|
||||
pattern_results = await mpfp_traverse_hop_synchronized(pool, pattern_jobs, self.config, cache)
|
||||
timings.traverse = time.time() - step_start
|
||||
|
||||
# Record edge loading stats from cache
|
||||
timings.edge_count = sum(len(neighbors) for g in cache.graphs.values() for neighbors in g.values())
|
||||
timings.db_queries = cache.db_queries
|
||||
timings.edge_load_time = cache.edge_load_time
|
||||
timings.hop_details = cache.hop_details
|
||||
|
||||
# Fuse results
|
||||
step_start = time.time()
|
||||
fused = rrf_fusion(pattern_results, top_k=budget)
|
||||
timings.fusion = time.time() - step_start
|
||||
|
||||
if not fused:
|
||||
logger.debug(f"[MPFP] No fused results after RRF fusion (pattern_count={len(pattern_results)})")
|
||||
return [], timings
|
||||
|
||||
# Get top result IDs
|
||||
result_ids = [node_id for node_id, score in fused][:budget]
|
||||
|
||||
# Fetch full details
|
||||
step_start = time.time()
|
||||
results = await fetch_memory_units_by_ids(pool, result_ids, fact_type)
|
||||
timings.fetch = time.time() - step_start
|
||||
|
||||
# Filter results by tags (graph traversal may have picked up unfiltered memories)
|
||||
if tags:
|
||||
from .tags import filter_results_by_tags
|
||||
|
||||
results = filter_results_by_tags(results, tags, match=tags_match)
|
||||
|
||||
# Apply compound tag group filtering (post-traversal)
|
||||
if tag_groups:
|
||||
from .tags import filter_results_by_tag_groups
|
||||
|
||||
results = filter_results_by_tag_groups(results, tag_groups)
|
||||
|
||||
timings.result_count = len(results)
|
||||
|
||||
# Add activation scores from fusion
|
||||
score_map = {node_id: score for node_id, score in fused}
|
||||
for result in results:
|
||||
result.activation = score_map.get(result.id, 0.0)
|
||||
|
||||
# Sort by activation
|
||||
results.sort(key=lambda r: r.activation or 0, reverse=True)
|
||||
|
||||
return results, timings
|
||||
|
||||
def _convert_seeds(
|
||||
self,
|
||||
seeds: list[RetrievalResult] | None,
|
||||
score_attr: str,
|
||||
) -> list[SeedNode]:
|
||||
"""Convert RetrievalResult seeds to SeedNode format."""
|
||||
if not seeds:
|
||||
return []
|
||||
|
||||
result = []
|
||||
for seed in seeds:
|
||||
score = getattr(seed, score_attr, None)
|
||||
if score is None:
|
||||
score = seed.activation or seed.similarity or 1.0
|
||||
result.append(SeedNode(node_id=seed.id, score=score))
|
||||
|
||||
return result
|
||||
|
||||
async def _find_semantic_seeds(
|
||||
self,
|
||||
pool,
|
||||
query_embedding_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
limit: int = 20,
|
||||
threshold: float = 0.3,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
) -> list[SeedNode]:
|
||||
"""Fallback: find semantic seeds via embedding search."""
|
||||
from .tags import build_tag_groups_where_clause, build_tags_where_clause_simple
|
||||
|
||||
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
|
||||
tag_groups_param_start = 6 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
params = [query_embedding_str, bank_id, fact_type, threshold, limit]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, 1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = $3
|
||||
AND (1 - (embedding <=> $1::vector)) >= $4
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $5
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
return [SeedNode(node_id=str(r["id"]), score=r["similarity"]) for r in rows]
|
||||
@@ -2,6 +2,7 @@
|
||||
Cross-encoder neural reranking for search results.
|
||||
"""
|
||||
|
||||
import math
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from .types import MergedCandidate, ScoredResult
|
||||
@@ -13,6 +14,7 @@ UTC = timezone.utc
|
||||
# so the max combined boost is (1 + alpha/2)^2 ≈ +21% and min is (1 - alpha/2)^2 ≈ -19%.
|
||||
_RECENCY_ALPHA: float = 0.2
|
||||
_TEMPORAL_ALPHA: float = 0.2
|
||||
_PROOF_COUNT_ALPHA: float = 0.1 # Conservative: max ±5% for evidence strength
|
||||
|
||||
|
||||
def apply_combined_scoring(
|
||||
@@ -20,32 +22,81 @@ def apply_combined_scoring(
|
||||
now: datetime,
|
||||
recency_alpha: float = _RECENCY_ALPHA,
|
||||
temporal_alpha: float = _TEMPORAL_ALPHA,
|
||||
proof_count_alpha: float = _PROOF_COUNT_ALPHA,
|
||||
is_passthrough_reranker: bool = False,
|
||||
) -> None:
|
||||
"""Apply combined scoring to a list of ScoredResults in-place.
|
||||
|
||||
Uses the cross-encoder score as the primary relevance signal, with recency
|
||||
and temporal proximity applied as multiplicative boosts. This ensures the
|
||||
influence of these secondary signals is always proportional to the base
|
||||
relevance score, regardless of the cross-encoder model's score calibration.
|
||||
Uses the cross-encoder score as the primary relevance signal, with recency,
|
||||
temporal proximity, and proof count applied as multiplicative boosts. This
|
||||
ensures the influence of these secondary signals is always proportional to
|
||||
the base relevance score, regardless of the cross-encoder model's score
|
||||
calibration.
|
||||
|
||||
Formula::
|
||||
|
||||
recency_boost = 1 + recency_alpha * (recency - 0.5) # in [1-α/2, 1+α/2]
|
||||
temporal_boost = 1 + temporal_alpha * (temporal - 0.5) # in [1-α/2, 1+α/2]
|
||||
combined_score = cross_encoder_score_normalized * recency_boost * temporal_boost
|
||||
recency_boost = 1 + recency_alpha * (recency - 0.5) # in [1-α/2, 1+α/2]
|
||||
temporal_boost = 1 + temporal_alpha * (temporal - 0.5) # in [1-α/2, 1+α/2]
|
||||
proof_count_boost = 1 + proof_count_alpha * (proof_norm - 0.5) # in [1-α/2, 1+α/2]
|
||||
combined_score = CE_normalized * recency_boost * temporal_boost * proof_count_boost
|
||||
|
||||
proof_norm maps proof_count using a smooth logarithmic curve centered at 0.5,
|
||||
clamped to [0, 1]:
|
||||
proof_count=1 → 0.5 + 0 = 0.5 (neutral multiplier)
|
||||
proof_count=150 → clamped to 1.0 (max +5% boost)
|
||||
|
||||
Temporal proximity is treated as neutral (0.5) when not set by temporal retrieval,
|
||||
so temporal_boost collapses to 1.0 for non-temporal queries.
|
||||
|
||||
Proof count is treated as neutral (0.5) when not available (non-observation facts),
|
||||
so proof_count_boost collapses to 1.0 for world/experience/opinion facts.
|
||||
|
||||
Args:
|
||||
scored_results: Results from the cross-encoder reranker. Mutated in place.
|
||||
now: Current UTC datetime for recency calculation.
|
||||
recency_alpha: Max relative recency adjustment (default 0.2 → ±10%).
|
||||
temporal_alpha: Max relative temporal adjustment (default 0.2 → ±10%).
|
||||
proof_count_alpha: Max relative proof count adjustment (default 0.1 → ±5%).
|
||||
"""
|
||||
if now.tzinfo is None:
|
||||
now = now.replace(tzinfo=UTC)
|
||||
|
||||
# When the configured cross-encoder is a passthrough (e.g.
|
||||
# RRFPassthroughCrossEncoder used by slim deployments), every
|
||||
# cross_encoder_score_normalized is identical and provides no relevance
|
||||
# signal. In that case the multiplicative recency / temporal / proof_count
|
||||
# boosts below become the *only* ranking signal — making the final order a
|
||||
# pure recency sort regardless of how relevant a candidate actually is.
|
||||
#
|
||||
# Detect that case and seed cross_encoder_score_normalized from the RRF
|
||||
# rank instead, so the boosts modulate a meaningful base score rather than
|
||||
# replacing it. This is a no-op for real cross-encoders, which produce
|
||||
# diverse scores.
|
||||
# When the reranker is a passthrough (e.g. RRFPassthroughCrossEncoder used
|
||||
# by slim deployments), every cross_encoder_score_normalized is identical
|
||||
# and provides no relevance signal. The multiplicative recency / temporal /
|
||||
# proof_count boosts below would then become the *only* ranking signal,
|
||||
# making the final order a pure recency sort regardless of how relevant a
|
||||
# candidate actually is.
|
||||
#
|
||||
# Seed cross_encoder_score_normalized from the RRF rank instead, so the
|
||||
# boosts modulate a meaningful base score. Caller passes is_passthrough
|
||||
# explicitly because "all scores identical" is too fragile a heuristic —
|
||||
# a real reranker can also tie scores (especially in tests with synthetic
|
||||
# data) and we'd corrupt legitimate single-result reranks.
|
||||
if is_passthrough_reranker and scored_results:
|
||||
n = len(scored_results)
|
||||
sorted_by_rrf = sorted(
|
||||
scored_results,
|
||||
key=lambda s: getattr(getattr(s, "candidate", None), "rrf_score", 0.0),
|
||||
reverse=True,
|
||||
)
|
||||
denom = max(1, n - 1)
|
||||
for new_rank, sr in enumerate(sorted_by_rrf):
|
||||
# Map rank → [0.1, 1.0] so the recency boost can still nudge
|
||||
# ordering between adjacent candidates without overpowering RRF.
|
||||
sr.cross_encoder_score_normalized = 1.0 - (0.9 * new_rank / denom)
|
||||
|
||||
for sr in scored_results:
|
||||
# Recency: linear decay over 365 days → [0.1, 1.0]; neutral 0.5 if no date.
|
||||
sr.recency = 0.5
|
||||
@@ -59,13 +110,23 @@ def apply_combined_scoring(
|
||||
# Temporal proximity: meaningful only for temporal queries; neutral otherwise.
|
||||
sr.temporal = sr.retrieval.temporal_proximity if sr.retrieval.temporal_proximity is not None else 0.5
|
||||
|
||||
# Proof count: log-normalized evidence strength; neutral for non-observations.
|
||||
proof_count = sr.retrieval.proof_count
|
||||
if proof_count is not None and proof_count >= 1:
|
||||
# Clamp to [0, 1] so extreme counts stay within documented ±5% range
|
||||
proof_norm = min(1.0, max(0.0, 0.5 + (math.log(proof_count) / 10.0)))
|
||||
else:
|
||||
# Neutral baseline is precisely 0.5, ensuring neutral multiplier (1.0)
|
||||
proof_norm = 0.5
|
||||
|
||||
# RRF: kept at 0.0 for trace continuity but excluded from scoring.
|
||||
# RRF is batch-relative (min-max normalised) and redundant after reranking.
|
||||
sr.rrf_normalized = 0.0
|
||||
|
||||
recency_boost = 1.0 + recency_alpha * (sr.recency - 0.5)
|
||||
temporal_boost = 1.0 + temporal_alpha * (sr.temporal - 0.5)
|
||||
sr.combined_score = sr.cross_encoder_score_normalized * recency_boost * temporal_boost
|
||||
proof_count_boost = 1.0 + proof_count_alpha * (proof_norm - 0.5)
|
||||
sr.combined_score = sr.cross_encoder_score_normalized * recency_boost * temporal_boost * proof_count_boost
|
||||
sr.weight = sr.combined_score
|
||||
|
||||
|
||||
@@ -153,6 +214,8 @@ class CrossEncoderReranker:
|
||||
|
||||
# Normalize scores using sigmoid to [0, 1] range
|
||||
# Cross-encoder returns logits which can be negative
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
def sigmoid(x):
|
||||
@@ -163,11 +226,20 @@ class CrossEncoderReranker:
|
||||
# Create ScoredResult objects with cross-encoder scores
|
||||
scored_results = []
|
||||
for candidate, raw_score, norm_score in zip(candidates, scores, normalized_scores):
|
||||
# Sanitize NaN scores (cross-encoder can return NaN for certain inputs).
|
||||
# NaN propagates through all downstream scoring and Pydantic serializes
|
||||
# NaN as JSON null, which breaks clients expecting numeric values.
|
||||
raw = float(raw_score)
|
||||
norm = float(norm_score)
|
||||
if math.isnan(raw):
|
||||
raw = 0.0
|
||||
if math.isnan(norm):
|
||||
norm = 0.0
|
||||
scored_result = ScoredResult(
|
||||
candidate=candidate,
|
||||
cross_encoder_score=float(raw_score),
|
||||
cross_encoder_score_normalized=float(norm_score),
|
||||
weight=float(norm_score), # Initial weight is just cross-encoder score
|
||||
cross_encoder_score=raw,
|
||||
cross_encoder_score_normalized=norm,
|
||||
weight=norm, # Initial weight is just cross-encoder score
|
||||
)
|
||||
scored_results.append(scored_result)
|
||||
|
||||
|
||||
@@ -10,22 +10,31 @@ Implements:
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import Optional
|
||||
from typing import Any, Optional
|
||||
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .graph_retrieval import BFSGraphRetriever, GraphRetriever
|
||||
from .graph_retrieval import GraphRetriever
|
||||
from .link_expansion_retrieval import LinkExpansionRetriever
|
||||
from .mpfp_retrieval import MPFPGraphRetriever
|
||||
from .tags import TagGroup, TagsMatch, build_tag_groups_where_clause, build_tags_where_clause_simple
|
||||
from .types import MPFPTimings, RetrievalResult
|
||||
from .types import GraphRetrievalTimings, RetrievalResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def tokenize_query(query_text: str) -> list[str]:
|
||||
"""Normalize query text and split into BM25 tokens.
|
||||
|
||||
Strips punctuation, lowercases, and splits on whitespace.
|
||||
Returns an empty list when the query contains no word characters.
|
||||
"""
|
||||
return re.sub(r"[^\w\s]", " ", query_text.lower()).split()
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParallelRetrievalResult:
|
||||
"""Result from parallel retrieval across all methods."""
|
||||
@@ -36,7 +45,9 @@ class ParallelRetrievalResult:
|
||||
temporal: list[RetrievalResult] | None
|
||||
timings: dict[str, float] = field(default_factory=dict)
|
||||
temporal_constraint: tuple | None = None # (start_date, end_date)
|
||||
mpfp_timings: list[MPFPTimings] = field(default_factory=list) # MPFP sub-step timings per fact type
|
||||
graph_timings: list[GraphRetrievalTimings] = field(
|
||||
default_factory=list
|
||||
) # Graph retrieval sub-step timings per fact type
|
||||
max_conn_wait: float = 0.0 # Maximum connection acquisition wait time across all methods
|
||||
|
||||
|
||||
@@ -62,15 +73,7 @@ def get_default_graph_retriever() -> GraphRetriever:
|
||||
if _default_graph_retriever is None:
|
||||
config = get_config()
|
||||
retriever_type = config.graph_retriever.lower()
|
||||
if retriever_type == "mpfp":
|
||||
_default_graph_retriever = MPFPGraphRetriever()
|
||||
logger.info(
|
||||
f"Using MPFP graph retriever (top_k_neighbors={_default_graph_retriever.config.top_k_neighbors})"
|
||||
)
|
||||
elif retriever_type == "bfs":
|
||||
_default_graph_retriever = BFSGraphRetriever()
|
||||
logger.info("Using BFS graph retriever")
|
||||
elif retriever_type == "link_expansion":
|
||||
if retriever_type == "link_expansion":
|
||||
_default_graph_retriever = LinkExpansionRetriever()
|
||||
logger.info("Using LinkExpansion graph retriever")
|
||||
else:
|
||||
@@ -95,6 +98,8 @@ async def retrieve_semantic_bm25_combined(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]]:
|
||||
"""
|
||||
Combined semantic + BM25 retrieval for multiple fact types in a single query.
|
||||
@@ -129,36 +134,52 @@ async def retrieve_semantic_bm25_combined(
|
||||
Returns:
|
||||
Dict mapping fact_type -> (semantic_results, bm25_results)
|
||||
"""
|
||||
import re
|
||||
|
||||
result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {ft: ([], []) for ft in fact_types}
|
||||
|
||||
sanitized_text = re.sub(r"[^\w\s]", " ", query_text.lower())
|
||||
tokens = [token for token in sanitized_text.split() if token]
|
||||
tokens = tokenize_query(query_text)
|
||||
|
||||
# Over-fetch for HNSW approximation; semantic results trimmed to limit in Python.
|
||||
hnsw_fetch = max(limit * 5, 100)
|
||||
|
||||
cols = (
|
||||
"id, text, context, event_date, occurred_start, occurred_end, mentioned_at, "
|
||||
"fact_type, document_id, chunk_id, tags"
|
||||
"fact_type, document_id, chunk_id, tags, metadata, proof_count"
|
||||
)
|
||||
table = fq_table("memory_units")
|
||||
|
||||
# --- Parameter layout ---
|
||||
# $1 = query_emb_str (semantic arms)
|
||||
# $2 = bank_id
|
||||
# $3 = limit (BM25 LIMIT; semantic uses inlined hnsw_fetch literal)
|
||||
# $4 = bm25_text (only when tokens present)
|
||||
# $N = tags (N=4 when no tokens, N=5 when tokens present)
|
||||
# $M+ = tag_groups params (one per leaf, starting after tags param)
|
||||
tags_param_idx = 5 if tokens else 4
|
||||
# When tokens present:
|
||||
# $3 = limit (BM25 LIMIT; semantic uses inlined hnsw_fetch literal)
|
||||
# $4 = bm25_text
|
||||
# $5 = tags (if present)
|
||||
# $6+ = tag_groups params (one per leaf)
|
||||
# When no tokens ($3 is skipped — not included in params to avoid type inference gap):
|
||||
# $3 = tags (if present)
|
||||
# $4+ = tag_groups params (one per leaf)
|
||||
tags_param_idx = 5 if tokens else 3
|
||||
tags_clause = build_tags_where_clause_simple(tags, tags_param_idx, match=tags_match)
|
||||
|
||||
# tag_groups params start immediately after the tags param slot
|
||||
tag_groups_param_start = tags_param_idx + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
# --- created_at time range filter (appended after tags/groups) ---
|
||||
# Param indices are computed relative to the final params list built below,
|
||||
# so we pre-compute the next available index after all preceding params.
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
# --- Semantic UNION ALL arms (one per fact_type) ---
|
||||
# Each arm has its own ORDER BY embedding <=> $1 LIMIT {hnsw_fetch}, which
|
||||
# lets the planner use the partial HNSW index for that fact_type.
|
||||
@@ -176,6 +197,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
f" AND (1 - (embedding <=> $1::vector)) >= 0.3"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {created_range_clause}"
|
||||
f" ORDER BY embedding <=> $1::vector"
|
||||
f" LIMIT {hnsw_fetch})"
|
||||
)
|
||||
@@ -216,18 +238,21 @@ async def retrieve_semantic_bm25_combined(
|
||||
f" {bm25_where_filter}"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {created_range_clause}"
|
||||
f" ORDER BY {bm25_order_by}"
|
||||
f" LIMIT $3)"
|
||||
)
|
||||
|
||||
query = "\nUNION ALL\n".join(arms)
|
||||
|
||||
params: list = [query_emb_str, bank_id, limit]
|
||||
params: list = [query_emb_str, bank_id]
|
||||
if tokens:
|
||||
params.append(bm25_text_param)
|
||||
params.append(limit) # $3: BM25 LIMIT (only referenced when tokens are present)
|
||||
params.append(bm25_text_param) # $4
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
rows = await conn.fetch(query, *params)
|
||||
|
||||
@@ -261,6 +286,8 @@ async def retrieve_temporal_combined(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, list[RetrievalResult]]:
|
||||
"""
|
||||
Temporal retrieval for multiple fact types in a single query.
|
||||
@@ -294,10 +321,25 @@ async def retrieve_temporal_combined(
|
||||
tags_clause = build_tags_where_clause_simple(tags, 7, match=tags_match)
|
||||
tag_groups_param_start = 7 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
# created_at time range filter (after tags/groups)
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
params: list = [query_emb_str, bank_id, fact_types, start_date, end_date, semantic_threshold]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
# Two-phase entry point query:
|
||||
# Phase 1 (date_ranked): rank by date only — no embedding computation — for all units in
|
||||
@@ -329,9 +371,10 @@ async def retrieve_temporal_combined(
|
||||
)
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{created_range_clause}
|
||||
),
|
||||
sim_ranked AS (
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.proof_count, mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity,
|
||||
ROW_NUMBER() OVER (PARTITION BY mu.fact_type ORDER BY mu.embedding <=> $1::vector) AS sim_rn
|
||||
FROM date_ranked dr
|
||||
@@ -339,7 +382,7 @@ async def retrieve_temporal_combined(
|
||||
WHERE dr.rn <= 50
|
||||
AND (1 - (mu.embedding <=> $1::vector)) >= $6
|
||||
)
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, similarity
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, proof_count, document_id, chunk_id, tags, metadata, similarity
|
||||
FROM sim_ranked
|
||||
WHERE sim_rn <= 10
|
||||
""",
|
||||
@@ -437,7 +480,7 @@ async def retrieve_temporal_combined(
|
||||
# bank_id on memory_units lets the planner use idx_memory_units_bank_fact_type.
|
||||
neighbors = await conn.fetch(
|
||||
f"""
|
||||
SELECT src.from_unit_id, mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
SELECT src.from_unit_id, mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
|
||||
l.weight, l.link_type,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity
|
||||
FROM unnest($2::uuid[]) AS src(from_unit_id)
|
||||
@@ -531,6 +574,8 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> MultiFactTypeRetrievalResult:
|
||||
"""
|
||||
Optimized retrieval for multiple fact types using batched queries.
|
||||
@@ -589,6 +634,8 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
semantic_bm25_time = time.time() - semantic_bm25_start
|
||||
|
||||
@@ -608,6 +655,8 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
temporal_time = time.time() - temporal_start
|
||||
|
||||
@@ -615,9 +664,11 @@ async def retrieve_all_fact_types_parallel(
|
||||
timings["temporal_combined"] = temporal_time
|
||||
|
||||
# Step 3: Run graph retrieval for each fact type in parallel
|
||||
async def run_graph_for_fact_type(ft: str) -> tuple[str, list[RetrievalResult], float, MPFPTimings | None]:
|
||||
async def run_graph_for_fact_type(
|
||||
ft: str,
|
||||
) -> tuple[str, list[RetrievalResult], float, GraphRetrievalTimings | None]:
|
||||
graph_start = time.time()
|
||||
results, mpfp_timing = await retriever.retrieve(
|
||||
results, graph_timing = await retriever.retrieve(
|
||||
pool=pool,
|
||||
query_embedding_str=query_embedding_str,
|
||||
bank_id=bank_id,
|
||||
@@ -629,8 +680,10 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
return ft, results, time.time() - graph_start, mpfp_timing
|
||||
return ft, results, time.time() - graph_start, graph_timing
|
||||
|
||||
# Run graph for all fact types in parallel
|
||||
graph_tasks = [run_graph_for_fact_type(ft) for ft in fact_types]
|
||||
@@ -639,7 +692,7 @@ async def retrieve_all_fact_types_parallel(
|
||||
# Organize results by fact type
|
||||
results_by_fact_type: dict[str, ParallelRetrievalResult] = {}
|
||||
max_conn_wait = conn_wait # Single connection for semantic+bm25+temporal
|
||||
all_mpfp_timings: list[MPFPTimings] = []
|
||||
all_graph_timings: list[GraphRetrievalTimings] = []
|
||||
|
||||
for ft in fact_types:
|
||||
# Get semantic + bm25 results for this fact type
|
||||
@@ -648,14 +701,14 @@ async def retrieve_all_fact_types_parallel(
|
||||
# Find graph results for this fact type
|
||||
graph_results = []
|
||||
graph_time = 0.0
|
||||
mpfp_timing = None
|
||||
graph_timing = None
|
||||
for gr in graph_results_list:
|
||||
if gr[0] == ft:
|
||||
graph_results = gr[1]
|
||||
graph_time = gr[2]
|
||||
mpfp_timing = gr[3]
|
||||
if mpfp_timing:
|
||||
all_mpfp_timings.append(mpfp_timing)
|
||||
graph_timing = gr[3]
|
||||
if graph_timing:
|
||||
all_graph_timings.append(graph_timing)
|
||||
break
|
||||
|
||||
# Get temporal results for this fact type from combined result
|
||||
@@ -676,7 +729,7 @@ async def retrieve_all_fact_types_parallel(
|
||||
"temporal_extraction": temporal_extraction_time,
|
||||
},
|
||||
temporal_constraint=temporal_constraint,
|
||||
mpfp_timings=[mpfp_timing] if mpfp_timing else [],
|
||||
graph_timings=[graph_timing] if graph_timing else [],
|
||||
max_conn_wait=max_conn_wait,
|
||||
)
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user