Compare commits
22
Commits
build-ci
...
python-deser
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
344ac8fae8 | ||
|
|
4b0c617ecf | ||
|
|
0a04770450 | ||
|
|
60574ee08f | ||
|
|
7d95a002c7 | ||
|
|
83ca669011 | ||
|
|
e798979733 | ||
|
|
43f9a8bec2 | ||
|
|
f641b30d83 | ||
|
|
90be7c6829 | ||
|
|
6eec83b20d | ||
|
|
dd1e0986a1 | ||
|
|
69dec8ec34 | ||
|
|
888b50de12 | ||
|
|
fb7be3eced | ||
|
|
4499254f6d | ||
|
|
9943957fb7 | ||
|
|
03f47e29c8 | ||
|
|
1240b82629 | ||
|
|
08f1cda3bf | ||
|
|
a3a9d7b37d | ||
|
|
c2607d7699 |
@@ -50,3 +50,18 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_RERANKER_LOCAL_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2
|
||||
# For TEI provider:
|
||||
# HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081
|
||||
|
||||
# Observability & Tracing (Optional - disabled by default)
|
||||
# Enable OpenTelemetry tracing for LLM calls (GenAI semantic conventions)
|
||||
# HINDSIGHT_API_OTEL_TRACES_ENABLED=true
|
||||
#
|
||||
# Local development with Grafana LGTM stack (recommended - see scripts/dev/grafana/README.md)
|
||||
# HINDSIGHT_API_OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318
|
||||
#
|
||||
# Cloud backends (Grafana Cloud, Langfuse, DataDog, etc.)
|
||||
# HINDSIGHT_API_OTEL_EXPORTER_OTLP_ENDPOINT=https://your-backend-url
|
||||
# HINDSIGHT_API_OTEL_EXPORTER_OTLP_HEADERS="Authorization=Bearer your-token"
|
||||
#
|
||||
# Custom service name and environment (optional, defaults: hindsight-api, development)
|
||||
# HINDSIGHT_API_OTEL_SERVICE_NAME=hindsight-production
|
||||
# HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT=production
|
||||
|
||||
@@ -334,8 +334,9 @@ jobs:
|
||||
push: false
|
||||
load: ${{ matrix.variant == 'slim' }}
|
||||
tags: hindsight-${{ matrix.name }}:test
|
||||
cache-from: type=gha,scope=${{ matrix.name }}
|
||||
cache-to: type=gha,mode=max,scope=${{ matrix.name }}
|
||||
# Removed GitHub Actions cache (type=gha) - it frequently returns 502 errors
|
||||
# causing buildx to fail with "failed to parse error response 502"
|
||||
# Build will be slower but more reliable
|
||||
|
||||
# Only test slim variants to save disk space (they're much smaller)
|
||||
# Slim variants require external embedding providers
|
||||
|
||||
+3
-1
@@ -53,4 +53,6 @@ hindsight-clients/rust/target
|
||||
whats-next.md
|
||||
TASK.md
|
||||
# Changelog is now tracked in hindsight-docs/src/pages/changelog.md
|
||||
# CHANGELOG.md
|
||||
# CHANGELOG.md
|
||||
|
||||
blog-post*
|
||||
@@ -45,6 +45,7 @@ cd hindsight-control-plane && npm run dev
|
||||
./scripts/dev/start-docs.sh
|
||||
```
|
||||
|
||||
|
||||
### Generating Clients/OpenAPI
|
||||
```bash
|
||||
# Regenerate OpenAPI spec after API changes (REQUIRED after changing endpoints)
|
||||
|
||||
@@ -42,46 +42,51 @@ If you need more control over how and when your agent stores and recalls memorie
|
||||
|
||||

|
||||
|
||||
---
|
||||
|
||||
> 🤖 **Using a coding agent?** Install the Hindsight documentation skill for instant access to docs while you code:
|
||||
> ```bash
|
||||
> npx skills add https://github.com/vectorize-io/hindsight --skill hindsight-docs
|
||||
> ```
|
||||
> Works with Claude Code, Cursor, and other AI coding assistants.
|
||||
|
||||
---
|
||||
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Docker (recommended)
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=your-key
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
|
||||
docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
|
||||
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
|
||||
-e HINDSIGHT_API_LLM_MODEL=o3-mini \
|
||||
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
|
||||
ghcr.io/vectorize-io/hindsight:latest
|
||||
```
|
||||
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, and `lmstudio`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
|
||||
### Docker compose
|
||||
|
||||
|
||||
### Docker (external PostgreSQL)
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
export HINDSIGHT_DB_PASSWORD=choose-a-password
|
||||
cd docker/docker-compose
|
||||
|
||||
# edit the docker compose file with your favorite editor
|
||||
nano docker-compose.yaml
|
||||
|
||||
# start hindsight with an external PostgeSQL
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
```bash
|
||||
# stop and cleanup the pg volume with the optional parameter -v
|
||||
docker compose down -v
|
||||
docker compose up
|
||||
```
|
||||
|
||||
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
API: http://localhost:8888
|
||||
UI: http://localhost:9999
|
||||
|
||||
Install client:
|
||||
### Client
|
||||
|
||||
```bash
|
||||
pip install hindsight-client -U
|
||||
@@ -89,7 +94,7 @@ pip install hindsight-client -U
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
Python example:
|
||||
#### Python
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
@@ -106,7 +111,29 @@ client.recall(bank_id="my-bank", query="What does Alice do?")
|
||||
client.reflect(bank_id="my-bank", query="Tell me about Alice")
|
||||
```
|
||||
|
||||
### Python (embedded, no Docker)
|
||||
#### Node.js / TypeScript
|
||||
|
||||
```bash
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
```javascript
|
||||
const { HindsightClient } = require('@vectorize-io/hindsight-client');
|
||||
|
||||
const main = async () => {
|
||||
const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
|
||||
|
||||
await client.retain('my-bank', 'Alice loves hiking in Yosemite');
|
||||
|
||||
const results = await client.recall('my-bank', 'What does Alice like?');
|
||||
console.log(results);
|
||||
}
|
||||
|
||||
main();
|
||||
```
|
||||
|
||||
|
||||
### Python Embedded (no server required)
|
||||
|
||||
```bash
|
||||
pip install hindsight-all -U
|
||||
@@ -126,26 +153,6 @@ with HindsightServer(
|
||||
results = client.recall(bank_id="my-bank", query="Where does Alice work?")
|
||||
```
|
||||
|
||||
### Node.js / TypeScript
|
||||
|
||||
```bash
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
```javascript
|
||||
const { HindsightClient } = require('@vectorize-io/hindsight-client');
|
||||
|
||||
const example = async () => {
|
||||
const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
|
||||
|
||||
await client.retain('my-bank', 'Alice loves hiking in Yosemite');
|
||||
|
||||
const results = await client.recall('my-bank', 'What does Alice like?');
|
||||
console.log(results);
|
||||
}
|
||||
|
||||
example();
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -1,135 +0,0 @@
|
||||
# Docker Testing
|
||||
|
||||
Scripts for testing Hindsight Docker images locally and in CI.
|
||||
|
||||
## Scripts
|
||||
|
||||
### `test-image.sh`
|
||||
|
||||
General-purpose Docker image test script. Starts a container and verifies it becomes healthy.
|
||||
|
||||
**Usage:**
|
||||
```bash
|
||||
./docker/test-image.sh <image> [target]
|
||||
```
|
||||
|
||||
**Arguments:**
|
||||
- `image` - Docker image to test (e.g., `hindsight:test`, `ghcr.io/vectorize-io/hindsight:latest`)
|
||||
- `target` - Optional: `cp-only` for control plane, `api-only` for API, or `standalone` (default)
|
||||
|
||||
**Environment Variables:**
|
||||
- `GROQ_API_KEY` - Required for API/standalone images
|
||||
- `HINDSIGHT_API_LLM_PROVIDER` - LLM provider (default: `groq`)
|
||||
- `HINDSIGHT_API_LLM_MODEL` - LLM model (default: `llama-3.3-70b-versatile`)
|
||||
- `HINDSIGHT_API_EMBEDDINGS_PROVIDER` - Embeddings provider (for slim images)
|
||||
- `HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY` - OpenAI API key for embeddings
|
||||
- `HINDSIGHT_API_RERANKER_PROVIDER` - Reranker provider (for slim images)
|
||||
- `HINDSIGHT_API_COHERE_API_KEY` - Cohere API key for reranking
|
||||
- `SMOKE_TEST_TIMEOUT` - Timeout in seconds (default: 120)
|
||||
|
||||
**Examples:**
|
||||
|
||||
Test a full image (with local ML models):
|
||||
```bash
|
||||
export GROQ_API_KEY=gsk_xxx
|
||||
./docker/test-image.sh hindsight:test
|
||||
```
|
||||
|
||||
Test a slim image (with external providers):
|
||||
```bash
|
||||
export GROQ_API_KEY=gsk_xxx
|
||||
export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai
|
||||
export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxx
|
||||
export HINDSIGHT_API_RERANKER_PROVIDER=cohere
|
||||
export HINDSIGHT_API_COHERE_API_KEY=xxx
|
||||
./docker/test-image.sh hindsight-slim:test
|
||||
```
|
||||
|
||||
### `test-slim-local.sh`
|
||||
|
||||
Convenience wrapper for testing slim images locally. Automatically configures external providers.
|
||||
|
||||
**Usage:**
|
||||
```bash
|
||||
# Set API keys
|
||||
export GROQ_API_KEY=gsk_xxx
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
export COHERE_API_KEY=xxx
|
||||
|
||||
# Run test
|
||||
./docker/test-slim-local.sh [image]
|
||||
```
|
||||
|
||||
**Or inline:**
|
||||
```bash
|
||||
GROQ_API_KEY=gsk_xxx \
|
||||
OPENAI_API_KEY=sk-xxx \
|
||||
COHERE_API_KEY=xxx \
|
||||
./docker/test-slim-local.sh hindsight-slim:test
|
||||
```
|
||||
|
||||
This script:
|
||||
- ✅ Validates API keys are set
|
||||
- ✅ Configures OpenAI embeddings automatically
|
||||
- ✅ Configures Cohere reranking automatically
|
||||
- ✅ Calls `test-image.sh` with the right configuration
|
||||
|
||||
## Building and Testing Locally
|
||||
|
||||
### Build a slim image
|
||||
|
||||
```bash
|
||||
docker build \
|
||||
--build-arg INCLUDE_LOCAL_MODELS=false \
|
||||
--build-arg PRELOAD_ML_MODELS=false \
|
||||
--target standalone \
|
||||
-t hindsight-slim:test \
|
||||
-f docker/standalone/Dockerfile \
|
||||
.
|
||||
```
|
||||
|
||||
### Test the slim image
|
||||
|
||||
```bash
|
||||
# With API keys
|
||||
export GROQ_API_KEY=gsk_xxx
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
export COHERE_API_KEY=xxx
|
||||
|
||||
# Run test
|
||||
./docker/test-slim-local.sh hindsight-slim:test
|
||||
```
|
||||
|
||||
## Expected Output
|
||||
|
||||
**Successful test:**
|
||||
```
|
||||
Starting smoke test for: hindsight-slim:test
|
||||
Target: standalone
|
||||
Health endpoint: http://localhost:8888/health
|
||||
Timeout: 120s
|
||||
|
||||
Starting container...
|
||||
Waiting for health endpoint at http://localhost:8888/health...
|
||||
Still waiting... (10s)
|
||||
Still waiting... (20s)
|
||||
|
||||
Container is healthy after 25s
|
||||
|
||||
=== Health Response ===
|
||||
{
|
||||
"status": "healthy",
|
||||
"database": "connected"
|
||||
}
|
||||
|
||||
Smoke test PASSED
|
||||
```
|
||||
|
||||
## CI Integration
|
||||
|
||||
These scripts are used in CI to validate Docker images on every PR:
|
||||
|
||||
- `.github/workflows/test.yml` - Runs `test-image.sh` for slim variants with OpenAI/Cohere
|
||||
- `.github/workflows/release.yml` - Can optionally run smoke tests during release
|
||||
|
||||
See the workflows for the exact configuration.
|
||||
@@ -35,44 +35,12 @@ services:
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
pull_policy: always
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM-configuration for Grog
|
||||
# - HINDSIGHT_API_LLM_PROVIDER=groq
|
||||
# - HINDSIGHT_API_LLM_API_KEY=${GROG_API_KEY?Please set the GROG_API_KEY env variable}
|
||||
# - HINDSIGHT_API_LLM_MODEL=${HINDSIGHT_API_LLM_MODEL:-openai/gpt-oss-20b}
|
||||
|
||||
# LLM-configuration for OpenAI
|
||||
# - HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
# - HINDSIGHT_API_LLM_API_KEY=${OPENAI_API_KEY?Please set the OPENAI_API_KEY env variable}
|
||||
# - HINDSIGHT_API_LLM_MODEL=${HINDSIGHT_API_LLM_MODEL:-gpt-4o}
|
||||
|
||||
# Gemini
|
||||
# - HINDSIGHT_API_LLM_PROVIDER=gemini
|
||||
# - HINDSIGHT_API_LLM_API_KEY=${GEMINI_API_KEY?Please set the GEMINI_API_KEY env variable}
|
||||
# - HINDSIGHT_API_LLM_MODEL=${HINDSIGHT_API_LLM_MODEL:-gemini-2.0-flash}
|
||||
|
||||
# Anthropic
|
||||
# - HINDSIGHT_API_LLM_PROVIDER=anthropic
|
||||
# - HINDSIGHT_API_LLM_API_KEY=${ANTHROPIC_API_KEY?Please set the ANTHROPIC_API_KEY env variable}
|
||||
# - HINDSIGHT_API_LLM_MODEL=${HINDSIGHT_API_LLM_MODEL:-claude-sonnet-4-20250514}
|
||||
|
||||
# LLM-configuration for Ollama (local, no API key)
|
||||
# - HINDSIGHT_API_LLM_PROVIDER=ollama
|
||||
# - HINDSIGHT_API_LLM_BASE_URL=${HINDSIGHT_API_LLM_BASE_URL:-http://127.0.0.1:11434/v1}
|
||||
# - HINDSIGHT_API_LLM_MODEL=${HINDSIGHT_API_LLM_MODEL:-llama3.2}
|
||||
|
||||
|
||||
# Configuration for the external Postgres database
|
||||
- HINDSIGHT_API_LLM_API_KEY=${OPENAI_API_KEY?Please set the OPENAI_API_KEY env variable}
|
||||
- HINDSIGHT_API_DATABASE_URL=postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:?Please set the HINDSIGHT_DB_PASSWORD env variable}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
# use public schema, otherwise the app start fails (2026-02-06)
|
||||
- HINDSIGHT_API_DATABASE_SCHEMA=public
|
||||
|
||||
# disable if you don't want automatic migrations on startup
|
||||
- HINDSIGHT_API_RUN_MIGRATIONS_ON_STARTUP=true
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
|
||||
@@ -8,6 +8,7 @@
|
||||
# Set to false when using external providers (TEI, OpenAI, Cohere)
|
||||
# PRELOAD_ML_MODELS=true/false - Pre-download ML models during build (default: true)
|
||||
# Only effective when INCLUDE_LOCAL_MODELS=true
|
||||
# NOTE: tiktoken encodings are ALWAYS preloaded (required for air-gapped deployments)
|
||||
#
|
||||
# Examples:
|
||||
# docker build -t hindsight . # Both (standalone)
|
||||
@@ -167,6 +168,28 @@ USER hindsight
|
||||
|
||||
ENV PATH="/app/api/.venv/bin:${PATH}"
|
||||
|
||||
# Pre-download tiktoken encoding (ALWAYS - required for token counting even in air-gapped envs)
|
||||
# Tiktoken is a core runtime dependency, not an optional ML model
|
||||
RUN MAX_RETRIES=3; \
|
||||
RETRY_DELAY=5; \
|
||||
for i in $(seq 1 $MAX_RETRIES); do \
|
||||
echo "Attempt $i/$MAX_RETRIES: Downloading tiktoken encoding..."; \
|
||||
/app/api/.venv/bin/python -c "\
|
||||
import tiktoken; \
|
||||
print('Downloading cl100k_base encoding...'); \
|
||||
tiktoken.get_encoding('cl100k_base'); \
|
||||
print('Tiktoken encoding cached successfully')" && break; \
|
||||
if [ $i -lt $MAX_RETRIES ]; then \
|
||||
echo "Attempt $i failed, retrying in ${RETRY_DELAY}s..."; \
|
||||
sleep $RETRY_DELAY; \
|
||||
RETRY_DELAY=$((RETRY_DELAY * 2)); \
|
||||
fi; \
|
||||
done; \
|
||||
if [ $i -eq $MAX_RETRIES ]; then \
|
||||
echo "ERROR: Failed to download tiktoken encoding after $MAX_RETRIES attempts"; \
|
||||
exit 1; \
|
||||
fi
|
||||
|
||||
# Pre-download ML models to avoid runtime download (conditional)
|
||||
# Only runs if both PRELOAD_ML_MODELS=true AND INCLUDE_LOCAL_MODELS=true
|
||||
# Includes retry logic with exponential backoff for transient network failures
|
||||
@@ -185,7 +208,6 @@ print('Downloading embedding model...'); \
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5'); \
|
||||
print('Downloading cross-encoder model...'); \
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \
|
||||
print('Downloading tiktoken encoding...'); import tiktoken; tiktoken.get_encoding('cl100k_base'); \
|
||||
print('Models cached successfully')" && break; \
|
||||
if [ $i -lt $MAX_RETRIES ]; then \
|
||||
echo "Attempt $i failed, retrying in ${RETRY_DELAY}s..."; \
|
||||
@@ -297,6 +319,28 @@ USER hindsight
|
||||
|
||||
ENV PATH="/app/api/.venv/bin:${PATH}"
|
||||
|
||||
# Pre-download tiktoken encoding (ALWAYS - required for token counting even in air-gapped envs)
|
||||
# Tiktoken is a core runtime dependency, not an optional ML model
|
||||
RUN MAX_RETRIES=3; \
|
||||
RETRY_DELAY=5; \
|
||||
for i in $(seq 1 $MAX_RETRIES); do \
|
||||
echo "Attempt $i/$MAX_RETRIES: Downloading tiktoken encoding..."; \
|
||||
/app/api/.venv/bin/python -c "\
|
||||
import tiktoken; \
|
||||
print('Downloading cl100k_base encoding...'); \
|
||||
tiktoken.get_encoding('cl100k_base'); \
|
||||
print('Tiktoken encoding cached successfully')" && break; \
|
||||
if [ $i -lt $MAX_RETRIES ]; then \
|
||||
echo "Attempt $i failed, retrying in ${RETRY_DELAY}s..."; \
|
||||
sleep $RETRY_DELAY; \
|
||||
RETRY_DELAY=$((RETRY_DELAY * 2)); \
|
||||
fi; \
|
||||
done; \
|
||||
if [ $i -eq $MAX_RETRIES ]; then \
|
||||
echo "ERROR: Failed to download tiktoken encoding after $MAX_RETRIES attempts"; \
|
||||
exit 1; \
|
||||
fi
|
||||
|
||||
# Pre-download ML models to avoid runtime download (conditional)
|
||||
# Only runs if both PRELOAD_ML_MODELS=true AND INCLUDE_LOCAL_MODELS=true
|
||||
# Includes retry logic with exponential backoff for transient network failures
|
||||
@@ -315,7 +359,6 @@ print('Downloading embedding model...'); \
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5'); \
|
||||
print('Downloading cross-encoder model...'); \
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \
|
||||
print('Downloading tiktoken encoding...'); import tiktoken; tiktoken.get_encoding('cl100k_base'); \
|
||||
print('Models cached successfully')" && break; \
|
||||
if [ $i -lt $MAX_RETRIES ]; then \
|
||||
echo "Attempt $i failed, retrying in ${RETRY_DELAY}s..."; \
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.4.9
|
||||
appVersion: "0.4.9"
|
||||
version: 0.4.10
|
||||
appVersion: "0.4.10"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -127,6 +127,38 @@ API URL for control plane
|
||||
{{- printf "http://%s-api:%d" (include "hindsight.fullname" .) (.Values.api.service.port | int) }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
TEI reranker labels
|
||||
*/}}
|
||||
{{- define "hindsight.tei.reranker.labels" -}}
|
||||
{{ include "hindsight.labels" . }}
|
||||
app.kubernetes.io/component: tei-reranker
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
TEI reranker selector labels
|
||||
*/}}
|
||||
{{- define "hindsight.tei.reranker.selectorLabels" -}}
|
||||
{{ include "hindsight.selectorLabels" . }}
|
||||
app.kubernetes.io/component: tei-reranker
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
TEI embedding labels
|
||||
*/}}
|
||||
{{- define "hindsight.tei.embedding.labels" -}}
|
||||
{{ include "hindsight.labels" . }}
|
||||
app.kubernetes.io/component: tei-embedding
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
TEI embedding selector labels
|
||||
*/}}
|
||||
{{- define "hindsight.tei.embedding.selectorLabels" -}}
|
||||
{{ include "hindsight.selectorLabels" . }}
|
||||
app.kubernetes.io/component: tei-embedding
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Get the name of the secret to use
|
||||
*/}}
|
||||
|
||||
@@ -33,7 +33,7 @@ spec:
|
||||
- name: api
|
||||
securityContext:
|
||||
{{- toYaml .Values.securityContext | nindent 10 }}
|
||||
image: "{{ .Values.api.image.repository }}:{{ .Values.api.image.tag | default .Values.version }}"
|
||||
image: "{{ .Values.api.image.repository }}:{{ .Values.api.image.tag | default .Values.version | default .Chart.AppVersion }}"
|
||||
imagePullPolicy: {{ .Values.api.image.pullPolicy }}
|
||||
ports:
|
||||
- name: http
|
||||
@@ -60,10 +60,25 @@ spec:
|
||||
- name: HINDSIGHT_API_WORKER_ENABLED
|
||||
value: "false"
|
||||
{{- end }}
|
||||
{{- /* Explicitly set port to override K8s service discovery env var (HINDSIGHT_API_PORT) */}}
|
||||
- name: HINDSIGHT_API_PORT
|
||||
value: {{ .Values.api.service.targetPort | quote }}
|
||||
{{- range $key, $value := .Values.api.env }}
|
||||
- name: {{ $key }}
|
||||
value: {{ $value | quote }}
|
||||
{{- end }}
|
||||
{{- if .Values.tei.reranker.enabled }}
|
||||
- name: HINDSIGHT_API_RERANKER_PROVIDER
|
||||
value: "tei"
|
||||
- name: HINDSIGHT_API_RERANKER_TEI_URL
|
||||
value: "http://{{ include "hindsight.fullname" . }}-tei-reranker:{{ .Values.tei.reranker.port }}"
|
||||
{{- end }}
|
||||
{{- if .Values.tei.embedding.enabled }}
|
||||
- name: HINDSIGHT_API_EMBEDDINGS_PROVIDER
|
||||
value: "tei"
|
||||
- name: HINDSIGHT_API_EMBEDDINGS_TEI_URL
|
||||
value: "http://{{ include "hindsight.fullname" . }}-tei-embedding:{{ .Values.tei.embedding.port }}"
|
||||
{{- end }}
|
||||
{{- /* Only use api.secrets when not using existingSecret (for chart-managed secrets) */}}
|
||||
{{- if not .Values.existingSecret }}
|
||||
{{- range $key, $value := .Values.api.secrets }}
|
||||
@@ -84,7 +99,7 @@ spec:
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.affinity }}
|
||||
{{- with (.Values.api.affinity | default .Values.affinity) }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
||||
@@ -33,7 +33,7 @@ spec:
|
||||
- name: control-plane
|
||||
securityContext:
|
||||
{{- toYaml .Values.securityContext | nindent 10 }}
|
||||
image: "{{ .Values.controlPlane.image.repository }}:{{ .Values.controlPlane.image.tag | default .Values.version }}"
|
||||
image: "{{ .Values.controlPlane.image.repository }}:{{ .Values.controlPlane.image.tag | default .Values.version | default .Chart.AppVersion }}"
|
||||
imagePullPolicy: {{ .Values.controlPlane.image.pullPolicy }}
|
||||
ports:
|
||||
- name: http
|
||||
@@ -71,7 +71,7 @@ spec:
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.affinity }}
|
||||
{{- with (.Values.controlPlane.affinity | default .Values.affinity) }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
{{- if and .Values.api.enabled .Values.api.podDisruptionBudget.enabled }}
|
||||
apiVersion: policy/v1
|
||||
kind: PodDisruptionBudget
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-api
|
||||
labels:
|
||||
{{- include "hindsight.api.labels" . | nindent 4 }}
|
||||
spec:
|
||||
{{- if .Values.api.podDisruptionBudget.minAvailable }}
|
||||
minAvailable: {{ .Values.api.podDisruptionBudget.minAvailable }}
|
||||
{{- end }}
|
||||
{{- if .Values.api.podDisruptionBudget.maxUnavailable }}
|
||||
maxUnavailable: {{ .Values.api.podDisruptionBudget.maxUnavailable }}
|
||||
{{- end }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "hindsight.api.selectorLabels" . | nindent 6 }}
|
||||
{{- end }}
|
||||
---
|
||||
{{- if and .Values.controlPlane.enabled .Values.controlPlane.podDisruptionBudget.enabled }}
|
||||
apiVersion: policy/v1
|
||||
kind: PodDisruptionBudget
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-control-plane
|
||||
labels:
|
||||
{{- include "hindsight.controlPlane.labels" . | nindent 4 }}
|
||||
spec:
|
||||
{{- if .Values.controlPlane.podDisruptionBudget.minAvailable }}
|
||||
minAvailable: {{ .Values.controlPlane.podDisruptionBudget.minAvailable }}
|
||||
{{- end }}
|
||||
{{- if .Values.controlPlane.podDisruptionBudget.maxUnavailable }}
|
||||
maxUnavailable: {{ .Values.controlPlane.podDisruptionBudget.maxUnavailable }}
|
||||
{{- end }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "hindsight.controlPlane.selectorLabels" . | nindent 6 }}
|
||||
{{- end }}
|
||||
---
|
||||
{{- if and .Values.worker.enabled .Values.worker.podDisruptionBudget.enabled }}
|
||||
apiVersion: policy/v1
|
||||
kind: PodDisruptionBudget
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-worker
|
||||
labels:
|
||||
{{- include "hindsight.worker.labels" . | nindent 4 }}
|
||||
spec:
|
||||
{{- if .Values.worker.podDisruptionBudget.minAvailable }}
|
||||
minAvailable: {{ .Values.worker.podDisruptionBudget.minAvailable }}
|
||||
{{- end }}
|
||||
{{- if .Values.worker.podDisruptionBudget.maxUnavailable }}
|
||||
maxUnavailable: {{ .Values.worker.podDisruptionBudget.maxUnavailable }}
|
||||
{{- end }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "hindsight.worker.selectorLabels" . | nindent 6 }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,76 @@
|
||||
{{- if .Values.tei.embedding.enabled }}
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-tei-embedding
|
||||
labels:
|
||||
{{- include "hindsight.tei.embedding.labels" . | nindent 4 }}
|
||||
spec:
|
||||
replicas: {{ .Values.tei.embedding.replicaCount }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "hindsight.tei.embedding.selectorLabels" . | nindent 6 }}
|
||||
template:
|
||||
metadata:
|
||||
{{- with .Values.podAnnotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
labels:
|
||||
{{- include "hindsight.tei.embedding.selectorLabels" . | nindent 8 }}
|
||||
spec:
|
||||
{{- if .Values.serviceAccount.create }}
|
||||
serviceAccountName: {{ include "hindsight.serviceAccountName" . }}
|
||||
{{- end }}
|
||||
securityContext:
|
||||
{{- toYaml .Values.podSecurityContext | nindent 8 }}
|
||||
containers:
|
||||
- name: tei-embedding
|
||||
securityContext:
|
||||
{{- toYaml .Values.securityContext | nindent 10 }}
|
||||
image: "{{ .Values.tei.embedding.image.repository }}:{{ .Values.tei.embedding.image.tag }}"
|
||||
imagePullPolicy: {{ .Values.tei.embedding.image.pullPolicy }}
|
||||
args:
|
||||
- "--model-id"
|
||||
- {{ .Values.tei.embedding.model | quote }}
|
||||
- "--hostname"
|
||||
- "0.0.0.0"
|
||||
{{- range .Values.tei.embedding.args }}
|
||||
- {{ . | quote }}
|
||||
{{- end }}
|
||||
ports:
|
||||
- name: http
|
||||
containerPort: {{ .Values.tei.embedding.port }}
|
||||
protocol: TCP
|
||||
env:
|
||||
- name: PORT
|
||||
value: {{ .Values.tei.embedding.port | quote }}
|
||||
{{- range $key, $value := .Values.tei.embedding.env }}
|
||||
- name: {{ $key }}
|
||||
value: {{ $value | quote }}
|
||||
{{- end }}
|
||||
livenessProbe:
|
||||
{{- toYaml .Values.tei.embedding.livenessProbe | nindent 10 }}
|
||||
readinessProbe:
|
||||
{{- toYaml .Values.tei.embedding.readinessProbe | nindent 10 }}
|
||||
resources:
|
||||
{{- toYaml .Values.tei.embedding.resources | nindent 10 }}
|
||||
volumeMounts:
|
||||
- name: model-cache
|
||||
mountPath: /data
|
||||
volumes:
|
||||
- name: model-cache
|
||||
emptyDir: {}
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.affinity }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.tolerations }}
|
||||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,17 @@
|
||||
{{- if .Values.tei.embedding.enabled }}
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-tei-embedding
|
||||
labels:
|
||||
{{- include "hindsight.tei.embedding.labels" . | nindent 4 }}
|
||||
spec:
|
||||
type: ClusterIP
|
||||
ports:
|
||||
- port: {{ .Values.tei.embedding.port }}
|
||||
targetPort: http
|
||||
protocol: TCP
|
||||
name: http
|
||||
selector:
|
||||
{{- include "hindsight.tei.embedding.selectorLabels" . | nindent 4 }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,76 @@
|
||||
{{- if .Values.tei.reranker.enabled }}
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-tei-reranker
|
||||
labels:
|
||||
{{- include "hindsight.tei.reranker.labels" . | nindent 4 }}
|
||||
spec:
|
||||
replicas: {{ .Values.tei.reranker.replicaCount }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "hindsight.tei.reranker.selectorLabels" . | nindent 6 }}
|
||||
template:
|
||||
metadata:
|
||||
{{- with .Values.podAnnotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
labels:
|
||||
{{- include "hindsight.tei.reranker.selectorLabels" . | nindent 8 }}
|
||||
spec:
|
||||
{{- if .Values.serviceAccount.create }}
|
||||
serviceAccountName: {{ include "hindsight.serviceAccountName" . }}
|
||||
{{- end }}
|
||||
securityContext:
|
||||
{{- toYaml .Values.podSecurityContext | nindent 8 }}
|
||||
containers:
|
||||
- name: tei-reranker
|
||||
securityContext:
|
||||
{{- toYaml .Values.securityContext | nindent 10 }}
|
||||
image: "{{ .Values.tei.reranker.image.repository }}:{{ .Values.tei.reranker.image.tag }}"
|
||||
imagePullPolicy: {{ .Values.tei.reranker.image.pullPolicy }}
|
||||
args:
|
||||
- "--model-id"
|
||||
- {{ .Values.tei.reranker.model | quote }}
|
||||
- "--hostname"
|
||||
- "0.0.0.0"
|
||||
{{- range .Values.tei.reranker.args }}
|
||||
- {{ . | quote }}
|
||||
{{- end }}
|
||||
ports:
|
||||
- name: http
|
||||
containerPort: {{ .Values.tei.reranker.port }}
|
||||
protocol: TCP
|
||||
env:
|
||||
- name: PORT
|
||||
value: {{ .Values.tei.reranker.port | quote }}
|
||||
{{- range $key, $value := .Values.tei.reranker.env }}
|
||||
- name: {{ $key }}
|
||||
value: {{ $value | quote }}
|
||||
{{- end }}
|
||||
livenessProbe:
|
||||
{{- toYaml .Values.tei.reranker.livenessProbe | nindent 10 }}
|
||||
readinessProbe:
|
||||
{{- toYaml .Values.tei.reranker.readinessProbe | nindent 10 }}
|
||||
resources:
|
||||
{{- toYaml .Values.tei.reranker.resources | nindent 10 }}
|
||||
volumeMounts:
|
||||
- name: model-cache
|
||||
mountPath: /data
|
||||
volumes:
|
||||
- name: model-cache
|
||||
emptyDir: {}
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.affinity }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.tolerations }}
|
||||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,17 @@
|
||||
{{- if .Values.tei.reranker.enabled }}
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-tei-reranker
|
||||
labels:
|
||||
{{- include "hindsight.tei.reranker.labels" . | nindent 4 }}
|
||||
spec:
|
||||
type: ClusterIP
|
||||
ports:
|
||||
- port: {{ .Values.tei.reranker.port }}
|
||||
targetPort: http
|
||||
protocol: TCP
|
||||
name: http
|
||||
selector:
|
||||
{{- include "hindsight.tei.reranker.selectorLabels" . | nindent 4 }}
|
||||
{{- end }}
|
||||
@@ -32,7 +32,7 @@ spec:
|
||||
- name: worker
|
||||
securityContext:
|
||||
{{- toYaml .Values.securityContext | nindent 10 }}
|
||||
image: "{{ .Values.worker.image.repository }}:{{ .Values.worker.image.tag | default .Values.version }}"
|
||||
image: "{{ .Values.worker.image.repository }}:{{ .Values.worker.image.tag | default .Values.version | default .Chart.AppVersion }}"
|
||||
imagePullPolicy: {{ .Values.worker.image.pullPolicy }}
|
||||
command: ["hindsight-worker"]
|
||||
ports:
|
||||
@@ -99,7 +99,7 @@ spec:
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.affinity }}
|
||||
{{- with (.Values.worker.affinity | default .Values.affinity) }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
||||
+110
-4
@@ -1,7 +1,8 @@
|
||||
# Default values for hindsight
|
||||
|
||||
# Chart version - use this to set a consistent image tag across all components
|
||||
version: "0.1.1"
|
||||
# Global version override - use this to set a consistent image tag across all components
|
||||
# If not set, defaults to Chart.appVersion from Chart.yaml
|
||||
# version: ""
|
||||
|
||||
# Use an existing secret instead of creating one from values
|
||||
# When set, all keys from this secret are injected as environment variables via envFrom
|
||||
@@ -57,6 +58,15 @@ api:
|
||||
timeoutSeconds: 3
|
||||
failureThreshold: 3
|
||||
|
||||
# Pod disruption budget
|
||||
podDisruptionBudget:
|
||||
enabled: false
|
||||
minAvailable: 1
|
||||
# maxUnavailable: 1
|
||||
|
||||
# Pod affinity/anti-affinity (overrides global affinity for this component)
|
||||
# affinity: {}
|
||||
|
||||
# Environment variables
|
||||
env:
|
||||
#HINDSIGHT_API_LLM_PROVIDER: "groq"
|
||||
@@ -75,7 +85,7 @@ worker:
|
||||
image:
|
||||
repository: ghcr.io/vectorize-io/hindsight-api
|
||||
pullPolicy: IfNotPresent
|
||||
# tag defaults to .Values.version if not specified
|
||||
# tag: "" # defaults to .Values.version, then Chart.appVersion if not specified
|
||||
|
||||
service:
|
||||
# Service for metrics scraping (headless for StatefulSet)
|
||||
@@ -121,6 +131,15 @@ worker:
|
||||
# HTTP port for metrics/health (matches service.targetPort)
|
||||
HINDSIGHT_API_WORKER_HTTP_PORT: "8889"
|
||||
|
||||
# Pod disruption budget
|
||||
podDisruptionBudget:
|
||||
enabled: false
|
||||
minAvailable: 1
|
||||
# maxUnavailable: 1
|
||||
|
||||
# Pod affinity/anti-affinity (overrides global affinity for this component)
|
||||
# affinity: {}
|
||||
|
||||
# Secret environment variables (inherited from api.secrets if not specified)
|
||||
secrets: {}
|
||||
|
||||
@@ -164,6 +183,15 @@ controlPlane:
|
||||
timeoutSeconds: 3
|
||||
failureThreshold: 3
|
||||
|
||||
# Pod disruption budget
|
||||
podDisruptionBudget:
|
||||
enabled: false
|
||||
minAvailable: 1
|
||||
# maxUnavailable: 1
|
||||
|
||||
# Pod affinity/anti-affinity (overrides global affinity for this component)
|
||||
# affinity: {}
|
||||
|
||||
# Environment variables
|
||||
env:
|
||||
NODE_ENV: "production"
|
||||
@@ -262,9 +290,87 @@ nodeSelector: {}
|
||||
# Tolerations
|
||||
tolerations: []
|
||||
|
||||
# Affinity
|
||||
# Affinity (applied to all components unless overridden per-component)
|
||||
affinity: {}
|
||||
|
||||
# TEI (Text Embeddings Inference) - optional standalone deployments
|
||||
# for reranking and/or embedding models
|
||||
tei:
|
||||
reranker:
|
||||
enabled: false
|
||||
replicaCount: 1
|
||||
image:
|
||||
repository: ghcr.io/huggingface/text-embeddings-inference
|
||||
tag: cpu-1.8.3
|
||||
pullPolicy: IfNotPresent
|
||||
model: "cross-encoder/ms-marco-MiniLM-L-6-v2"
|
||||
port: 8090
|
||||
args:
|
||||
- "--auto-truncate"
|
||||
env:
|
||||
PAYLOAD_LIMIT: "10000000"
|
||||
MAX_CLIENT_BATCH_SIZE: "256"
|
||||
resources:
|
||||
limits:
|
||||
cpu: 2000m
|
||||
memory: 2Gi
|
||||
requests:
|
||||
cpu: 500m
|
||||
memory: 1Gi
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /health
|
||||
port: 8090
|
||||
initialDelaySeconds: 30
|
||||
periodSeconds: 10
|
||||
timeoutSeconds: 5
|
||||
failureThreshold: 6
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /health
|
||||
port: 8090
|
||||
initialDelaySeconds: 15
|
||||
periodSeconds: 5
|
||||
timeoutSeconds: 3
|
||||
failureThreshold: 3
|
||||
|
||||
embedding:
|
||||
enabled: false
|
||||
replicaCount: 1
|
||||
image:
|
||||
repository: ghcr.io/huggingface/text-embeddings-inference
|
||||
tag: cpu-1.8.3
|
||||
pullPolicy: IfNotPresent
|
||||
model: "sentence-transformers/all-MiniLM-L6-v2"
|
||||
port: 8091
|
||||
args: []
|
||||
env:
|
||||
PAYLOAD_LIMIT: "10000000"
|
||||
MAX_CLIENT_BATCH_SIZE: "256"
|
||||
resources:
|
||||
limits:
|
||||
cpu: 2000m
|
||||
memory: 2Gi
|
||||
requests:
|
||||
cpu: 500m
|
||||
memory: 1Gi
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /health
|
||||
port: 8091
|
||||
initialDelaySeconds: 30
|
||||
periodSeconds: 10
|
||||
timeoutSeconds: 5
|
||||
failureThreshold: 6
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /health
|
||||
port: 8091
|
||||
initialDelaySeconds: 15
|
||||
periodSeconds: 5
|
||||
timeoutSeconds: 3
|
||||
failureThreshold: 3
|
||||
|
||||
# Autoscaling
|
||||
autoscaling:
|
||||
enabled: false
|
||||
|
||||
@@ -46,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.4.9"
|
||||
__version__ = "0.4.10"
|
||||
|
||||
@@ -6,7 +6,6 @@ Provides both HTTP REST API and MCP (Model Context Protocol) server.
|
||||
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
@@ -46,14 +45,14 @@ def create_app(
|
||||
# Both HTTP and MCP
|
||||
app = create_app(memory, mcp_api_enabled=True)
|
||||
"""
|
||||
mcp_app = None
|
||||
mcp_servers = None
|
||||
|
||||
# Create MCP app first if enabled (we need its lifespan for chaining)
|
||||
# Create MCP servers first if enabled (we need their lifespans for chaining)
|
||||
if mcp_api_enabled:
|
||||
try:
|
||||
from .mcp import create_mcp_app
|
||||
from .mcp import MCPMiddleware, create_mcp_servers
|
||||
|
||||
mcp_app = create_mcp_app(memory=memory)
|
||||
mcp_servers = create_mcp_servers(memory=memory)
|
||||
except ImportError as e:
|
||||
logger.error(f"MCP server requested but dependencies not available: {e}")
|
||||
logger.error("Install with: pip install hindsight-api[mcp]")
|
||||
@@ -70,11 +69,9 @@ def create_app(
|
||||
app = FastAPI(title="Hindsight API", version="0.0.7")
|
||||
logger.info("HTTP REST API disabled")
|
||||
|
||||
# Mount MCP server and chain its lifespan if enabled
|
||||
if mcp_app is not None:
|
||||
# Get both MCP apps' underlying Starlette apps for lifespan access
|
||||
multi_bank_starlette_app = mcp_app.multi_bank_app
|
||||
single_bank_starlette_app = mcp_app.single_bank_app
|
||||
# Add MCP middleware and chain its lifespan if enabled
|
||||
if mcp_servers is not None:
|
||||
multi_bank_server, single_bank_server, multi_bank_starlette_app, single_bank_starlette_app = mcp_servers
|
||||
|
||||
# Store the original lifespan
|
||||
original_lifespan = app.router.lifespan_context
|
||||
@@ -94,8 +91,19 @@ def create_app(
|
||||
# Replace the app's lifespan with the chained version
|
||||
app.router.lifespan_context = chained_lifespan
|
||||
|
||||
# Mount the MCP middleware
|
||||
app.mount(mcp_mount_path, mcp_app)
|
||||
# Add MCP as a wrapping middleware — intercepts /mcp* requests directly,
|
||||
# passes everything else through to the FastAPI app. No Starlette Mount
|
||||
# means no 307 redirect for /mcp (no trailing slash).
|
||||
app.add_middleware(
|
||||
MCPMiddleware,
|
||||
memory=memory,
|
||||
prefix=mcp_mount_path,
|
||||
multi_bank_app=multi_bank_starlette_app,
|
||||
single_bank_app=single_bank_starlette_app,
|
||||
multi_bank_server=multi_bank_server,
|
||||
single_bank_server=single_bank_server,
|
||||
)
|
||||
|
||||
logger.info(f"MCP server enabled at {mcp_mount_path}/")
|
||||
|
||||
return app
|
||||
|
||||
@@ -32,9 +32,44 @@ def _parse_metadata(metadata: Any) -> dict[str, Any]:
|
||||
return {}
|
||||
|
||||
|
||||
from typing import Callable
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
|
||||
|
||||
def FieldWithDefault(default_factory: Callable, **kwargs) -> Any:
|
||||
"""
|
||||
Field wrapper that ensures default_factory values appear in OpenAPI schema.
|
||||
|
||||
Pydantic doesn't include default_factory in OpenAPI schemas, causing OpenAPI
|
||||
Generator to make fields Optional with default=None instead of non-optional
|
||||
with the correct default value.
|
||||
|
||||
This wrapper adds json_schema_extra to include the default in the schema.
|
||||
"""
|
||||
# Determine the default value for the schema based on the factory
|
||||
if default_factory is list:
|
||||
schema_default = []
|
||||
elif default_factory is dict:
|
||||
schema_default = {}
|
||||
else:
|
||||
# For custom factories (like IncludeOptions), use empty dict as placeholder
|
||||
schema_default = {}
|
||||
|
||||
# Add or merge json_schema_extra
|
||||
json_extra = kwargs.pop("json_schema_extra", {})
|
||||
if isinstance(json_extra, dict):
|
||||
json_extra["default"] = schema_default
|
||||
else:
|
||||
# If json_schema_extra was a function, we can't merge easily
|
||||
# Fall back to just setting default
|
||||
json_extra = {"default": schema_default}
|
||||
|
||||
return Field(default_factory=default_factory, json_schema_extra=json_extra, **kwargs)
|
||||
|
||||
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
from hindsight_api.engine.memory_engine import Budget, _get_tiktoken_encoding, fq_table
|
||||
from hindsight_api.engine.reflect.observations import Observation
|
||||
@@ -103,8 +138,8 @@ class RecallRequest(BaseModel):
|
||||
query_timestamp: str | None = Field(
|
||||
default=None, description="ISO format date string (e.g., '2023-05-30T23:40:00')"
|
||||
)
|
||||
include: IncludeOptions = Field(
|
||||
default_factory=IncludeOptions,
|
||||
include: IncludeOptions = FieldWithDefault(
|
||||
IncludeOptions,
|
||||
description="Options for including additional data (entities are included by default)",
|
||||
)
|
||||
tags: list[str] | None = Field(
|
||||
@@ -570,18 +605,16 @@ class ReflectLLMCall(BaseModel):
|
||||
class ReflectBasedOn(BaseModel):
|
||||
"""Evidence the response is based on: memories, mental models, and directives."""
|
||||
|
||||
memories: list[ReflectFact] = Field(default_factory=list, description="Memory facts used to generate the response")
|
||||
mental_models: list[ReflectMentalModel] = Field(
|
||||
default_factory=list, description="Mental models used during reflection"
|
||||
)
|
||||
directives: list[ReflectDirective] = Field(default_factory=list, description="Directives applied during reflection")
|
||||
memories: list[ReflectFact] = FieldWithDefault(list, description="Memory facts used to generate the response")
|
||||
mental_models: list[ReflectMentalModel] = FieldWithDefault(list, description="Mental models used during reflection")
|
||||
directives: list[ReflectDirective] = FieldWithDefault(list, description="Directives applied during reflection")
|
||||
|
||||
|
||||
class ReflectTrace(BaseModel):
|
||||
"""Execution trace of LLM and tool calls during reflection."""
|
||||
|
||||
tool_calls: list[ReflectToolCall] = Field(default_factory=list, description="Tool calls made during reflection")
|
||||
llm_calls: list[ReflectLLMCall] = Field(default_factory=list, description="LLM calls made during reflection")
|
||||
tool_calls: list[ReflectToolCall] = FieldWithDefault(list, description="Tool calls made during reflection")
|
||||
llm_calls: list[ReflectLLMCall] = FieldWithDefault(list, description="LLM calls made during reflection")
|
||||
|
||||
|
||||
class ReflectResponse(BaseModel):
|
||||
@@ -942,7 +975,7 @@ class DocumentResponse(BaseModel):
|
||||
created_at: str
|
||||
updated_at: str
|
||||
memory_unit_count: int
|
||||
tags: list[str] = Field(default_factory=list, description="Tags associated with this document")
|
||||
tags: list[str] = FieldWithDefault(list, description="Tags associated with this document")
|
||||
|
||||
|
||||
class DeleteDocumentResponse(BaseModel):
|
||||
@@ -1066,7 +1099,7 @@ class DirectiveResponse(BaseModel):
|
||||
content: str
|
||||
priority: int = 0
|
||||
is_active: bool = True
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
tags: list[str] = FieldWithDefault(list)
|
||||
created_at: str | None = None
|
||||
updated_at: str | None = None
|
||||
|
||||
@@ -1084,7 +1117,7 @@ class CreateDirectiveRequest(BaseModel):
|
||||
content: str = Field(description="The directive text to inject into prompts")
|
||||
priority: int = Field(default=0, description="Higher priority directives are injected first")
|
||||
is_active: bool = Field(default=True, description="Whether this directive is active")
|
||||
tags: list[str] = Field(default_factory=list, description="Tags for filtering")
|
||||
tags: list[str] = FieldWithDefault(list, description="Tags for filtering")
|
||||
|
||||
|
||||
class UpdateDirectiveRequest(BaseModel):
|
||||
@@ -1121,9 +1154,9 @@ class MentalModelResponse(BaseModel):
|
||||
content: str = Field(
|
||||
description="The mental model content as well-formatted markdown (auto-generated from reflect endpoint)"
|
||||
)
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
tags: list[str] = FieldWithDefault(list)
|
||||
max_tokens: int = Field(default=2048)
|
||||
trigger: MentalModelTrigger = Field(default_factory=MentalModelTrigger)
|
||||
trigger: MentalModelTrigger = FieldWithDefault(MentalModelTrigger)
|
||||
last_refreshed_at: str | None = None
|
||||
created_at: str | None = None
|
||||
reflect_response: dict | None = Field(
|
||||
@@ -1159,9 +1192,9 @@ class CreateMentalModelRequest(BaseModel):
|
||||
)
|
||||
name: str = Field(description="Human-readable name for the mental model")
|
||||
source_query: str = Field(description="The query to run to generate content")
|
||||
tags: list[str] = Field(default_factory=list, description="Tags for scoped visibility")
|
||||
tags: list[str] = FieldWithDefault(list, description="Tags for scoped visibility")
|
||||
max_tokens: int = Field(default=2048, ge=256, le=8192, description="Maximum tokens for generated content")
|
||||
trigger: MentalModelTrigger = Field(default_factory=MentalModelTrigger, description="Trigger settings")
|
||||
trigger: MentalModelTrigger = FieldWithDefault(MentalModelTrigger, description="Trigger settings")
|
||||
|
||||
|
||||
class CreateMentalModelResponse(BaseModel):
|
||||
@@ -1400,6 +1433,26 @@ def create_app(
|
||||
app.state.prometheus_reader = None
|
||||
# Metrics collector is already initialized as no-op by default
|
||||
|
||||
# Initialize OpenTelemetry tracing if enabled
|
||||
if config.otel_traces_enabled:
|
||||
if not config.otel_exporter_otlp_endpoint:
|
||||
logging.warning("OTEL tracing enabled but no endpoint configured. Tracing disabled.")
|
||||
else:
|
||||
from hindsight_api.tracing import create_span_recorder, initialize_tracing
|
||||
|
||||
try:
|
||||
initialize_tracing(
|
||||
service_name=config.otel_service_name,
|
||||
endpoint=config.otel_exporter_otlp_endpoint,
|
||||
headers=config.otel_exporter_otlp_headers,
|
||||
deployment_environment=config.otel_deployment_environment,
|
||||
)
|
||||
create_span_recorder()
|
||||
logging.info("OpenTelemetry tracing enabled and configured")
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to initialize tracing: {e}")
|
||||
logging.warning("Continuing without tracing")
|
||||
|
||||
# Startup: Initialize database and memory system (migrations run inside initialize if enabled)
|
||||
if initialize_memory:
|
||||
await memory.initialize()
|
||||
@@ -2334,23 +2387,6 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""Get a mental model by ID."""
|
||||
try:
|
||||
# Pre-operation validation hook
|
||||
validator = app.state.memory._operation_validator
|
||||
if validator:
|
||||
from hindsight_api.extensions.operation_validator import MentalModelGetContext
|
||||
|
||||
ctx = MentalModelGetContext(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
validation = await validator.validate_mental_model_get(ctx)
|
||||
if not validation.allowed:
|
||||
raise OperationValidationError(
|
||||
validation.reason or "Operation not allowed",
|
||||
status_code=validation.status_code,
|
||||
)
|
||||
|
||||
mental_model = await app.state.memory.get_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
@@ -2359,25 +2395,6 @@ def _register_routes(app: FastAPI):
|
||||
if mental_model is None:
|
||||
raise HTTPException(status_code=404, detail=f"Mental model '{mental_model_id}' not found")
|
||||
|
||||
# Post-operation hook
|
||||
if validator:
|
||||
from hindsight_api.extensions.operation_validator import MentalModelGetResult
|
||||
|
||||
content = mental_model.get("content", "")
|
||||
output_tokens = len(content) // 4 if content else 0
|
||||
|
||||
result_ctx = MentalModelGetResult(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=request_context,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
try:
|
||||
await validator.on_mental_model_get_complete(result_ctx)
|
||||
except Exception as hook_err:
|
||||
logger.warning(f"Post-mental-model-get hook error (non-fatal): {hook_err}")
|
||||
|
||||
return MentalModelResponse(**mental_model)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
@@ -2407,23 +2424,6 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""Create a mental model (async - returns operation_id)."""
|
||||
try:
|
||||
# Pre-operation validation hook
|
||||
validator = app.state.memory._operation_validator
|
||||
if validator:
|
||||
from hindsight_api.extensions.operation_validator import MentalModelRefreshContext
|
||||
|
||||
ctx = MentalModelRefreshContext(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=None, # Not yet created
|
||||
request_context=request_context,
|
||||
)
|
||||
validation = await validator.validate_mental_model_refresh(ctx)
|
||||
if not validation.allowed:
|
||||
raise OperationValidationError(
|
||||
validation.reason or "Operation not allowed",
|
||||
status_code=validation.status_code,
|
||||
)
|
||||
|
||||
# 1. Create the mental model with placeholder content
|
||||
mental_model = await app.state.memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
@@ -2471,23 +2471,6 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""Refresh a mental model by re-running its source query (async)."""
|
||||
try:
|
||||
# Pre-operation validation hook
|
||||
validator = app.state.memory._operation_validator
|
||||
if validator:
|
||||
from hindsight_api.extensions.operation_validator import MentalModelRefreshContext
|
||||
|
||||
ctx = MentalModelRefreshContext(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
validation = await validator.validate_mental_model_refresh(ctx)
|
||||
if not validation.allowed:
|
||||
raise OperationValidationError(
|
||||
validation.reason or "Operation not allowed",
|
||||
status_code=validation.status_code,
|
||||
)
|
||||
|
||||
result = await app.state.memory.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
|
||||
@@ -43,6 +43,10 @@ _current_bank_id: ContextVar[str | None] = ContextVar("current_bank_id", default
|
||||
# Context variable to hold the current API key (for tenant auth propagation)
|
||||
_current_api_key: ContextVar[str | None] = ContextVar("current_api_key", default=None)
|
||||
|
||||
# Context variables for tenant_id and api_key_id (set by authenticate, used by usage metering)
|
||||
_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)
|
||||
|
||||
|
||||
def get_current_bank_id() -> str | None:
|
||||
"""Get the current bank_id from context."""
|
||||
@@ -54,6 +58,16 @@ def get_current_api_key() -> str | None:
|
||||
return _current_api_key.get()
|
||||
|
||||
|
||||
def get_current_tenant_id() -> str | None:
|
||||
"""Get the current tenant_id from context."""
|
||||
return _current_tenant_id.get()
|
||||
|
||||
|
||||
def get_current_api_key_id() -> str | None:
|
||||
"""Get the current api_key_id from context."""
|
||||
return _current_api_key_id.get()
|
||||
|
||||
|
||||
def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"""
|
||||
Create and configure the Hindsight MCP server.
|
||||
@@ -73,8 +87,22 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=get_current_bank_id,
|
||||
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
|
||||
include_bank_id_param=multi_bank,
|
||||
tools=None if multi_bank else {"retain", "recall", "reflect"}, # Scoped tools for single-bank mode
|
||||
tools=None
|
||||
if multi_bank
|
||||
else {
|
||||
"retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"create_mental_model",
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
}, # Scoped tools for single-bank mode (excludes bank management: list_banks, create_bank)
|
||||
retain_fire_and_forget=False, # HTTP MCP supports sync/async modes
|
||||
)
|
||||
|
||||
@@ -90,7 +118,10 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
|
||||
|
||||
class MCPMiddleware:
|
||||
"""ASGI middleware that handles authentication and routes to appropriate MCP server.
|
||||
"""ASGI middleware that intercepts MCP requests and routes to appropriate MCP server.
|
||||
|
||||
This middleware wraps the main FastAPI app and intercepts requests matching the
|
||||
configured prefix (default: /mcp). Non-MCP requests pass through to the inner app.
|
||||
|
||||
Authentication:
|
||||
1. If HINDSIGHT_API_MCP_AUTH_TOKEN is set (legacy), validates against that token
|
||||
@@ -121,27 +152,33 @@ class MCPMiddleware:
|
||||
--header "X-Bank-Id: my-bank" --header "Authorization: Bearer <token>"
|
||||
"""
|
||||
|
||||
def __init__(self, app, memory: MemoryEngine):
|
||||
def __init__(
|
||||
self,
|
||||
app,
|
||||
memory: MemoryEngine,
|
||||
prefix: str = "/mcp",
|
||||
multi_bank_app=None,
|
||||
single_bank_app=None,
|
||||
multi_bank_server=None,
|
||||
single_bank_server=None,
|
||||
):
|
||||
self.app = app
|
||||
self.prefix = prefix
|
||||
self.memory = memory
|
||||
self.tenant_extension = memory._tenant_extension
|
||||
|
||||
# Create two server instances:
|
||||
# 1. Multi-bank server (for /mcp/ root endpoint)
|
||||
self.multi_bank_server = create_mcp_server(memory, multi_bank=True)
|
||||
self.multi_bank_app = self.multi_bank_server.http_app(path="/")
|
||||
|
||||
# 2. Single-bank server (for /mcp/{bank_id}/ endpoints)
|
||||
self.single_bank_server = create_mcp_server(memory, multi_bank=False)
|
||||
self.single_bank_app = self.single_bank_server.http_app(path="/")
|
||||
|
||||
# Backward compatibility: expose multi_bank_app as mcp_app
|
||||
self.mcp_app = self.multi_bank_app
|
||||
|
||||
# Expose the lifespan for the parent app to chain (use multi-bank as default)
|
||||
self.lifespan = (
|
||||
self.multi_bank_app.lifespan_handler if hasattr(self.multi_bank_app, "lifespan_handler") else None
|
||||
)
|
||||
if multi_bank_app and single_bank_app:
|
||||
# Pre-created servers (used when called via add_middleware from create_app)
|
||||
self.multi_bank_app = multi_bank_app
|
||||
self.single_bank_app = single_bank_app
|
||||
self.multi_bank_server = multi_bank_server
|
||||
self.single_bank_server = single_bank_server
|
||||
else:
|
||||
# Create servers internally (for direct construction / tests)
|
||||
self.multi_bank_server = create_mcp_server(memory, multi_bank=True)
|
||||
self.multi_bank_app = self.multi_bank_server.http_app(path="/")
|
||||
self.single_bank_server = create_mcp_server(memory, multi_bank=False)
|
||||
self.single_bank_app = self.single_bank_server.http_app(path="/")
|
||||
|
||||
def _get_header(self, scope: dict, name: str) -> str | None:
|
||||
"""Extract a header value from ASGI scope."""
|
||||
@@ -153,9 +190,20 @@ class MCPMiddleware:
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] != "http":
|
||||
await self.multi_bank_app(scope, receive, send)
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
path = scope.get("path", "")
|
||||
|
||||
# Check if this is an MCP request (matches prefix)
|
||||
if not (path == self.prefix or path.startswith(self.prefix + "/")):
|
||||
# Not an MCP request — pass through to the inner app
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
# Strip prefix from path
|
||||
path = path[len(self.prefix) :] or "/"
|
||||
|
||||
# Extract auth token from header (for tenant auth propagation)
|
||||
auth_header = self._get_header(scope, "Authorization")
|
||||
auth_token: str | None = None
|
||||
@@ -165,6 +213,8 @@ class MCPMiddleware:
|
||||
|
||||
# Authenticate: check legacy MCP_AUTH_TOKEN first, then TenantExtension
|
||||
tenant_context = None
|
||||
auth_tenant_id: str | None = None
|
||||
auth_api_key_id: str | None = None
|
||||
if MCP_AUTH_TOKEN:
|
||||
# Legacy authentication mode - validate against static token
|
||||
if not auth_token:
|
||||
@@ -178,7 +228,11 @@ class MCPMiddleware:
|
||||
else:
|
||||
# Use TenantExtension.authenticate_mcp() for auth
|
||||
try:
|
||||
tenant_context = await self.tenant_extension.authenticate_mcp(RequestContext(api_key=auth_token))
|
||||
auth_context = RequestContext(api_key=auth_token)
|
||||
tenant_context = await self.tenant_extension.authenticate_mcp(auth_context)
|
||||
# Capture tenant_id and api_key_id set by authenticate() for usage metering
|
||||
auth_tenant_id = auth_context.tenant_id
|
||||
auth_api_key_id = auth_context.api_key_id
|
||||
except AuthenticationError as e:
|
||||
await self._send_error(send, 401, str(e))
|
||||
return
|
||||
@@ -188,36 +242,15 @@ class MCPMiddleware:
|
||||
_current_schema.set(tenant_context.schema_name) if tenant_context and tenant_context.schema_name else None
|
||||
)
|
||||
|
||||
path = scope.get("path", "")
|
||||
|
||||
# Strip any mount prefix (e.g., /mcp) that FastAPI might not have stripped
|
||||
root_path = scope.get("root_path", "")
|
||||
if root_path and path.startswith(root_path):
|
||||
path = path[len(root_path) :] or "/"
|
||||
|
||||
# Also handle case where mount path wasn't stripped (e.g., /mcp/...)
|
||||
if path.startswith("/mcp/"):
|
||||
path = path[4:] # Remove /mcp prefix
|
||||
elif path == "/mcp":
|
||||
path = "/"
|
||||
|
||||
# Ensure path has leading slash (needed after stripping mount path)
|
||||
if path and not path.startswith("/"):
|
||||
path = "/" + path
|
||||
|
||||
# Try to get bank_id from header first (for Claude Code compatibility)
|
||||
bank_id = self._get_header(scope, "X-Bank-Id")
|
||||
bank_id_from_path = False
|
||||
|
||||
# MCP endpoint paths that should not be treated as bank_ids
|
||||
MCP_ENDPOINTS = {"sse", "messages"}
|
||||
|
||||
# If no header, try to extract from path: /{bank_id}/...
|
||||
new_path = path
|
||||
if not bank_id and path.startswith("/") and len(path) > 1:
|
||||
parts = path[1:].split("/", 1)
|
||||
# Don't treat MCP endpoints as bank_ids
|
||||
if parts[0] and parts[0] not in MCP_ENDPOINTS:
|
||||
if parts[0]:
|
||||
# First segment looks like a bank_id
|
||||
bank_id = parts[0]
|
||||
bank_id_from_path = True
|
||||
@@ -233,19 +266,32 @@ 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 and api_key context
|
||||
# Set bank_id, api_key, tenant_id, and api_key_id 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
|
||||
try:
|
||||
new_scope = scope.copy()
|
||||
new_scope["path"] = new_path
|
||||
# Clear root_path since we're passing directly to the app
|
||||
new_scope["root_path"] = ""
|
||||
|
||||
# Wrap send to rewrite the SSE endpoint URL to include bank_id if using path-based routing
|
||||
# Wrap send to rewrite the SSE endpoint URL to include bank_id if using path-based routing.
|
||||
# Only rewrite SSE (text/event-stream) responses to avoid corrupting tool results
|
||||
# that might contain the literal string "data: /messages".
|
||||
is_sse_response = False
|
||||
|
||||
async def send_wrapper(message):
|
||||
if message["type"] == "http.response.body" and bank_id_from_path:
|
||||
nonlocal is_sse_response
|
||||
if message["type"] == "http.response.start":
|
||||
for header_name, header_value in message.get("headers", []):
|
||||
if header_name == b"content-type" and b"text/event-stream" in header_value:
|
||||
is_sse_response = True
|
||||
break
|
||||
if message["type"] == "http.response.body" and bank_id_from_path and is_sse_response:
|
||||
body = message.get("body", b"")
|
||||
if body and b"/messages" in body:
|
||||
# Rewrite /messages to /{bank_id}/messages in SSE endpoint event
|
||||
@@ -258,6 +304,10 @@ class MCPMiddleware:
|
||||
_current_bank_id.reset(bank_id_token)
|
||||
if api_key_token is not None:
|
||||
_current_api_key.reset(api_key_token)
|
||||
if tenant_id_token is not None:
|
||||
_current_tenant_id.reset(tenant_id_token)
|
||||
if api_key_id_token is not None:
|
||||
_current_api_key_id.reset(api_key_id_token)
|
||||
if schema_token is not None:
|
||||
_current_schema.reset(schema_token)
|
||||
|
||||
@@ -279,30 +329,19 @@ class MCPMiddleware:
|
||||
)
|
||||
|
||||
|
||||
def create_mcp_app(memory: MemoryEngine):
|
||||
"""
|
||||
Create an ASGI app that handles MCP requests with dynamic tool exposure.
|
||||
def create_mcp_servers(memory: MemoryEngine):
|
||||
"""Create multi-bank and single-bank MCP servers and their Starlette apps.
|
||||
|
||||
Authentication:
|
||||
Uses the TenantExtension from the MemoryEngine (same auth as REST API).
|
||||
|
||||
Two modes based on URL structure:
|
||||
|
||||
1. Single-bank mode (recommended for agent isolation):
|
||||
- URL: /mcp/{bank_id}/
|
||||
- Tools: retain, recall, reflect (no bank_id parameter)
|
||||
- Example: claude mcp add --transport http my-agent http://localhost:8888/mcp/my-agent-bank/
|
||||
|
||||
2. Multi-bank mode (for cross-bank operations):
|
||||
- URL: /mcp/
|
||||
- Tools: retain, recall, reflect, list_banks, create_bank (all with bank_id parameter)
|
||||
- Bank ID from: X-Bank-Id header or HINDSIGHT_MCP_BANK_ID env var (default: "default")
|
||||
- Example: claude mcp add --transport http hindsight http://localhost:8888/mcp --header "X-Bank-Id: my-bank"
|
||||
|
||||
Args:
|
||||
memory: MemoryEngine instance
|
||||
Returns the servers and apps separately so lifespans can be chained before
|
||||
the middleware wraps the main app.
|
||||
|
||||
Returns:
|
||||
ASGI application
|
||||
Tuple of (multi_bank_server, single_bank_server, multi_bank_app, single_bank_app)
|
||||
"""
|
||||
return MCPMiddleware(None, memory)
|
||||
multi_bank_server = create_mcp_server(memory, multi_bank=True)
|
||||
multi_bank_app = multi_bank_server.http_app(path="/")
|
||||
|
||||
single_bank_server = create_mcp_server(memory, multi_bank=False)
|
||||
single_bank_app = single_bank_server.http_app(path="/")
|
||||
|
||||
return multi_bank_server, single_bank_server, multi_bank_app, single_bank_app
|
||||
|
||||
@@ -66,27 +66,40 @@ ENV_CONSOLIDATION_LLM_TIMEOUT = "HINDSIGHT_API_CONSOLIDATION_LLM_TIMEOUT"
|
||||
ENV_EMBEDDINGS_PROVIDER = "HINDSIGHT_API_EMBEDDINGS_PROVIDER"
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL"
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU = "HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU"
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE = "HINDSIGHT_API_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE"
|
||||
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_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY"
|
||||
# Cohere configuration (separate for embeddings and reranker)
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY = "HINDSIGHT_API_EMBEDDINGS_COHERE_API_KEY"
|
||||
ENV_EMBEDDINGS_COHERE_MODEL = "HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL"
|
||||
ENV_EMBEDDINGS_COHERE_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL"
|
||||
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"
|
||||
|
||||
# LiteLLM gateway configuration (for embeddings and reranker via LiteLLM proxy)
|
||||
# Deprecated: Legacy shared Cohere API key (for backward compatibility)
|
||||
ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY"
|
||||
|
||||
# LiteLLM configuration (separate for embeddings and reranker)
|
||||
ENV_EMBEDDINGS_LITELLM_API_BASE = "HINDSIGHT_API_EMBEDDINGS_LITELLM_API_BASE"
|
||||
ENV_EMBEDDINGS_LITELLM_API_KEY = "HINDSIGHT_API_EMBEDDINGS_LITELLM_API_KEY"
|
||||
ENV_EMBEDDINGS_LITELLM_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL"
|
||||
ENV_RERANKER_LITELLM_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_API_BASE"
|
||||
ENV_RERANKER_LITELLM_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_API_KEY"
|
||||
ENV_RERANKER_LITELLM_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_MODEL"
|
||||
|
||||
# Deprecated: Legacy shared LiteLLM config (for backward compatibility)
|
||||
ENV_LITELLM_API_BASE = "HINDSIGHT_API_LITELLM_API_BASE"
|
||||
ENV_LITELLM_API_KEY = "HINDSIGHT_API_LITELLM_API_KEY"
|
||||
ENV_EMBEDDINGS_LITELLM_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL"
|
||||
ENV_RERANKER_LITELLM_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_MODEL"
|
||||
|
||||
ENV_RERANKER_PROVIDER = "HINDSIGHT_API_RERANKER_PROVIDER"
|
||||
ENV_RERANKER_LOCAL_MODEL = "HINDSIGHT_API_RERANKER_LOCAL_MODEL"
|
||||
ENV_RERANKER_LOCAL_FORCE_CPU = "HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU"
|
||||
ENV_RERANKER_LOCAL_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_LOCAL_MAX_CONCURRENT"
|
||||
ENV_RERANKER_LOCAL_TRUST_REMOTE_CODE = "HINDSIGHT_API_RERANKER_LOCAL_TRUST_REMOTE_CODE"
|
||||
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"
|
||||
@@ -108,6 +121,13 @@ ENV_MCP_LOCAL_BANK_ID = "HINDSIGHT_API_MCP_LOCAL_BANK_ID"
|
||||
ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
|
||||
ENV_MENTAL_MODEL_REFRESH_CONCURRENCY = "HINDSIGHT_API_MENTAL_MODEL_REFRESH_CONCURRENCY"
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
ENV_OTEL_TRACES_ENABLED = "HINDSIGHT_API_OTEL_TRACES_ENABLED"
|
||||
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"
|
||||
|
||||
# Vertex AI configuration
|
||||
ENV_LLM_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"
|
||||
ENV_LLM_VERTEXAI_REGION = "HINDSIGHT_API_LLM_VERTEXAI_REGION"
|
||||
@@ -183,6 +203,7 @@ DEFAULT_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY = None # Optional, uses ADC if not set
|
||||
DEFAULT_EMBEDDINGS_PROVIDER = "local"
|
||||
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_EMBEDDING_DIMENSION = 384
|
||||
|
||||
@@ -190,6 +211,9 @@ DEFAULT_RERANKER_PROVIDER = "local"
|
||||
DEFAULT_RERANKER_LOCAL_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
|
||||
DEFAULT_RERANKER_LOCAL_FORCE_CPU = False # Force CPU mode for local reranker (avoids MPS/XPC issues on macOS)
|
||||
DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT = 4 # Limit concurrent CPU-bound reranking to prevent thrashing
|
||||
DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE = (
|
||||
False # Security: disabled by default, required for some models like jina-reranker-v2
|
||||
)
|
||||
DEFAULT_RERANKER_TEI_BATCH_SIZE = 128
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT = 8
|
||||
DEFAULT_RERANKER_MAX_CANDIDATES = 300
|
||||
@@ -251,6 +275,11 @@ DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks
|
||||
# Reflect agent settings
|
||||
DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
DEFAULT_OTEL_TRACES_ENABLED = False # Disabled by default for backward compatibility
|
||||
DEFAULT_OTEL_SERVICE_NAME = "hindsight-api"
|
||||
DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT = "development"
|
||||
|
||||
# Default MCP tool descriptions (can be customized via env vars)
|
||||
DEFAULT_MCP_RETAIN_DESCRIPTION = """Store important information to long-term memory.
|
||||
|
||||
@@ -381,20 +410,32 @@ class HindsightConfig:
|
||||
embeddings_provider: str
|
||||
embeddings_local_model: str
|
||||
embeddings_local_force_cpu: bool
|
||||
embeddings_local_trust_remote_code: bool
|
||||
embeddings_tei_url: str | None
|
||||
embeddings_openai_base_url: str | None
|
||||
embeddings_cohere_api_key: str | None
|
||||
embeddings_cohere_model: str
|
||||
embeddings_cohere_base_url: str | None
|
||||
embeddings_litellm_api_base: str
|
||||
embeddings_litellm_api_key: str | None
|
||||
embeddings_litellm_model: str
|
||||
|
||||
# Reranker
|
||||
reranker_provider: str
|
||||
reranker_local_model: str
|
||||
reranker_local_force_cpu: bool
|
||||
reranker_local_max_concurrent: int
|
||||
reranker_local_trust_remote_code: bool
|
||||
reranker_tei_url: str | None
|
||||
reranker_tei_batch_size: int
|
||||
reranker_tei_max_concurrent: int
|
||||
reranker_max_candidates: int
|
||||
reranker_cohere_api_key: str | None
|
||||
reranker_cohere_model: str
|
||||
reranker_cohere_base_url: str | None
|
||||
reranker_litellm_api_base: str
|
||||
reranker_litellm_api_key: str | None
|
||||
reranker_litellm_model: str
|
||||
|
||||
# Server
|
||||
host: str
|
||||
@@ -447,6 +488,13 @@ class HindsightConfig:
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations: int
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
otel_traces_enabled: bool
|
||||
otel_exporter_otlp_endpoint: str | None
|
||||
otel_exporter_otlp_headers: str | None
|
||||
otel_service_name: str
|
||||
otel_deployment_environment: str
|
||||
|
||||
def validate(self) -> None:
|
||||
"""Validate configuration values and raise errors for invalid combinations."""
|
||||
# RETAIN_MAX_COMPLETION_TOKENS must be greater than RETAIN_CHUNK_SIZE
|
||||
@@ -567,9 +615,21 @@ class HindsightConfig:
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU, str(DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU)
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
embeddings_local_trust_remote_code=os.getenv(
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE, str(DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE)
|
||||
).lower()
|
||||
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,
|
||||
# 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,
|
||||
# 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),
|
||||
embeddings_litellm_api_key=os.getenv(ENV_EMBEDDINGS_LITELLM_API_KEY) or os.getenv(ENV_LITELLM_API_KEY),
|
||||
embeddings_litellm_model=os.getenv(ENV_EMBEDDINGS_LITELLM_MODEL, DEFAULT_EMBEDDINGS_LITELLM_MODEL),
|
||||
# Reranker
|
||||
reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER),
|
||||
reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL),
|
||||
@@ -580,13 +640,25 @@ class HindsightConfig:
|
||||
reranker_local_max_concurrent=int(
|
||||
os.getenv(ENV_RERANKER_LOCAL_MAX_CONCURRENT, str(DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT))
|
||||
),
|
||||
reranker_local_trust_remote_code=os.getenv(
|
||||
ENV_RERANKER_LOCAL_TRUST_REMOTE_CODE, str(DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE)
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
reranker_tei_url=os.getenv(ENV_RERANKER_TEI_URL),
|
||||
reranker_tei_batch_size=int(os.getenv(ENV_RERANKER_TEI_BATCH_SIZE, str(DEFAULT_RERANKER_TEI_BATCH_SIZE))),
|
||||
reranker_tei_max_concurrent=int(
|
||||
os.getenv(ENV_RERANKER_TEI_MAX_CONCURRENT, str(DEFAULT_RERANKER_TEI_MAX_CONCURRENT))
|
||||
),
|
||||
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,
|
||||
# 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),
|
||||
reranker_litellm_api_key=os.getenv(ENV_RERANKER_LITELLM_API_KEY) or os.getenv(ENV_LITELLM_API_KEY),
|
||||
reranker_litellm_model=os.getenv(ENV_RERANKER_LITELLM_MODEL, DEFAULT_RERANKER_LITELLM_MODEL),
|
||||
# Server
|
||||
host=os.getenv(ENV_HOST, DEFAULT_HOST),
|
||||
port=int(os.getenv(ENV_PORT, DEFAULT_PORT)),
|
||||
@@ -646,6 +718,13 @@ class HindsightConfig:
|
||||
),
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
|
||||
# OpenTelemetry tracing configuration
|
||||
otel_traces_enabled=os.getenv(ENV_OTEL_TRACES_ENABLED, str(DEFAULT_OTEL_TRACES_ENABLED)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
otel_exporter_otlp_endpoint=os.getenv(ENV_OTEL_EXPORTER_OTLP_ENDPOINT) or None,
|
||||
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),
|
||||
)
|
||||
config.validate()
|
||||
return config
|
||||
|
||||
@@ -426,94 +426,109 @@ async def _process_memory(
|
||||
Returns:
|
||||
Dict with action summary: created/updated/merged counts
|
||||
"""
|
||||
from ...tracing import get_tracer, is_tracing_enabled
|
||||
|
||||
fact_text = memory["text"]
|
||||
memory_id = memory["id"]
|
||||
fact_tags = memory.get("tags") or []
|
||||
|
||||
# Find related observations using the full recall system
|
||||
# SECURITY: Pass tags to ensure observations don't leak across security boundaries
|
||||
t0 = time.time()
|
||||
related_observations = await _find_related_observations(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
query=fact_text,
|
||||
request_context=request_context,
|
||||
tags=fact_tags, # Pass source memory's tags for security
|
||||
)
|
||||
if perf:
|
||||
perf.record_timing("recall", time.time() - t0)
|
||||
# Create parent span for this memory's consolidation
|
||||
tracer = get_tracer()
|
||||
if is_tracing_enabled():
|
||||
consolidation_span = tracer.start_span("hindsight.consolidation")
|
||||
consolidation_span.set_attribute("hindsight.memory_id", str(memory_id))
|
||||
consolidation_span.set_attribute("hindsight.bank_id", bank_id)
|
||||
else:
|
||||
consolidation_span = None
|
||||
|
||||
# Single LLM call handles ALL cases (with or without existing observations)
|
||||
# Note: Tags are NOT passed to LLM - they are handled algorithmically
|
||||
t0 = time.time()
|
||||
actions = await _consolidate_with_llm(
|
||||
memory_engine=memory_engine,
|
||||
fact_text=fact_text,
|
||||
observations=related_observations, # Can be empty list
|
||||
mission=mission,
|
||||
)
|
||||
if perf:
|
||||
perf.record_timing("llm", time.time() - t0)
|
||||
try:
|
||||
# Find related observations using the full recall system
|
||||
# SECURITY: Pass tags to ensure observations don't leak across security boundaries
|
||||
t0 = time.time()
|
||||
related_observations = await _find_related_observations(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
query=fact_text,
|
||||
request_context=request_context,
|
||||
tags=fact_tags, # Pass source memory's tags for security
|
||||
)
|
||||
if perf:
|
||||
perf.record_timing("recall", time.time() - t0)
|
||||
|
||||
if not actions:
|
||||
# LLM returned empty array - fact is purely ephemeral, skip
|
||||
return {"action": "skipped", "reason": "no_durable_knowledge"}
|
||||
# Single LLM call handles ALL cases (with or without existing observations)
|
||||
# Note: Tags are NOT passed to LLM - they are handled algorithmically
|
||||
t0 = time.time()
|
||||
actions = await _consolidate_with_llm(
|
||||
memory_engine=memory_engine,
|
||||
fact_text=fact_text,
|
||||
observations=related_observations, # Can be empty list
|
||||
mission=mission,
|
||||
)
|
||||
if perf:
|
||||
perf.record_timing("llm", time.time() - t0)
|
||||
|
||||
# Execute all actions and collect results
|
||||
results = []
|
||||
for action in actions:
|
||||
action_type = action.get("action")
|
||||
if action_type == "update":
|
||||
result = await _execute_update_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
memory_id=memory_id,
|
||||
action=action,
|
||||
observations=related_observations,
|
||||
source_fact_tags=fact_tags, # Pass source fact's tags for security
|
||||
source_occurred_start=memory.get("occurred_start"),
|
||||
source_occurred_end=memory.get("occurred_end"),
|
||||
source_mentioned_at=memory.get("mentioned_at"),
|
||||
perf=perf,
|
||||
)
|
||||
results.append(result)
|
||||
elif action_type == "create":
|
||||
result = await _execute_create_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
memory_id=memory_id,
|
||||
action=action,
|
||||
source_fact_tags=fact_tags, # Pass source fact's tags for security
|
||||
event_date=memory.get("event_date"),
|
||||
occurred_start=memory.get("occurred_start"),
|
||||
occurred_end=memory.get("occurred_end"),
|
||||
mentioned_at=memory.get("mentioned_at"),
|
||||
perf=perf,
|
||||
)
|
||||
results.append(result)
|
||||
if not actions:
|
||||
# LLM returned empty array - fact is purely ephemeral, skip
|
||||
return {"action": "skipped", "reason": "no_durable_knowledge"}
|
||||
|
||||
if not results:
|
||||
# No valid actions executed
|
||||
return {"action": "skipped", "reason": "no_valid_actions"}
|
||||
# Execute all actions and collect results
|
||||
results = []
|
||||
for action in actions:
|
||||
action_type = action.get("action")
|
||||
if action_type == "update":
|
||||
result = await _execute_update_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
memory_id=memory_id,
|
||||
action=action,
|
||||
observations=related_observations,
|
||||
source_fact_tags=fact_tags, # Pass source fact's tags for security
|
||||
source_occurred_start=memory.get("occurred_start"),
|
||||
source_occurred_end=memory.get("occurred_end"),
|
||||
source_mentioned_at=memory.get("mentioned_at"),
|
||||
perf=perf,
|
||||
)
|
||||
results.append(result)
|
||||
elif action_type == "create":
|
||||
result = await _execute_create_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
memory_id=memory_id,
|
||||
action=action,
|
||||
source_fact_tags=fact_tags, # Pass source fact's tags for security
|
||||
event_date=memory.get("event_date"),
|
||||
occurred_start=memory.get("occurred_start"),
|
||||
occurred_end=memory.get("occurred_end"),
|
||||
mentioned_at=memory.get("mentioned_at"),
|
||||
perf=perf,
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
# Summarize results
|
||||
created = sum(1 for r in results if r.get("action") == "created")
|
||||
updated = sum(1 for r in results if r.get("action") == "updated")
|
||||
merged = sum(1 for r in results if r.get("action") == "merged")
|
||||
if not results:
|
||||
# No valid actions executed
|
||||
return {"action": "skipped", "reason": "no_valid_actions"}
|
||||
|
||||
if len(results) == 1:
|
||||
return results[0]
|
||||
# Summarize results
|
||||
created = sum(1 for r in results if r.get("action") == "created")
|
||||
updated = sum(1 for r in results if r.get("action") == "updated")
|
||||
merged = sum(1 for r in results if r.get("action") == "merged")
|
||||
|
||||
return {
|
||||
"action": "multiple",
|
||||
"created": created,
|
||||
"updated": updated,
|
||||
"merged": merged,
|
||||
"total_actions": len(results),
|
||||
}
|
||||
if len(results) == 1:
|
||||
return results[0]
|
||||
|
||||
return {
|
||||
"action": "multiple",
|
||||
"created": created,
|
||||
"updated": updated,
|
||||
"merged": merged,
|
||||
"total_actions": len(results),
|
||||
}
|
||||
finally:
|
||||
if consolidation_span:
|
||||
consolidation_span.end()
|
||||
|
||||
|
||||
async def _execute_update_action(
|
||||
@@ -733,22 +748,37 @@ async def _find_related_observations(
|
||||
# Use recall to find related observations with token budget
|
||||
# max_tokens naturally limits how many observations are returned
|
||||
from ...config import get_config
|
||||
from ...tracing import get_tracer, is_tracing_enabled
|
||||
|
||||
config = get_config()
|
||||
|
||||
# SECURITY: Use all_strict matching if tags provided to prevent cross-scope consolidation
|
||||
tags_match = "all_strict" if tags else "any"
|
||||
|
||||
recall_result = await memory_engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
max_tokens=config.consolidation_max_tokens, # Token budget for observations (configurable)
|
||||
fact_type=["observation"], # Only retrieve observations
|
||||
request_context=request_context,
|
||||
tags=tags, # Filter by source memory's tags
|
||||
tags_match=tags_match, # Use strict matching for security
|
||||
_quiet=True, # Suppress logging
|
||||
)
|
||||
# Create span for recall operation within consolidation
|
||||
tracer = get_tracer()
|
||||
if is_tracing_enabled():
|
||||
recall_span = tracer.start_span("hindsight.consolidation_recall")
|
||||
recall_span.set_attribute("hindsight.bank_id", bank_id)
|
||||
recall_span.set_attribute("hindsight.query", query[:100]) # Truncate for brevity
|
||||
recall_span.set_attribute("hindsight.fact_type", "observation")
|
||||
else:
|
||||
recall_span = None
|
||||
|
||||
try:
|
||||
recall_result = await memory_engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
max_tokens=config.consolidation_max_tokens, # Token budget for observations (configurable)
|
||||
fact_type=["observation"], # Only retrieve observations
|
||||
request_context=request_context,
|
||||
tags=tags, # Filter by source memory's tags
|
||||
tags_match=tags_match, # Use strict matching for security
|
||||
_quiet=True, # Suppress logging
|
||||
)
|
||||
finally:
|
||||
if recall_span:
|
||||
recall_span.end()
|
||||
|
||||
# If no observations returned, return empty list
|
||||
if not recall_result.results:
|
||||
|
||||
@@ -24,20 +24,18 @@ from ..config import (
|
||||
DEFAULT_RERANKER_LOCAL_FORCE_CPU,
|
||||
DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_LOCAL_MODEL,
|
||||
DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_RERANKER_PROVIDER,
|
||||
DEFAULT_RERANKER_TEI_BATCH_SIZE,
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
|
||||
ENV_COHERE_API_KEY,
|
||||
ENV_LITELLM_API_BASE,
|
||||
ENV_LITELLM_API_KEY,
|
||||
ENV_RERANKER_COHERE_BASE_URL,
|
||||
ENV_RERANKER_COHERE_API_KEY,
|
||||
ENV_RERANKER_COHERE_MODEL,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
ENV_RERANKER_FLASHRANK_MODEL,
|
||||
ENV_RERANKER_LITELLM_MODEL,
|
||||
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_TEI_BATCH_SIZE,
|
||||
ENV_RERANKER_TEI_MAX_CONCURRENT,
|
||||
@@ -102,7 +100,13 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
_executor: ThreadPoolExecutor | None = None
|
||||
_max_concurrent: int = 4 # Limit concurrent CPU-bound reranking calls
|
||||
|
||||
def __init__(self, model_name: str | None = None, max_concurrent: int = 4, force_cpu: bool = False):
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str | None = None,
|
||||
max_concurrent: int = 4,
|
||||
force_cpu: bool = False,
|
||||
trust_remote_code: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize local SentenceTransformers cross-encoder.
|
||||
|
||||
@@ -113,9 +117,13 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
Higher values may cause CPU thrashing under load.
|
||||
force_cpu: Force CPU mode (avoids MPS/XPC issues on macOS in daemon mode).
|
||||
Default: False
|
||||
trust_remote_code: Allow loading models with custom code (security risk).
|
||||
Required for some models like jina-reranker-v2-base-multilingual.
|
||||
Default: False (disabled for security)
|
||||
"""
|
||||
self.model_name = model_name or DEFAULT_RERANKER_LOCAL_MODEL
|
||||
self.force_cpu = force_cpu
|
||||
self.trust_remote_code = trust_remote_code
|
||||
self._model = None
|
||||
LocalSTCrossEncoder._max_concurrent = max_concurrent
|
||||
|
||||
@@ -181,6 +189,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
self.model_name,
|
||||
device=device,
|
||||
model_kwargs={"low_cpu_mem_usage": False},
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
)
|
||||
finally:
|
||||
# Restore original logging level
|
||||
@@ -847,23 +856,27 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
model_name=config.reranker_local_model,
|
||||
max_concurrent=config.reranker_local_max_concurrent,
|
||||
force_cpu=config.reranker_local_force_cpu,
|
||||
trust_remote_code=config.reranker_local_trust_remote_code,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = os.environ.get(ENV_COHERE_API_KEY)
|
||||
api_key = config.reranker_cohere_api_key
|
||||
if not api_key:
|
||||
raise ValueError(f"{ENV_COHERE_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'cohere'")
|
||||
model = os.environ.get(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL)
|
||||
base_url = os.environ.get(ENV_RERANKER_COHERE_BASE_URL) or None
|
||||
return CohereCrossEncoder(api_key=api_key, model=model, base_url=base_url)
|
||||
raise ValueError(f"{ENV_RERANKER_COHERE_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'cohere'")
|
||||
return CohereCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_cohere_model,
|
||||
base_url=config.reranker_cohere_base_url,
|
||||
)
|
||||
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)
|
||||
return FlashRankCrossEncoder(model_name=model, cache_dir=cache_dir)
|
||||
elif provider == "litellm":
|
||||
api_base = os.environ.get(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE)
|
||||
api_key = os.environ.get(ENV_LITELLM_API_KEY)
|
||||
model = os.environ.get(ENV_RERANKER_LITELLM_MODEL, DEFAULT_RERANKER_LITELLM_MODEL)
|
||||
return LiteLLMCrossEncoder(api_base=api_base, api_key=api_key, model=model)
|
||||
return LiteLLMCrossEncoder(
|
||||
api_base=config.reranker_litellm_api_base,
|
||||
api_key=config.reranker_litellm_api_key,
|
||||
model=config.reranker_litellm_model,
|
||||
)
|
||||
elif provider == "rrf":
|
||||
return RRFPassthroughCrossEncoder()
|
||||
else:
|
||||
|
||||
@@ -21,22 +21,19 @@ from ..config import (
|
||||
DEFAULT_EMBEDDINGS_LITELLM_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
DEFAULT_EMBEDDINGS_PROVIDER,
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
ENV_COHERE_API_KEY,
|
||||
ENV_EMBEDDINGS_COHERE_BASE_URL,
|
||||
ENV_EMBEDDINGS_COHERE_MODEL,
|
||||
ENV_EMBEDDINGS_LITELLM_MODEL,
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY,
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL,
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
ENV_EMBEDDINGS_OPENAI_API_KEY,
|
||||
ENV_EMBEDDINGS_OPENAI_BASE_URL,
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL,
|
||||
ENV_EMBEDDINGS_PROVIDER,
|
||||
ENV_EMBEDDINGS_TEI_URL,
|
||||
ENV_LITELLM_API_BASE,
|
||||
ENV_LITELLM_API_KEY,
|
||||
ENV_LLM_API_KEY,
|
||||
)
|
||||
|
||||
@@ -95,7 +92,7 @@ class LocalSTEmbeddings(Embeddings):
|
||||
The embedding dimension is auto-detected from the model.
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str | None = None, force_cpu: bool = False):
|
||||
def __init__(self, model_name: str | None = None, force_cpu: bool = False, trust_remote_code: bool = False):
|
||||
"""
|
||||
Initialize local SentenceTransformers embeddings.
|
||||
|
||||
@@ -104,9 +101,13 @@ class LocalSTEmbeddings(Embeddings):
|
||||
Default: BAAI/bge-small-en-v1.5
|
||||
force_cpu: Force CPU mode (avoids MPS/XPC issues on macOS in daemon mode).
|
||||
Default: False
|
||||
trust_remote_code: Allow loading models with custom code (security risk).
|
||||
Required for some models with custom architectures.
|
||||
Default: False (disabled for security)
|
||||
"""
|
||||
self.model_name = model_name or DEFAULT_EMBEDDINGS_LOCAL_MODEL
|
||||
self.force_cpu = force_cpu
|
||||
self.trust_remote_code = trust_remote_code
|
||||
self._model = None
|
||||
self._dimension: int | None = None
|
||||
|
||||
@@ -176,6 +177,7 @@ class LocalSTEmbeddings(Embeddings):
|
||||
self.model_name,
|
||||
device=device,
|
||||
model_kwargs={"low_cpu_mem_usage": False},
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
)
|
||||
finally:
|
||||
# Restore original logging level
|
||||
@@ -741,6 +743,7 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
return LocalSTEmbeddings(
|
||||
model_name=config.embeddings_local_model,
|
||||
force_cpu=config.embeddings_local_force_cpu,
|
||||
trust_remote_code=config.embeddings_local_trust_remote_code,
|
||||
)
|
||||
elif provider == "openai":
|
||||
# Use dedicated embeddings API key, or fall back to LLM API key
|
||||
@@ -754,17 +757,20 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
base_url = os.environ.get(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None
|
||||
return OpenAIEmbeddings(api_key=api_key, model=model, base_url=base_url)
|
||||
elif provider == "cohere":
|
||||
api_key = os.environ.get(ENV_COHERE_API_KEY)
|
||||
api_key = config.embeddings_cohere_api_key
|
||||
if not api_key:
|
||||
raise ValueError(f"{ENV_COHERE_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'cohere'")
|
||||
model = os.environ.get(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL)
|
||||
base_url = os.environ.get(ENV_EMBEDDINGS_COHERE_BASE_URL) or None
|
||||
return CohereEmbeddings(api_key=api_key, model=model, base_url=base_url)
|
||||
raise ValueError(f"{ENV_EMBEDDINGS_COHERE_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'cohere'")
|
||||
return CohereEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_cohere_model,
|
||||
base_url=config.embeddings_cohere_base_url,
|
||||
)
|
||||
elif provider == "litellm":
|
||||
api_base = os.environ.get(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE)
|
||||
api_key = os.environ.get(ENV_LITELLM_API_KEY)
|
||||
model = os.environ.get(ENV_EMBEDDINGS_LITELLM_MODEL, DEFAULT_EMBEDDINGS_LITELLM_MODEL)
|
||||
return LiteLLMEmbeddings(api_base=api_base, api_key=api_key, model=model)
|
||||
return LiteLLMEmbeddings(
|
||||
api_base=config.embeddings_litellm_api_base,
|
||||
api_key=config.embeddings_litellm_api_key,
|
||||
model=config.embeddings_litellm_model,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei', 'openai', 'cohere', 'litellm'"
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -84,7 +84,7 @@ class AnthropicLLM(LLMInterface):
|
||||
messages=test_messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=0.0,
|
||||
scope="test",
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
logger.info("Anthropic connection verified successfully")
|
||||
@@ -223,6 +223,24 @@ class AnthropicLLM(LLMInterface):
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
|
||||
|
||||
finish_reason = response.stop_reason if hasattr(response, "stop_reason") else None
|
||||
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,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0:
|
||||
logger.info(
|
||||
@@ -397,16 +415,41 @@ class AnthropicLLM(LLMInterface):
|
||||
|
||||
# Record metrics
|
||||
metrics = get_metrics_collector()
|
||||
duration = time.time() - start_time
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=time.time() - start_time,
|
||||
duration=duration,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record OpenTelemetry span
|
||||
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=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,
|
||||
|
||||
@@ -95,7 +95,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
messages=test_messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=0.0,
|
||||
scope="test",
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
logger.info("Claude Code connection verified successfully")
|
||||
@@ -237,6 +237,23 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
from hindsight_api.tracing import 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=result if isinstance(result, str) else json.dumps(result),
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
duration=duration,
|
||||
finish_reason=None,
|
||||
error=None,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0:
|
||||
logger.info(
|
||||
|
||||
@@ -136,6 +136,7 @@ class CodexLLM(LLMInterface):
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
scope="verification",
|
||||
)
|
||||
logger.info(f"Codex LLM verified: {self.model}")
|
||||
except Exception as e:
|
||||
@@ -261,6 +262,26 @@ class CodexLLM(LLMInterface):
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
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,
|
||||
)
|
||||
|
||||
if return_usage:
|
||||
# Codex doesn't provide token counts, estimate based on content
|
||||
estimated_input = sum(len(m.get("content", "")) for m in messages) // 4
|
||||
@@ -504,6 +525,28 @@ class CodexLLM(LLMInterface):
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record OpenTelemetry span
|
||||
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,
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
|
||||
@@ -136,6 +136,7 @@ class GeminiLLM(LLMInterface):
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
scope="verification",
|
||||
)
|
||||
logger.info(f"{self.provider.upper()} connection verified successfully")
|
||||
except Exception as e:
|
||||
@@ -275,6 +276,29 @@ class GeminiLLM(LLMInterface):
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
|
||||
finish_reason = None
|
||||
if hasattr(response, "candidates") and response.candidates:
|
||||
if hasattr(response.candidates[0], "finish_reason"):
|
||||
finish_reason = str(response.candidates[0].finish_reason)
|
||||
span_recorder = get_span_recorder()
|
||||
from hindsight_api.tracing import _serialize_for_span
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0 and input_tokens > 0:
|
||||
logger.info(
|
||||
@@ -466,6 +490,30 @@ class GeminiLLM(LLMInterface):
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record OpenTelemetry span
|
||||
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=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,
|
||||
|
||||
@@ -129,6 +129,23 @@ class MockLLM(LLMInterface):
|
||||
if self._mock_exception is not None:
|
||||
raise self._mock_exception
|
||||
|
||||
# Record trace span (minimal for mock provider)
|
||||
from hindsight_api.tracing import 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="mock response",
|
||||
input_tokens=10,
|
||||
output_tokens=5,
|
||||
duration=0.001, # Mock calls are instant
|
||||
finish_reason="stop",
|
||||
error=None,
|
||||
)
|
||||
|
||||
# Return mock response
|
||||
if self._mock_response is not None:
|
||||
result = self._mock_response
|
||||
@@ -192,20 +209,50 @@ class MockLLM(LLMInterface):
|
||||
if self._mock_exception is not None:
|
||||
raise self._mock_exception
|
||||
|
||||
# Record OpenTelemetry span
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
|
||||
span_recorder = get_span_recorder()
|
||||
|
||||
if self._mock_response is not None:
|
||||
if isinstance(self._mock_response, LLMToolCallResult):
|
||||
return self._mock_response
|
||||
# Allow setting just tool calls as a list
|
||||
if isinstance(self._mock_response, list):
|
||||
return LLMToolCallResult(
|
||||
result = self._mock_response
|
||||
elif isinstance(self._mock_response, list):
|
||||
# Allow setting just tool calls as a list
|
||||
result = LLMToolCallResult(
|
||||
tool_calls=[
|
||||
LLMToolCall(id=f"mock_{i}", name=tc["name"], arguments=tc.get("arguments", {}))
|
||||
for i, tc in enumerate(self._mock_response)
|
||||
],
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
else:
|
||||
result = LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
else:
|
||||
result = LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
|
||||
return LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
# Record span with mock values
|
||||
# Convert LLMToolCall objects to dicts for span recording
|
||||
tool_calls_dict = (
|
||||
[{"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in result.tool_calls]
|
||||
if result.tool_calls
|
||||
else None
|
||||
)
|
||||
span_recorder.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=result.content,
|
||||
input_tokens=10, # Mock value
|
||||
output_tokens=5, # Mock value
|
||||
duration=0.1, # Mock value
|
||||
finish_reason=result.finish_reason,
|
||||
error=None,
|
||||
tool_calls=tool_calls_dict,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (no-op for mock provider)."""
|
||||
|
||||
@@ -130,6 +130,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
scope="verification",
|
||||
)
|
||||
logger.info(f"Connection verified: {self.provider}/{self.model}")
|
||||
except Exception as e:
|
||||
@@ -368,6 +369,24 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
|
||||
|
||||
finish_reason = response.choices[0].finish_reason if response.choices else None
|
||||
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,
|
||||
)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0 and usage:
|
||||
ratio = max(1, output_tokens) / max(1, input_tokens)
|
||||
@@ -556,6 +575,30 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record OpenTelemetry span
|
||||
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=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,
|
||||
|
||||
@@ -402,7 +402,7 @@ async def run_reflect_agent(
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect_agent_final",
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
@@ -447,7 +447,7 @@ async def run_reflect_agent(
|
||||
result = await llm_config.call_with_tools(
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
scope="reflect_agent",
|
||||
scope="reflect_tool_call",
|
||||
tool_choice="required" if iteration == 0 else "auto", # Force tool use on first iteration
|
||||
)
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
@@ -479,7 +479,7 @@ async def run_reflect_agent(
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect_agent_final",
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
@@ -550,7 +550,7 @@ async def run_reflect_agent(
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect_agent_final",
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
@@ -617,23 +617,30 @@ async def run_reflect_agent(
|
||||
)
|
||||
continue
|
||||
|
||||
# Process done tool
|
||||
return await _process_done_tool(
|
||||
done_call,
|
||||
available_memory_ids,
|
||||
available_mental_model_ids,
|
||||
available_observation_ids,
|
||||
iteration + 1,
|
||||
total_tools_called,
|
||||
tool_trace,
|
||||
_get_llm_trace(),
|
||||
_get_usage(),
|
||||
_log_completion,
|
||||
reflect_id,
|
||||
directives_applied=directives_applied,
|
||||
llm_config=llm_config,
|
||||
response_schema=response_schema,
|
||||
)
|
||||
# Process done tool - wrap with tool call span
|
||||
from hindsight_api.tracing import get_tracer
|
||||
|
||||
tracer = get_tracer()
|
||||
span_name = "hindsight.reflect_tool_call"
|
||||
with tracer.start_as_current_span(span_name) as span:
|
||||
span.set_attribute("hindsight.scope", "reflect_tool_call")
|
||||
span.set_attribute("hindsight.operation", "reflect_tool_call")
|
||||
return await _process_done_tool(
|
||||
done_call,
|
||||
available_memory_ids,
|
||||
available_mental_model_ids,
|
||||
available_observation_ids,
|
||||
iteration + 1,
|
||||
total_tools_called,
|
||||
tool_trace,
|
||||
_get_llm_trace(),
|
||||
_get_usage(),
|
||||
_log_completion,
|
||||
reflect_id,
|
||||
directives_applied=directives_applied,
|
||||
llm_config=llm_config,
|
||||
response_schema=response_schema,
|
||||
)
|
||||
|
||||
# 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)]
|
||||
@@ -842,17 +849,67 @@ async def _execute_tool_with_timing(
|
||||
expand_fn: Callable[[list[str], str], Awaitable[dict[str, Any]]],
|
||||
) -> tuple[dict[str, Any], int]:
|
||||
"""Execute a tool call and return result with timing."""
|
||||
start = time.time()
|
||||
result = await _execute_tool(
|
||||
tc.name,
|
||||
tc.arguments,
|
||||
search_mental_models_fn,
|
||||
search_observations_fn,
|
||||
recall_fn,
|
||||
expand_fn,
|
||||
)
|
||||
duration_ms = int((time.time() - start) * 1000)
|
||||
return result, duration_ms
|
||||
from hindsight_api.tracing import get_tracer
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Create span for tool execution
|
||||
tracer = get_tracer()
|
||||
# Normalize tool name for span
|
||||
normalized_name = _normalize_tool_name(tc.name)
|
||||
span_name = f"hindsight.reflect_tool_exec.{normalized_name}"
|
||||
|
||||
# Calculate timestamps
|
||||
start_time_ns = time.time_ns()
|
||||
|
||||
with tracer.start_as_current_span(
|
||||
span_name,
|
||||
start_time=start_time_ns,
|
||||
end_on_exit=False,
|
||||
) as span:
|
||||
# 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))
|
||||
|
||||
try:
|
||||
result = await _execute_tool(
|
||||
tc.name,
|
||||
tc.arguments,
|
||||
search_mental_models_fn,
|
||||
search_observations_fn,
|
||||
recall_fn,
|
||||
expand_fn,
|
||||
)
|
||||
|
||||
# Set success attributes
|
||||
if isinstance(result, dict) and "error" in result:
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
span.set_status(Status(StatusCode.ERROR, result["error"]))
|
||||
else:
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
span.set_status(Status(StatusCode.OK))
|
||||
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
span.set_attribute("hindsight.tool.duration_ms", duration_ms)
|
||||
|
||||
# End span with correct timestamp
|
||||
end_time_ns = time.time_ns()
|
||||
span.end(end_time=end_time_ns)
|
||||
|
||||
return result, duration_ms
|
||||
except Exception as e:
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
span.set_status(Status(StatusCode.ERROR, str(e)))
|
||||
span.record_exception(e)
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
span.set_attribute("hindsight.tool.duration_ms", duration_ms)
|
||||
end_time_ns = time.time_ns()
|
||||
span.end(end_time=end_time_ns)
|
||||
raise
|
||||
|
||||
|
||||
async def _execute_tool(
|
||||
|
||||
@@ -802,7 +802,7 @@ Text:
|
||||
extraction_response_json, call_usage = await llm_config.call(
|
||||
messages=[{"role": "system", "content": prompt}, {"role": "user", "content": user_message}],
|
||||
response_format=response_schema,
|
||||
scope="memory_extract_facts",
|
||||
scope="retain_extract_facts",
|
||||
temperature=0.1,
|
||||
max_completion_tokens=config.retain_max_completion_tokens,
|
||||
max_retries=max_retries,
|
||||
|
||||
@@ -8,6 +8,16 @@ This extension enables multi-tenant memory isolation for applications using
|
||||
Supabase Auth - each authenticated user's memories are stored in a separate
|
||||
schema, ensuring complete data isolation.
|
||||
|
||||
Features:
|
||||
- Local JWT Verification: Validates tokens locally using JWKS public keys
|
||||
(no network call per request)
|
||||
- Automatic Schema Isolation: Each user gets {prefix}_{user_id} schema
|
||||
- Zero User Management: Leverages your existing Supabase Auth setup
|
||||
- Production Ready: Includes health checks, timeouts, key rotation handling,
|
||||
and error handling
|
||||
- Built-in: Ships with Hindsight, no extra installation needed
|
||||
- Legacy Support: Falls back to /auth/v1/user endpoint for HS256 projects
|
||||
|
||||
JWT Verification Strategy:
|
||||
By default, JWTs are verified locally using public keys from the Supabase
|
||||
JWKS endpoint (/auth/v1/.well-known/jwks.json). This is the Supabase-recommended
|
||||
|
||||
@@ -132,6 +132,10 @@ class RetainResult:
|
||||
unit_ids: list[list[str]] # List of unit IDs per content item
|
||||
success: bool = True
|
||||
error: str | None = None
|
||||
# Actual LLM token usage (populated by engine when available)
|
||||
llm_input_tokens: int | None = None
|
||||
llm_output_tokens: int | None = None
|
||||
llm_total_tokens: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -197,18 +197,30 @@ def main():
|
||||
embeddings_provider=config.embeddings_provider,
|
||||
embeddings_local_model=config.embeddings_local_model,
|
||||
embeddings_local_force_cpu=config.embeddings_local_force_cpu,
|
||||
embeddings_local_trust_remote_code=config.embeddings_local_trust_remote_code,
|
||||
embeddings_tei_url=config.embeddings_tei_url,
|
||||
embeddings_openai_base_url=config.embeddings_openai_base_url,
|
||||
embeddings_cohere_api_key=config.embeddings_cohere_api_key,
|
||||
embeddings_cohere_model=config.embeddings_cohere_model,
|
||||
embeddings_cohere_base_url=config.embeddings_cohere_base_url,
|
||||
embeddings_litellm_api_base=config.embeddings_litellm_api_base,
|
||||
embeddings_litellm_api_key=config.embeddings_litellm_api_key,
|
||||
embeddings_litellm_model=config.embeddings_litellm_model,
|
||||
reranker_provider=config.reranker_provider,
|
||||
reranker_local_model=config.reranker_local_model,
|
||||
reranker_local_force_cpu=config.reranker_local_force_cpu,
|
||||
reranker_local_max_concurrent=config.reranker_local_max_concurrent,
|
||||
reranker_local_trust_remote_code=config.reranker_local_trust_remote_code,
|
||||
reranker_tei_url=config.reranker_tei_url,
|
||||
reranker_tei_batch_size=config.reranker_tei_batch_size,
|
||||
reranker_tei_max_concurrent=config.reranker_tei_max_concurrent,
|
||||
reranker_max_candidates=config.reranker_max_candidates,
|
||||
reranker_cohere_api_key=config.reranker_cohere_api_key,
|
||||
reranker_cohere_model=config.reranker_cohere_model,
|
||||
reranker_cohere_base_url=config.reranker_cohere_base_url,
|
||||
reranker_litellm_api_base=config.reranker_litellm_api_base,
|
||||
reranker_litellm_api_key=config.reranker_litellm_api_key,
|
||||
reranker_litellm_model=config.reranker_litellm_model,
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
log_level=args.log_level,
|
||||
@@ -242,6 +254,11 @@ def main():
|
||||
worker_consolidation_max_slots=config.worker_consolidation_max_slots,
|
||||
reflect_max_iterations=config.reflect_max_iterations,
|
||||
mental_model_refresh_concurrency=config.mental_model_refresh_concurrency,
|
||||
otel_traces_enabled=config.otel_traces_enabled,
|
||||
otel_exporter_otlp_endpoint=config.otel_exporter_otlp_endpoint,
|
||||
otel_exporter_otlp_headers=config.otel_exporter_otlp_headers,
|
||||
otel_service_name=config.otel_service_name,
|
||||
otel_deployment_environment=config.otel_deployment_environment,
|
||||
)
|
||||
config.configure_logging()
|
||||
if not args.daemon:
|
||||
|
||||
@@ -35,6 +35,12 @@ class MCPToolsConfig:
|
||||
# How to resolve API key for tenant auth (optional)
|
||||
api_key_resolver: Callable[[], str | None] | None = None
|
||||
|
||||
# How to resolve tenant_id for usage metering (set by MCP middleware after auth)
|
||||
tenant_id_resolver: Callable[[], str | None] | None = None
|
||||
|
||||
# How to resolve api_key_id for usage metering (set by MCP middleware after auth)
|
||||
api_key_id_resolver: Callable[[], str | None] | None = None
|
||||
|
||||
# Whether to include bank_id as a parameter on tools (for multi-bank support)
|
||||
include_bank_id_param: bool = False
|
||||
|
||||
@@ -50,13 +56,15 @@ class MCPToolsConfig:
|
||||
|
||||
|
||||
def _get_request_context(config: MCPToolsConfig) -> RequestContext:
|
||||
"""Create RequestContext with API key from resolver if available.
|
||||
"""Create RequestContext with auth details from resolvers.
|
||||
|
||||
This enables tenant auth to work with MCP tools by propagating
|
||||
the Bearer token from the MCP middleware to the memory engine.
|
||||
This enables tenant auth and usage metering to work with MCP tools by propagating
|
||||
the authentication results from the MCP middleware to the memory engine.
|
||||
"""
|
||||
api_key = config.api_key_resolver() if config.api_key_resolver else None
|
||||
return RequestContext(api_key=api_key)
|
||||
tenant_id = config.tenant_id_resolver() if config.tenant_id_resolver else None
|
||||
api_key_id = config.api_key_id_resolver() if config.api_key_id_resolver else None
|
||||
return RequestContext(api_key=api_key, tenant_id=tenant_id, api_key_id=api_key_id)
|
||||
|
||||
|
||||
def parse_timestamp(timestamp: str) -> datetime | None:
|
||||
@@ -119,7 +127,19 @@ def register_mcp_tools(
|
||||
memory: MemoryEngine instance
|
||||
config: Tool configuration
|
||||
"""
|
||||
tools_to_register = config.tools or {"retain", "recall", "reflect", "list_banks", "create_bank"}
|
||||
tools_to_register = config.tools or {
|
||||
"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",
|
||||
}
|
||||
|
||||
if "retain" in tools_to_register:
|
||||
_register_retain(mcp, memory, config)
|
||||
@@ -136,6 +156,25 @@ def register_mcp_tools(
|
||||
if "create_bank" in tools_to_register:
|
||||
_register_create_bank(mcp, memory, config)
|
||||
|
||||
# Mental model tools
|
||||
if "list_mental_models" in tools_to_register:
|
||||
_register_list_mental_models(mcp, memory, config)
|
||||
|
||||
if "get_mental_model" in tools_to_register:
|
||||
_register_get_mental_model(mcp, memory, config)
|
||||
|
||||
if "create_mental_model" in tools_to_register:
|
||||
_register_create_mental_model(mcp, memory, config)
|
||||
|
||||
if "update_mental_model" in tools_to_register:
|
||||
_register_update_mental_model(mcp, memory, config)
|
||||
|
||||
if "delete_mental_model" in tools_to_register:
|
||||
_register_delete_mental_model(mcp, memory, config)
|
||||
|
||||
if "refresh_mental_model" in tools_to_register:
|
||||
_register_refresh_mental_model(mcp, memory, config)
|
||||
|
||||
|
||||
def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the retain tool."""
|
||||
@@ -511,3 +550,567 @@ def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating bank: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
|
||||
def _validate_mental_model_inputs(
|
||||
name: str | None = None, source_query: str | None = None, max_tokens: int | None = None
|
||||
) -> str | None:
|
||||
"""Validate mental model inputs, returning an error message or None if valid."""
|
||||
if name is not None and not name.strip():
|
||||
return "name cannot be empty"
|
||||
if source_query is not None and not source_query.strip():
|
||||
return "source_query cannot be empty"
|
||||
if max_tokens is not None and (max_tokens < 256 or max_tokens > 8192):
|
||||
return f"max_tokens must be between 256 and 8192, got {max_tokens}"
|
||||
return None
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# MENTAL MODEL TOOLS
|
||||
# =========================================================================
|
||||
|
||||
|
||||
def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the list_mental_models tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def list_mental_models(
|
||||
tags: list[str] | None = None,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
List mental models (pinned reflections) for a memory bank.
|
||||
|
||||
Mental models are living documents that stay current by periodically re-running
|
||||
a source query through reflect. Use them to maintain up-to-date summaries,
|
||||
preferences, or synthesized knowledge.
|
||||
|
||||
Args:
|
||||
tags: Optional tags to filter by (returns models matching any tag)
|
||||
bank_id: Optional bank to list from (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured", "items": []}'
|
||||
|
||||
models = await memory.list_mental_models(
|
||||
bank_id=target_bank,
|
||||
tags=tags,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps({"items": models}, indent=2, default=str)
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing mental models: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}", "items": []}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def list_mental_models(
|
||||
tags: list[str] | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
List mental models (pinned reflections) for this memory bank.
|
||||
|
||||
Mental models are living documents that stay current by periodically re-running
|
||||
a source query through reflect. Use them to maintain up-to-date summaries,
|
||||
preferences, or synthesized knowledge.
|
||||
|
||||
Args:
|
||||
tags: Optional tags to filter by (returns models matching any tag)
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured", "items": []}
|
||||
|
||||
models = await memory.list_mental_models(
|
||||
bank_id=target_bank,
|
||||
tags=tags,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return {"items": models}
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing mental models: {e}", exc_info=True)
|
||||
return {"error": str(e), "items": []}
|
||||
|
||||
|
||||
def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the get_mental_model tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def get_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Get a specific mental model by ID.
|
||||
|
||||
Returns the full mental model including its generated content, source query,
|
||||
and metadata. Use list_mental_models first to discover available model IDs.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to retrieve
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
model = await memory.get_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if model is None:
|
||||
return json.dumps({"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"})
|
||||
return json.dumps(model, indent=2, default=str)
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting mental model: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def get_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Get a specific mental model by ID.
|
||||
|
||||
Returns the full mental model including its generated content, source query,
|
||||
and metadata. Use list_mental_models first to discover available model IDs.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to retrieve
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
model = await memory.get_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if model is None:
|
||||
return {"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"}
|
||||
return model
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting mental model: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the create_mental_model tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def create_mental_model(
|
||||
name: str,
|
||||
source_query: str,
|
||||
mental_model_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
max_tokens: int = 2048,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Create a new mental model (pinned reflection).
|
||||
|
||||
A mental model is a living document generated by running the source_query through
|
||||
reflect. The content is auto-generated asynchronously - use the returned operation_id
|
||||
to track progress.
|
||||
|
||||
EXAMPLES:
|
||||
- name="Coding Preferences", source_query="What coding patterns and tools does the user prefer?"
|
||||
- name="Project Goals", source_query="What are the user's current project goals and priorities?"
|
||||
- name="Communication Style", source_query="How does the user prefer to communicate?"
|
||||
|
||||
Args:
|
||||
name: Human-readable name for the mental model
|
||||
source_query: The query to run through reflect to generate content
|
||||
mental_model_id: Optional custom ID (alphanumeric lowercase with hyphens). Auto-generated if not provided.
|
||||
tags: Optional tags for scoped visibility filtering
|
||||
max_tokens: Maximum tokens for generated content (256-8192, default: 2048)
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
validation_error = _validate_mental_model_inputs(
|
||||
name=name, source_query=source_query, max_tokens=max_tokens
|
||||
)
|
||||
if validation_error:
|
||||
return json.dumps({"error": validation_error})
|
||||
|
||||
request_context = _get_request_context(config)
|
||||
|
||||
# Create with placeholder content
|
||||
model = await memory.create_mental_model(
|
||||
bank_id=target_bank,
|
||||
name=name,
|
||||
source_query=source_query,
|
||||
content="Generating content...",
|
||||
mental_model_id=mental_model_id,
|
||||
tags=tags,
|
||||
max_tokens=max_tokens,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Schedule async refresh to generate actual content
|
||||
result = await memory.submit_async_refresh_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=model["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
return json.dumps(
|
||||
{
|
||||
"mental_model_id": model["id"],
|
||||
"operation_id": result["operation_id"],
|
||||
"status": "created",
|
||||
"message": f"Mental model '{name}' created. Content is being generated asynchronously.",
|
||||
}
|
||||
)
|
||||
except ValueError as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating mental model: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def create_mental_model(
|
||||
name: str,
|
||||
source_query: str,
|
||||
mental_model_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
max_tokens: int = 2048,
|
||||
) -> dict:
|
||||
"""
|
||||
Create a new mental model (pinned reflection).
|
||||
|
||||
A mental model is a living document generated by running the source_query through
|
||||
reflect. The content is auto-generated asynchronously - use the returned operation_id
|
||||
to track progress.
|
||||
|
||||
EXAMPLES:
|
||||
- name="Coding Preferences", source_query="What coding patterns and tools does the user prefer?"
|
||||
- name="Project Goals", source_query="What are the user's current project goals and priorities?"
|
||||
- name="Communication Style", source_query="How does the user prefer to communicate?"
|
||||
|
||||
Args:
|
||||
name: Human-readable name for the mental model
|
||||
source_query: The query to run through reflect to generate content
|
||||
mental_model_id: Optional custom ID (alphanumeric lowercase with hyphens). Auto-generated if not provided.
|
||||
tags: Optional tags for scoped visibility filtering
|
||||
max_tokens: Maximum tokens for generated content (256-8192, default: 2048)
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
validation_error = _validate_mental_model_inputs(
|
||||
name=name, source_query=source_query, max_tokens=max_tokens
|
||||
)
|
||||
if validation_error:
|
||||
return {"error": validation_error}
|
||||
|
||||
request_context = _get_request_context(config)
|
||||
|
||||
model = await memory.create_mental_model(
|
||||
bank_id=target_bank,
|
||||
name=name,
|
||||
source_query=source_query,
|
||||
content="Generating content...",
|
||||
mental_model_id=mental_model_id,
|
||||
tags=tags,
|
||||
max_tokens=max_tokens,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.submit_async_refresh_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=model["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
return {
|
||||
"mental_model_id": model["id"],
|
||||
"operation_id": result["operation_id"],
|
||||
"status": "created",
|
||||
"message": f"Mental model '{name}' created. Content is being generated asynchronously.",
|
||||
}
|
||||
except ValueError as e:
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating mental model: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_update_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the update_mental_model tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def update_mental_model(
|
||||
mental_model_id: str,
|
||||
name: str | None = None,
|
||||
source_query: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
tags: list[str] | None = None,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Update a mental model's metadata.
|
||||
|
||||
Changes the name, source query, or tags of an existing mental model.
|
||||
To regenerate the content, use refresh_mental_model after updating the source query.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to update
|
||||
name: New name (leave None to keep current)
|
||||
source_query: New source query (leave None to keep current)
|
||||
max_tokens: New max tokens for content generation (256-8192, leave None to keep current)
|
||||
tags: New tags (leave None to keep current)
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
validation_error = _validate_mental_model_inputs(
|
||||
name=name, source_query=source_query, max_tokens=max_tokens
|
||||
)
|
||||
if validation_error:
|
||||
return json.dumps({"error": validation_error})
|
||||
|
||||
model = await memory.update_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
name=name,
|
||||
source_query=source_query,
|
||||
max_tokens=max_tokens,
|
||||
tags=tags,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if model is None:
|
||||
return json.dumps({"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"})
|
||||
return json.dumps(model, indent=2, default=str)
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating mental model: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def update_mental_model(
|
||||
mental_model_id: str,
|
||||
name: str | None = None,
|
||||
source_query: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
tags: list[str] | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Update a mental model's metadata.
|
||||
|
||||
Changes the name, source query, or tags of an existing mental model.
|
||||
To regenerate the content, use refresh_mental_model after updating the source query.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to update
|
||||
name: New name (leave None to keep current)
|
||||
source_query: New source query (leave None to keep current)
|
||||
max_tokens: New max tokens for content generation (256-8192, leave None to keep current)
|
||||
tags: New tags (leave None to keep current)
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
validation_error = _validate_mental_model_inputs(
|
||||
name=name, source_query=source_query, max_tokens=max_tokens
|
||||
)
|
||||
if validation_error:
|
||||
return {"error": validation_error}
|
||||
|
||||
model = await memory.update_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
name=name,
|
||||
source_query=source_query,
|
||||
max_tokens=max_tokens,
|
||||
tags=tags,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if model is None:
|
||||
return {"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"}
|
||||
return model
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating mental model: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_delete_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the delete_mental_model tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def delete_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Delete a mental model.
|
||||
|
||||
Permanently removes a mental model and its generated content.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to delete
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
deleted = await memory.delete_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if not deleted:
|
||||
return json.dumps({"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"})
|
||||
return json.dumps({"status": "deleted", "mental_model_id": mental_model_id})
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting mental model: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def delete_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Delete a mental model.
|
||||
|
||||
Permanently removes a mental model and its generated content.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to delete
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
deleted = await memory.delete_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if not deleted:
|
||||
return {"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"}
|
||||
return {"status": "deleted", "mental_model_id": mental_model_id}
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting mental model: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the refresh_mental_model tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def refresh_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Refresh a mental model by re-running its source query.
|
||||
|
||||
Schedules an async task to re-run the source query through reflect and update the
|
||||
mental model's content with fresh results. Use this after adding new memories or
|
||||
when the mental model's content may be stale.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to refresh
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await memory.submit_async_refresh_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(
|
||||
{
|
||||
"operation_id": result["operation_id"],
|
||||
"status": "queued",
|
||||
"message": f"Refresh queued for mental model '{mental_model_id}'.",
|
||||
}
|
||||
)
|
||||
except ValueError as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error refreshing mental model: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def refresh_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Refresh a mental model by re-running its source query.
|
||||
|
||||
Schedules an async task to re-run the source query through reflect and update the
|
||||
mental model's content with fresh results. Use this after adding new memories or
|
||||
when the mental model's content may be stale.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to refresh
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
result = await memory.submit_async_refresh_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return {
|
||||
"operation_id": result["operation_id"],
|
||||
"status": "queued",
|
||||
"message": f"Refresh queued for mental model '{mental_model_id}'.",
|
||||
}
|
||||
except ValueError as e:
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error refreshing mental model: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
@@ -24,6 +24,7 @@ from alembic import command
|
||||
from alembic.config import Config
|
||||
from alembic.script.revision import ResolutionError
|
||||
from sqlalchemy import create_engine, text
|
||||
|
||||
from .utils import mask_network_location
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -20,7 +20,8 @@ class RequestContext:
|
||||
api_key: str | None = None
|
||||
api_key_id: str | None = None # UUID of the API key used for authentication
|
||||
tenant_id: str | None = None # Tenant identifier (set by extension after auth)
|
||||
internal: bool = False # True for background/internal operations (not user-visible)
|
||||
internal: bool = False # True for background/internal operations (skips extension auth)
|
||||
user_initiated: bool = False # True for async operations that originated from a user request
|
||||
|
||||
|
||||
from pgvector.sqlalchemy import Vector
|
||||
|
||||
@@ -0,0 +1,480 @@
|
||||
"""
|
||||
OpenTelemetry distributed tracing instrumentation for Hindsight API.
|
||||
|
||||
This module provides tracing for:
|
||||
- LLM API calls with full prompts/completions following GenAI semantic conventions
|
||||
- Token usage and model information
|
||||
- Error tracking and finish reasons
|
||||
|
||||
Tracing is conditional and disabled by default. When enabled, traces are exported
|
||||
to Langfuse (or any OTLP-compatible backend) via OTLP HTTP protocol.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Optional
|
||||
|
||||
from opentelemetry import trace
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _serialize_for_span(obj: Any) -> str:
|
||||
"""Serialize an object for span recording, handling Pydantic models."""
|
||||
if isinstance(obj, str):
|
||||
return obj
|
||||
if hasattr(obj, "model_dump_json"):
|
||||
# Pydantic v2 model
|
||||
return obj.model_dump_json()
|
||||
if hasattr(obj, "json"):
|
||||
# Pydantic v1 model
|
||||
return obj.json()
|
||||
if hasattr(obj, "model_dump"):
|
||||
# Pydantic v2 model - convert to dict then json
|
||||
return json.dumps(obj.model_dump())
|
||||
if hasattr(obj, "dict"):
|
||||
# Pydantic v1 model - convert to dict then json
|
||||
return json.dumps(obj.dict())
|
||||
# Fallback to json.dumps for dicts and other types
|
||||
return json.dumps(obj)
|
||||
|
||||
|
||||
# No-op tracer for when tracing is disabled
|
||||
class NoOpTracer:
|
||||
"""No-op tracer that provides the same interface as OpenTelemetry Tracer but does nothing."""
|
||||
|
||||
def start_as_current_span(self, name: str, **kwargs):
|
||||
"""Return a no-op context manager that yields a NoOpSpan."""
|
||||
from contextlib import contextmanager
|
||||
|
||||
@contextmanager
|
||||
def noop_span_context():
|
||||
yield NoOpSpan()
|
||||
|
||||
return noop_span_context()
|
||||
|
||||
def start_span(self, name: str, **kwargs):
|
||||
"""Return a no-op span."""
|
||||
return NoOpSpan()
|
||||
|
||||
|
||||
class NoOpSpan:
|
||||
"""No-op span that provides the same interface as OpenTelemetry Span but does nothing."""
|
||||
|
||||
def set_attribute(self, key: str, value: Any) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
def set_status(self, status: Any) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
def record_exception(self, exception: Exception) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
def add_event(self, name: str, attributes: dict | None = None) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
def end(self, end_time: int | None = None) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
|
||||
# Global tracer instance
|
||||
_tracer: trace.Tracer | NoOpTracer = NoOpTracer()
|
||||
_tracing_enabled: bool = False
|
||||
|
||||
|
||||
# GenAI semantic convention attribute names (based on v1.37 spec)
|
||||
class GenAIAttributes:
|
||||
"""GenAI semantic convention attribute names."""
|
||||
|
||||
# Operation and provider
|
||||
OPERATION_NAME = "gen_ai.operation.name"
|
||||
PROVIDER_NAME = "gen_ai.provider.name"
|
||||
|
||||
# Model information
|
||||
REQUEST_MODEL = "gen_ai.request.model"
|
||||
RESPONSE_MODEL = "gen_ai.response.model"
|
||||
|
||||
# Token usage
|
||||
USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens"
|
||||
USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
|
||||
|
||||
# Messages and prompts
|
||||
SYSTEM_INSTRUCTIONS = "gen_ai.system_instructions"
|
||||
INPUT_MESSAGES = "gen_ai.input.messages"
|
||||
OUTPUT_MESSAGES = "gen_ai.output.messages"
|
||||
|
||||
# Response metadata
|
||||
FINISH_REASONS = "gen_ai.response.finish_reasons"
|
||||
|
||||
# Error tracking
|
||||
ERROR_TYPE = "error.type"
|
||||
|
||||
|
||||
# Provider name mapping (Hindsight internal -> GenAI semantic convention)
|
||||
PROVIDER_NAME_MAPPING = {
|
||||
"openai": "openai",
|
||||
"anthropic": "anthropic",
|
||||
"gemini": "google",
|
||||
"vertexai": "google",
|
||||
"groq": "groq",
|
||||
"ollama": "ollama",
|
||||
"lmstudio": "lmstudio",
|
||||
"openai-codex": "openai",
|
||||
"claude-code": "anthropic",
|
||||
"mock": "mock",
|
||||
}
|
||||
|
||||
|
||||
def initialize_tracing(
|
||||
service_name: str,
|
||||
endpoint: str,
|
||||
headers: Optional[str] = None,
|
||||
deployment_environment: str = "development",
|
||||
) -> None:
|
||||
"""
|
||||
Initialize OpenTelemetry tracing with OTLP exporter.
|
||||
|
||||
Args:
|
||||
service_name: Name of the service for resource attributes
|
||||
endpoint: OTLP endpoint URL (e.g., https://cloud.langfuse.com/api/public/otel)
|
||||
headers: Optional headers in format "key1=value1,key2=value2"
|
||||
deployment_environment: Deployment environment (e.g., development, staging, production)
|
||||
"""
|
||||
global _tracer, _tracing_enabled
|
||||
|
||||
# Create resource with service information
|
||||
resource = Resource.create(
|
||||
{
|
||||
"service.name": service_name,
|
||||
"service.version": "0.4.8", # Could import from __version__
|
||||
"deployment.environment.name": deployment_environment,
|
||||
}
|
||||
)
|
||||
|
||||
# Parse headers
|
||||
headers_dict = {}
|
||||
if headers:
|
||||
for pair in headers.split(","):
|
||||
if "=" in pair:
|
||||
key, value = pair.split("=", 1)
|
||||
headers_dict[key.strip()] = value.strip()
|
||||
|
||||
# Create OTLP HTTP exporter
|
||||
# Note: Langfuse expects /v1/traces path appended to base endpoint
|
||||
otlp_endpoint = endpoint if endpoint.endswith("/v1/traces") else f"{endpoint}/v1/traces"
|
||||
otlp_exporter = OTLPSpanExporter(
|
||||
endpoint=otlp_endpoint,
|
||||
headers=headers_dict,
|
||||
)
|
||||
|
||||
# Create tracer provider with batch processor
|
||||
provider = TracerProvider(resource=resource)
|
||||
provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
|
||||
|
||||
# Set global tracer provider
|
||||
trace.set_tracer_provider(provider)
|
||||
|
||||
# Get tracer for this application
|
||||
_tracer = trace.get_tracer(__name__)
|
||||
_tracing_enabled = True
|
||||
|
||||
logger.info(f"Tracing initialized: endpoint={otlp_endpoint}, service={service_name}")
|
||||
|
||||
|
||||
def get_tracer() -> trace.Tracer | NoOpTracer:
|
||||
"""
|
||||
Get the global tracer instance.
|
||||
|
||||
Returns a no-op tracer if tracing is disabled, so callers don't need to check for None.
|
||||
This improves code readability by allowing direct use without null checks.
|
||||
"""
|
||||
return _tracer
|
||||
|
||||
|
||||
def create_operation_span(operation: str, bank_id: str | None = None):
|
||||
"""
|
||||
Create a parent span for a Hindsight operation (retain, reflect, consolidation, etc.).
|
||||
|
||||
This creates the span hierarchy:
|
||||
- hindsight.{operation} (parent)
|
||||
- chat {model} (child LLM calls)
|
||||
|
||||
Args:
|
||||
operation: Operation name (retain, reflect, consolidation, mental_model_refresh)
|
||||
bank_id: Optional bank ID for context
|
||||
|
||||
Returns:
|
||||
Span context manager
|
||||
"""
|
||||
if not _tracing_enabled or _tracer is None:
|
||||
# Return a no-op context manager
|
||||
from contextlib import nullcontext
|
||||
|
||||
return nullcontext()
|
||||
|
||||
span_name = f"hindsight.{operation}"
|
||||
span = _tracer.start_as_current_span(span_name)
|
||||
|
||||
# Add operation-specific attributes
|
||||
if span and hasattr(span, "set_attribute"):
|
||||
span.set_attribute("hindsight.operation", operation)
|
||||
if bank_id:
|
||||
span.set_attribute("hindsight.bank_id", bank_id)
|
||||
|
||||
return span
|
||||
|
||||
|
||||
def is_tracing_enabled() -> bool:
|
||||
"""Check if tracing is enabled."""
|
||||
return _tracing_enabled
|
||||
|
||||
|
||||
# Maximum content length before truncation (to stay within span size limits)
|
||||
MAX_CONTENT_LENGTH = 100_000 # characters
|
||||
|
||||
|
||||
def _truncate_content(content: str) -> str:
|
||||
"""Truncate content if too large for span."""
|
||||
if len(content) > MAX_CONTENT_LENGTH:
|
||||
return content[:MAX_CONTENT_LENGTH] + f"\n\n[TRUNCATED: {len(content) - MAX_CONTENT_LENGTH} chars omitted]"
|
||||
return content
|
||||
|
||||
|
||||
class LLMSpanRecorder:
|
||||
"""
|
||||
Records OpenTelemetry spans for LLM calls following GenAI semantic conventions.
|
||||
"""
|
||||
|
||||
def __init__(self, tracer: trace.Tracer):
|
||||
self.tracer = tracer
|
||||
|
||||
def record_llm_call(
|
||||
self,
|
||||
provider: str,
|
||||
model: str,
|
||||
scope: str,
|
||||
messages: list[dict[str, str]],
|
||||
response_content: Optional[str],
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
duration: float,
|
||||
finish_reason: Optional[str] = None,
|
||||
error: Optional[Exception] = None,
|
||||
tool_calls: Optional[list[dict[str, Any]]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Record a completed LLM call as a span with GenAI semantic conventions.
|
||||
|
||||
This creates a span AFTER the call completes, using timestamps to
|
||||
set the correct start/end times. This approach works better with
|
||||
the existing sync metrics recording pattern.
|
||||
|
||||
Args:
|
||||
provider: Hindsight provider name
|
||||
model: Model name
|
||||
scope: Scope identifier (memory, reflect, consolidation, etc.)
|
||||
messages: Input messages (chat history)
|
||||
response_content: Response text from LLM
|
||||
input_tokens: Input token count
|
||||
output_tokens: Output token count
|
||||
duration: Call duration in seconds
|
||||
finish_reason: Reason the model stopped (stop, length, tool_calls, etc.)
|
||||
error: Exception if call failed
|
||||
tool_calls: List of tool calls made (for function calling)
|
||||
"""
|
||||
try:
|
||||
# Map provider name to GenAI semantic convention
|
||||
genai_provider = PROVIDER_NAME_MAPPING.get(provider.lower(), provider.lower())
|
||||
|
||||
# Determine operation name based on scope/context
|
||||
operation_name = "chat" # Default for GenAI semantic conventions
|
||||
|
||||
# Create span name: "hindsight.{scope}" for consistency with parent spans
|
||||
# Model info is available in span attributes (gen_ai.request.model)
|
||||
if scope:
|
||||
span_name = f"hindsight.{scope}"
|
||||
else:
|
||||
# Fallback to chat {model} if no scope provided
|
||||
span_name = f"{operation_name} {model}"
|
||||
|
||||
# Calculate timestamps
|
||||
end_time_ns = time.time_ns()
|
||||
start_time_ns = end_time_ns - int(duration * 1_000_000_000)
|
||||
|
||||
# Create span with explicit timestamps
|
||||
with self.tracer.start_as_current_span(
|
||||
span_name,
|
||||
start_time=start_time_ns,
|
||||
end_on_exit=False, # We'll set end time manually
|
||||
) as span:
|
||||
# Set required attributes
|
||||
span.set_attribute(GenAIAttributes.OPERATION_NAME, operation_name)
|
||||
span.set_attribute(GenAIAttributes.PROVIDER_NAME, genai_provider)
|
||||
span.set_attribute(GenAIAttributes.REQUEST_MODEL, model)
|
||||
span.set_attribute(GenAIAttributes.RESPONSE_MODEL, model)
|
||||
span.set_attribute(GenAIAttributes.USAGE_INPUT_TOKENS, input_tokens)
|
||||
span.set_attribute(GenAIAttributes.USAGE_OUTPUT_TOKENS, output_tokens)
|
||||
|
||||
# Add custom attributes for Hindsight context
|
||||
span.set_attribute("hindsight.scope", scope)
|
||||
span.set_attribute("hindsight.provider.internal", provider)
|
||||
|
||||
# Add tool call information if present
|
||||
if tool_calls:
|
||||
span.set_attribute("gen_ai.tool_calls.count", len(tool_calls))
|
||||
# Add tool names as comma-separated list
|
||||
tool_names = [tc.get("name", "") for tc in tool_calls]
|
||||
span.set_attribute("gen_ai.tool_calls.names", ",".join(tool_names))
|
||||
|
||||
# Format messages for GenAI conventions (as JSON)
|
||||
input_messages_json = self._format_messages(messages)
|
||||
output_messages_json = self._format_output(response_content, finish_reason)
|
||||
|
||||
# Extract system instructions if present
|
||||
system_instructions = self._extract_system_instructions(messages)
|
||||
|
||||
# Add event with prompts/completions following v1.37 conventions
|
||||
event_attrs = {}
|
||||
if input_messages_json:
|
||||
event_attrs[GenAIAttributes.INPUT_MESSAGES] = input_messages_json
|
||||
if output_messages_json:
|
||||
event_attrs[GenAIAttributes.OUTPUT_MESSAGES] = output_messages_json
|
||||
if system_instructions:
|
||||
event_attrs[GenAIAttributes.SYSTEM_INSTRUCTIONS] = system_instructions
|
||||
if finish_reason:
|
||||
event_attrs[GenAIAttributes.FINISH_REASONS] = json.dumps([finish_reason])
|
||||
|
||||
span.add_event(
|
||||
"gen_ai.client.inference.operation.details",
|
||||
attributes=event_attrs,
|
||||
)
|
||||
|
||||
# Add individual tool call events with details
|
||||
if tool_calls:
|
||||
for i, tc in enumerate(tool_calls):
|
||||
tool_event_attrs = {
|
||||
"tool.name": tc.get("name", ""),
|
||||
"tool.id": tc.get("id", ""),
|
||||
"tool.arguments": json.dumps(tc.get("arguments", {})),
|
||||
}
|
||||
span.add_event(f"gen_ai.tool_call.{i}", attributes=tool_event_attrs)
|
||||
|
||||
# Handle errors
|
||||
if error:
|
||||
span.set_status(Status(StatusCode.ERROR, str(error)))
|
||||
span.set_attribute(GenAIAttributes.ERROR_TYPE, type(error).__name__)
|
||||
span.record_exception(error)
|
||||
else:
|
||||
span.set_status(Status(StatusCode.OK))
|
||||
|
||||
# Set end time
|
||||
span.end(end_time=end_time_ns)
|
||||
|
||||
except Exception as e:
|
||||
# Don't let tracing errors break LLM calls
|
||||
logger.error(f"Failed to record LLM span: {e}", exc_info=True)
|
||||
|
||||
def _format_messages(self, messages: list[dict[str, str]]) -> str:
|
||||
"""
|
||||
Format messages into GenAI semantic convention format (JSON array).
|
||||
|
||||
Returns JSON string representation of message array.
|
||||
"""
|
||||
try:
|
||||
formatted = []
|
||||
for msg in messages:
|
||||
content = msg.get("content", "")
|
||||
# Truncate if needed
|
||||
if isinstance(content, str):
|
||||
content = _truncate_content(content)
|
||||
|
||||
formatted.append(
|
||||
{
|
||||
"role": msg.get("role", "user"),
|
||||
"content": content,
|
||||
}
|
||||
)
|
||||
|
||||
return json.dumps(formatted)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to format input messages: {e}")
|
||||
return "[]"
|
||||
|
||||
def _format_output(
|
||||
self,
|
||||
content: Optional[str],
|
||||
finish_reason: Optional[str],
|
||||
) -> str:
|
||||
"""Format output message into GenAI semantic convention format."""
|
||||
try:
|
||||
if content is None:
|
||||
return "[]"
|
||||
|
||||
# Truncate if needed
|
||||
if isinstance(content, str):
|
||||
content = _truncate_content(content)
|
||||
|
||||
return json.dumps(
|
||||
[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
}
|
||||
]
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to format output message: {e}")
|
||||
return "[]"
|
||||
|
||||
def _extract_system_instructions(self, messages: list[dict[str, str]]) -> Optional[str]:
|
||||
"""Extract system instructions from messages if present."""
|
||||
try:
|
||||
for msg in messages:
|
||||
if msg.get("role") == "system":
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str):
|
||||
return _truncate_content(content)
|
||||
return str(content)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to extract system instructions: {e}")
|
||||
return None
|
||||
|
||||
|
||||
class NoOpLLMSpanRecorder:
|
||||
"""No-op span recorder for when tracing is disabled."""
|
||||
|
||||
def record_llm_call(self, **kwargs) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
|
||||
# Global span recorder instance
|
||||
_span_recorder: Optional[LLMSpanRecorder] = None
|
||||
|
||||
|
||||
def get_span_recorder() -> LLMSpanRecorder | NoOpLLMSpanRecorder:
|
||||
"""Get the global span recorder (NoOp if tracing disabled)."""
|
||||
if _span_recorder is None:
|
||||
return NoOpLLMSpanRecorder()
|
||||
return _span_recorder
|
||||
|
||||
|
||||
def create_span_recorder() -> LLMSpanRecorder:
|
||||
"""Create and set the global span recorder."""
|
||||
global _span_recorder
|
||||
tracer = get_tracer()
|
||||
if tracer is None:
|
||||
raise RuntimeError("Tracing not initialized. Call initialize_tracing() first.")
|
||||
_span_recorder = LLMSpanRecorder(tracer)
|
||||
return _span_recorder
|
||||
@@ -1,5 +1,6 @@
|
||||
from urllib.parse import urlparse, urlunparse
|
||||
|
||||
|
||||
def mask_network_location(url):
|
||||
if not url:
|
||||
return url
|
||||
@@ -9,4 +10,4 @@ def mask_network_location(url):
|
||||
masked_network_location += f":{parsed_url.port}"
|
||||
if parsed_url.username or parsed_url.password:
|
||||
masked_network_location = f"***:***@{masked_network_location}"
|
||||
return urlunparse(parsed_url._replace(netloc=masked_network_location))
|
||||
return urlunparse(parsed_url._replace(netloc=masked_network_location))
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api"
|
||||
version = "0.4.9"
|
||||
version = "0.4.10"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -33,6 +33,8 @@ dependencies = [
|
||||
"opentelemetry-sdk>=1.20.0",
|
||||
"opentelemetry-instrumentation-fastapi>=0.41b0",
|
||||
"opentelemetry-exporter-prometheus>=0.41b0",
|
||||
"opentelemetry-exporter-otlp-proto-http>=1.20.0",
|
||||
"opentelemetry-semantic-conventions>=0.41b0",
|
||||
"dateparser>=1.2.2",
|
||||
"google-genai>=1.0.0",
|
||||
"google-auth>=2.0.0",
|
||||
|
||||
@@ -353,6 +353,14 @@ class TestOperationHooksParameters:
|
||||
assert post_result.error is None
|
||||
assert post_result.unit_ids == result # Should match the return value
|
||||
|
||||
# Verify actual LLM token usage is populated
|
||||
assert post_result.llm_input_tokens is not None
|
||||
assert post_result.llm_input_tokens > 0
|
||||
assert post_result.llm_output_tokens is not None
|
||||
assert post_result.llm_output_tokens > 0
|
||||
assert post_result.llm_total_tokens is not None
|
||||
assert post_result.llm_total_tokens == post_result.llm_input_tokens + post_result.llm_output_tokens
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
|
||||
"""Pre-recall hook receives all user-provided parameters."""
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Integration test for MCP endpoint routing.
|
||||
|
||||
This test verifies that /mcp/ and /mcp/{bank_id}/ expose different tool sets.
|
||||
This test verifies that /mcp/ and /mcp/{bank_id}/ expose different tool sets,
|
||||
and that URLs with or without trailing slashes both work (no 307 redirect).
|
||||
"""
|
||||
|
||||
import httpx
|
||||
@@ -39,12 +40,18 @@ async def test_mcp_endpoint_routing_integration(memory):
|
||||
|
||||
multi_tools = {t.name for t in multi_result.tools}
|
||||
|
||||
# Multi-bank should have all tools including bank management
|
||||
# Multi-bank should have all tools including bank management and mental models
|
||||
assert "retain" in multi_tools
|
||||
assert "recall" in multi_tools
|
||||
assert "reflect" in multi_tools
|
||||
assert "list_banks" in multi_tools, "Multi-bank should expose list_banks"
|
||||
assert "create_bank" in multi_tools, "Multi-bank should expose create_bank"
|
||||
assert "list_mental_models" in multi_tools, "Multi-bank should expose list_mental_models"
|
||||
assert "create_mental_model" in multi_tools, "Multi-bank should expose create_mental_model"
|
||||
assert "get_mental_model" in multi_tools, "Multi-bank should expose get_mental_model"
|
||||
assert "update_mental_model" in multi_tools, "Multi-bank should expose update_mental_model"
|
||||
assert "delete_mental_model" in multi_tools, "Multi-bank should expose delete_mental_model"
|
||||
assert "refresh_mental_model" in multi_tools, "Multi-bank should expose refresh_mental_model"
|
||||
|
||||
# Multi-bank retain should have bank_id parameter
|
||||
retain_tool = next((t for t in multi_result.tools if t.name == "retain"), None)
|
||||
@@ -64,10 +71,12 @@ async def test_mcp_endpoint_routing_integration(memory):
|
||||
|
||||
single_tools = {t.name for t in single_result.tools}
|
||||
|
||||
# Single-bank should only have scoped tools (no bank management)
|
||||
# Single-bank should have scoped tools including mental models (no bank management)
|
||||
assert "retain" in single_tools
|
||||
assert "recall" in single_tools
|
||||
assert "reflect" in single_tools
|
||||
assert "list_mental_models" in single_tools, "Single-bank should expose list_mental_models"
|
||||
assert "create_mental_model" in single_tools, "Single-bank should expose create_mental_model"
|
||||
assert "list_banks" not in single_tools, "Single-bank should NOT expose list_banks"
|
||||
assert "create_bank" not in single_tools, "Single-bank should NOT expose create_bank"
|
||||
|
||||
@@ -76,3 +85,196 @@ async def test_mcp_endpoint_routing_integration(memory):
|
||||
assert retain_tool is not None
|
||||
single_params = set(retain_tool.inputSchema.get("properties", {}).keys())
|
||||
assert "bank_id" not in single_params, "Single-bank retain should NOT have bank_id parameter"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_no_trailing_slash_works(memory):
|
||||
"""Test that /mcp (no trailing slash) discovers tools without 307 redirect.
|
||||
|
||||
Starlette's Mount class redirects /mcp to /mcp/ with a 307 Temporary Redirect.
|
||||
Many MCP clients don't follow POST redirects, causing 0 tools to be discovered.
|
||||
MCPMiddleware wraps the app directly (no Mount), so the redirect never happens.
|
||||
"""
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
app = create_app(memory, mcp_api_enabled=True, initialize_memory=False)
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
from httpx import ASGITransport
|
||||
|
||||
async with httpx.AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as http_client:
|
||||
# /mcp (no slash) should work the same as /mcp/
|
||||
async with streamable_http_client("http://test/mcp", http_client=http_client) as (
|
||||
read_stream,
|
||||
write_stream,
|
||||
_,
|
||||
):
|
||||
async with ClientSession(read_stream, write_stream) as session:
|
||||
await session.initialize()
|
||||
result = await session.list_tools()
|
||||
|
||||
tools = {t.name for t in result.tools}
|
||||
assert len(tools) >= 11, f"Expected at least 11 tools from /mcp, got {len(tools)}: {tools}"
|
||||
assert "retain" in tools
|
||||
assert "recall" in tools
|
||||
assert "list_banks" in tools
|
||||
|
||||
# /mcp/my-bank (single-bank, no slash) should also work
|
||||
async with streamable_http_client("http://test/mcp/my-bank", http_client=http_client) as (
|
||||
read_stream,
|
||||
write_stream,
|
||||
_,
|
||||
):
|
||||
async with ClientSession(read_stream, write_stream) as session:
|
||||
await session.initialize()
|
||||
result = await session.list_tools()
|
||||
|
||||
tools = {t.name for t in result.tools}
|
||||
assert "retain" in tools
|
||||
assert "list_banks" not in tools, "Single-bank /mcp/my-bank should NOT expose list_banks"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tool_execution_through_client(memory):
|
||||
"""Test that tools can be called (not just discovered) through the MCP client.
|
||||
|
||||
This verifies the full pipeline: HTTP → middleware → FastMCP → tool → engine → response.
|
||||
Previous tests only checked tool discovery (list_tools), not actual execution.
|
||||
"""
|
||||
from httpx import ASGITransport
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
app = create_app(memory, mcp_api_enabled=True, initialize_memory=False)
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
async with httpx.AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as http_client:
|
||||
async with streamable_http_client("http://test/mcp/", http_client=http_client) as (
|
||||
read_stream,
|
||||
write_stream,
|
||||
_,
|
||||
):
|
||||
async with ClientSession(read_stream, write_stream) as session:
|
||||
await session.initialize()
|
||||
|
||||
# Execute list_banks tool
|
||||
result = await session.call_tool("list_banks", arguments={})
|
||||
assert result is not None
|
||||
assert len(result.content) > 0
|
||||
# The result text should be valid JSON with a "banks" key
|
||||
import json
|
||||
|
||||
response_text = result.content[0].text
|
||||
parsed = json.loads(response_text)
|
||||
assert "banks" in parsed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_mental_model_validation_through_client(memory):
|
||||
"""Test that input validation works through the real MCP transport.
|
||||
|
||||
Verifies that invalid inputs return error messages without crashing,
|
||||
and that the engine is never called with invalid data.
|
||||
"""
|
||||
from httpx import ASGITransport
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
app = create_app(memory, mcp_api_enabled=True, initialize_memory=False)
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
async with httpx.AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as http_client:
|
||||
async with streamable_http_client("http://test/mcp/", http_client=http_client) as (
|
||||
read_stream,
|
||||
write_stream,
|
||||
_,
|
||||
):
|
||||
async with ClientSession(read_stream, write_stream) as session:
|
||||
await session.initialize()
|
||||
|
||||
# Test: empty name should return validation error
|
||||
import json
|
||||
|
||||
result = await session.call_tool(
|
||||
"create_mental_model",
|
||||
arguments={"name": "", "source_query": "test query"},
|
||||
)
|
||||
assert result is not None
|
||||
parsed = json.loads(result.content[0].text)
|
||||
assert "error" in parsed
|
||||
assert "name cannot be empty" in parsed["error"]
|
||||
|
||||
# Test: max_tokens out of range should return validation error
|
||||
result = await session.call_tool(
|
||||
"create_mental_model",
|
||||
arguments={"name": "Test", "source_query": "test query", "max_tokens": 0},
|
||||
)
|
||||
parsed = json.loads(result.content[0].text)
|
||||
assert "error" in parsed
|
||||
assert "max_tokens must be between 256 and 8192" in parsed["error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_bank_named_sse_routes_to_single_bank(memory):
|
||||
"""Test that a bank named 'sse' routes to single-bank mode.
|
||||
|
||||
Regression test: the old MCP_ENDPOINTS blocklist prevented banks named 'sse'
|
||||
or 'messages' from being accessed via path routing. They fell through to
|
||||
multi-bank mode instead.
|
||||
"""
|
||||
from httpx import ASGITransport
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
app = create_app(memory, mcp_api_enabled=True, initialize_memory=False)
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
async with httpx.AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as http_client:
|
||||
async with streamable_http_client("http://test/mcp/sse/", http_client=http_client) as (
|
||||
read_stream,
|
||||
write_stream,
|
||||
_,
|
||||
):
|
||||
async with ClientSession(read_stream, write_stream) as session:
|
||||
await session.initialize()
|
||||
result = await session.list_tools()
|
||||
tools = {t.name for t in result.tools}
|
||||
|
||||
# Should be single-bank mode (no bank management tools)
|
||||
assert "retain" in tools
|
||||
assert "recall" in tools
|
||||
assert "list_banks" not in tools, "Bank 'sse' should route to single-bank mode"
|
||||
assert "create_bank" not in tools
|
||||
|
||||
# retain should NOT have bank_id parameter (single-bank mode)
|
||||
retain_tool = next(t for t in result.tools if t.name == "retain")
|
||||
params = set(retain_tool.inputSchema.get("properties", {}).keys())
|
||||
assert "bank_id" not in params
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_bank_named_messages_routes_to_single_bank(memory):
|
||||
"""Test that a bank named 'messages' routes to single-bank mode.
|
||||
|
||||
Same regression test as test_mcp_bank_named_sse_routes_to_single_bank but for 'messages'.
|
||||
"""
|
||||
from httpx import ASGITransport
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
app = create_app(memory, mcp_api_enabled=True, initialize_memory=False)
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
async with httpx.AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as http_client:
|
||||
async with streamable_http_client("http://test/mcp/messages/", http_client=http_client) as (
|
||||
read_stream,
|
||||
write_stream,
|
||||
_,
|
||||
):
|
||||
async with ClientSession(read_stream, write_stream) as session:
|
||||
await session.initialize()
|
||||
result = await session.list_tools()
|
||||
tools = {t.name for t in result.tools}
|
||||
|
||||
assert "retain" in tools
|
||||
assert "list_banks" not in tools, "Bank 'messages' should route to single-bank mode"
|
||||
|
||||
@@ -165,5 +165,5 @@ class TestMCPExtensionIntegration:
|
||||
assert "create_bank" in tools
|
||||
# Extension tool also present
|
||||
assert "test_extension_tool" in tools
|
||||
# Total: 5 core + 1 extension = 6 tools
|
||||
assert len(tools) == 6
|
||||
# At least 11 core + 1 extension = 12 tools (may grow as new tools are added)
|
||||
assert len(tools) >= 12
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
"""Test MCP server routing with dynamic bank_id."""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_memory():
|
||||
@@ -17,7 +18,7 @@ def mock_memory():
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_context_variable():
|
||||
"""Test that context variable works correctly."""
|
||||
from hindsight_api.api.mcp import get_current_bank_id, _current_bank_id
|
||||
from hindsight_api.api.mcp import _current_bank_id, get_current_bank_id
|
||||
|
||||
# Initially None
|
||||
assert get_current_bank_id() is None
|
||||
@@ -36,7 +37,7 @@ async def test_mcp_context_variable():
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tools_use_context_bank_id(mock_memory):
|
||||
"""Test that MCP tools use bank_id from context."""
|
||||
from hindsight_api.api.mcp import create_mcp_server, _current_bank_id
|
||||
from hindsight_api.api.mcp import _current_bank_id, create_mcp_server
|
||||
|
||||
mcp_server = create_mcp_server(mock_memory)
|
||||
|
||||
@@ -62,6 +63,7 @@ async def test_mcp_tools_use_context_bank_id(mock_memory):
|
||||
|
||||
def test_path_parsing_logic():
|
||||
"""Test the path parsing logic for bank_id extraction."""
|
||||
|
||||
def parse_path(path):
|
||||
"""Simulate the path parsing logic from MCPMiddleware."""
|
||||
if not path.startswith("/") or len(path) <= 1:
|
||||
@@ -102,7 +104,7 @@ def test_path_parsing_logic():
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_key_context_variable():
|
||||
"""Test that API key context variable works correctly."""
|
||||
from hindsight_api.api.mcp import get_current_api_key, _current_api_key
|
||||
from hindsight_api.api.mcp import _current_api_key, get_current_api_key
|
||||
|
||||
# Initially None
|
||||
assert get_current_api_key() is None
|
||||
@@ -121,7 +123,7 @@ async def test_api_key_context_variable():
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tools_propagate_api_key(mock_memory):
|
||||
"""Test that MCP tools propagate API key to RequestContext."""
|
||||
from hindsight_api.api.mcp import create_mcp_server, _current_bank_id, _current_api_key
|
||||
from hindsight_api.api.mcp import _current_api_key, _current_bank_id, create_mcp_server
|
||||
|
||||
mcp_server = create_mcp_server(mock_memory)
|
||||
tools = mcp_server._tool_manager._tools
|
||||
@@ -143,21 +145,99 @@ async def test_mcp_tools_propagate_api_key(mock_memory):
|
||||
_current_api_key.reset(api_key_token)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tenant_id_context_variable():
|
||||
"""Test that tenant_id and api_key_id context variables work correctly."""
|
||||
from hindsight_api.api.mcp import (
|
||||
_current_api_key_id,
|
||||
_current_tenant_id,
|
||||
get_current_api_key_id,
|
||||
get_current_tenant_id,
|
||||
)
|
||||
|
||||
# Initially None
|
||||
assert get_current_tenant_id() is None
|
||||
assert get_current_api_key_id() is None
|
||||
|
||||
# Set and verify
|
||||
tenant_token = _current_tenant_id.set("org-123")
|
||||
key_id_token = _current_api_key_id.set("key-456")
|
||||
try:
|
||||
assert get_current_tenant_id() == "org-123"
|
||||
assert get_current_api_key_id() == "key-456"
|
||||
finally:
|
||||
_current_tenant_id.reset(tenant_token)
|
||||
_current_api_key_id.reset(key_id_token)
|
||||
|
||||
# Back to None after reset
|
||||
assert get_current_tenant_id() is None
|
||||
assert get_current_api_key_id() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tools_propagate_tenant_id_and_api_key_id(mock_memory):
|
||||
"""Test that MCP tools propagate tenant_id and api_key_id to RequestContext.
|
||||
|
||||
This is the critical test for usage metering: the UsageMeteringValidator reads
|
||||
request_context.tenant_id to identify the org for billing. Without this,
|
||||
MCP operations get tenant_id="unknown" and billing is skipped entirely.
|
||||
"""
|
||||
from hindsight_api.api.mcp import (
|
||||
_current_api_key,
|
||||
_current_api_key_id,
|
||||
_current_bank_id,
|
||||
_current_tenant_id,
|
||||
create_mcp_server,
|
||||
)
|
||||
|
||||
mcp_server = create_mcp_server(mock_memory)
|
||||
tools = mcp_server._tool_manager._tools
|
||||
|
||||
# Set all context vars (simulating what MCPMiddleware does after authenticate_mcp)
|
||||
bank_token = _current_bank_id.set("test-bank")
|
||||
api_key_token = _current_api_key.set("hsk_test_key")
|
||||
tenant_token = _current_tenant_id.set("org-billing-123")
|
||||
key_id_token = _current_api_key_id.set("key-uuid-456")
|
||||
try:
|
||||
retain_tool = tools["retain"]
|
||||
await retain_tool.fn(content="test content", context="test_context", async_processing=False)
|
||||
|
||||
# Verify the RequestContext passed to memory engine has all auth fields
|
||||
mock_memory.retain_batch_async.assert_called_once()
|
||||
request_context = mock_memory.retain_batch_async.call_args.kwargs["request_context"]
|
||||
assert request_context.api_key == "hsk_test_key"
|
||||
assert request_context.tenant_id == "org-billing-123"
|
||||
assert request_context.api_key_id == "key-uuid-456"
|
||||
finally:
|
||||
_current_bank_id.reset(bank_token)
|
||||
_current_api_key.reset(api_key_token)
|
||||
_current_tenant_id.reset(tenant_token)
|
||||
_current_api_key_id.reset(key_id_token)
|
||||
|
||||
|
||||
def test_multi_bank_mode_exposes_all_tools(mock_memory):
|
||||
"""Test that multi-bank mode exposes all tools including bank management."""
|
||||
"""Test that multi-bank mode exposes all tools including bank management and mental models."""
|
||||
from hindsight_api.api.mcp import create_mcp_server
|
||||
|
||||
# Create server in multi-bank mode (default)
|
||||
mcp_server = create_mcp_server(mock_memory, multi_bank=True)
|
||||
tools = mcp_server._tool_manager._tools
|
||||
|
||||
# Should have all tools
|
||||
# Core tools
|
||||
assert "retain" in tools
|
||||
assert "recall" in tools
|
||||
assert "reflect" in tools
|
||||
assert "list_banks" in tools
|
||||
assert "create_bank" in tools
|
||||
|
||||
# Mental model tools
|
||||
assert "list_mental_models" in tools
|
||||
assert "get_mental_model" in tools
|
||||
assert "create_mental_model" in tools
|
||||
assert "update_mental_model" in tools
|
||||
assert "delete_mental_model" in tools
|
||||
assert "refresh_mental_model" in tools
|
||||
|
||||
|
||||
def test_single_bank_mode_excludes_bank_management_tools(mock_memory):
|
||||
"""Test that single-bank mode only exposes bank-scoped tools."""
|
||||
@@ -167,11 +247,19 @@ def test_single_bank_mode_excludes_bank_management_tools(mock_memory):
|
||||
mcp_server = create_mcp_server(mock_memory, multi_bank=False)
|
||||
tools = mcp_server._tool_manager._tools
|
||||
|
||||
# Should only have bank-scoped tools
|
||||
# Should have bank-scoped tools
|
||||
assert "retain" in tools
|
||||
assert "recall" in tools
|
||||
assert "reflect" in tools
|
||||
|
||||
# Mental model tools should also be present (they're bank-scoped)
|
||||
assert "list_mental_models" in tools
|
||||
assert "get_mental_model" in tools
|
||||
assert "create_mental_model" in tools
|
||||
assert "update_mental_model" in tools
|
||||
assert "delete_mental_model" in tools
|
||||
assert "refresh_mental_model" in tools
|
||||
|
||||
# Should NOT have bank management tools
|
||||
assert "list_banks" not in tools
|
||||
assert "create_bank" not in tools
|
||||
@@ -179,46 +267,56 @@ def test_single_bank_mode_excludes_bank_management_tools(mock_memory):
|
||||
|
||||
def test_multi_bank_mode_tools_have_bank_id_param(mock_memory):
|
||||
"""Test that multi-bank mode tools include bank_id parameter."""
|
||||
from hindsight_api.api.mcp import create_mcp_server
|
||||
import inspect
|
||||
|
||||
from hindsight_api.api.mcp import create_mcp_server
|
||||
|
||||
mcp_server = create_mcp_server(mock_memory, multi_bank=True)
|
||||
tools = mcp_server._tool_manager._tools
|
||||
|
||||
# Check that tools have bank_id parameter
|
||||
retain_tool = tools["retain"]
|
||||
retain_sig = inspect.signature(retain_tool.fn)
|
||||
assert "bank_id" in retain_sig.parameters
|
||||
|
||||
recall_tool = tools["recall"]
|
||||
recall_sig = inspect.signature(recall_tool.fn)
|
||||
assert "bank_id" in recall_sig.parameters
|
||||
|
||||
reflect_tool = tools["reflect"]
|
||||
reflect_sig = inspect.signature(reflect_tool.fn)
|
||||
assert "bank_id" in reflect_sig.parameters
|
||||
# All bank-scoped tools should have bank_id parameter in multi-bank mode
|
||||
bank_scoped_tools = [
|
||||
"retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"create_mental_model",
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
]
|
||||
for tool_name in bank_scoped_tools:
|
||||
tool = tools[tool_name]
|
||||
sig = inspect.signature(tool.fn)
|
||||
assert "bank_id" in sig.parameters, f"{tool_name} should have bank_id param in multi-bank mode"
|
||||
|
||||
|
||||
def test_single_bank_mode_tools_no_bank_id_param(mock_memory):
|
||||
"""Test that single-bank mode tools do NOT include bank_id parameter."""
|
||||
from hindsight_api.api.mcp import create_mcp_server
|
||||
import inspect
|
||||
|
||||
from hindsight_api.api.mcp import create_mcp_server
|
||||
|
||||
mcp_server = create_mcp_server(mock_memory, multi_bank=False)
|
||||
tools = mcp_server._tool_manager._tools
|
||||
|
||||
# Check that tools do NOT have bank_id parameter
|
||||
retain_tool = tools["retain"]
|
||||
retain_sig = inspect.signature(retain_tool.fn)
|
||||
assert "bank_id" not in retain_sig.parameters
|
||||
|
||||
recall_tool = tools["recall"]
|
||||
recall_sig = inspect.signature(recall_tool.fn)
|
||||
assert "bank_id" not in recall_sig.parameters
|
||||
|
||||
reflect_tool = tools["reflect"]
|
||||
reflect_sig = inspect.signature(reflect_tool.fn)
|
||||
assert "bank_id" not in reflect_sig.parameters
|
||||
# No bank-scoped tool should have bank_id parameter in single-bank mode
|
||||
bank_scoped_tools = [
|
||||
"retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"create_mental_model",
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
]
|
||||
for tool_name in bank_scoped_tools:
|
||||
tool = tools[tool_name]
|
||||
sig = inspect.signature(tool.fn)
|
||||
assert "bank_id" not in sig.parameters, f"{tool_name} should NOT have bank_id param in single-bank mode"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -242,19 +340,26 @@ async def test_middleware_handles_both_endpoints(mock_memory):
|
||||
assert "recall" in multi_bank_tools
|
||||
assert "list_banks" in multi_bank_tools
|
||||
assert "create_bank" in multi_bank_tools
|
||||
assert "list_mental_models" in multi_bank_tools
|
||||
assert "create_mental_model" in multi_bank_tools
|
||||
|
||||
# Single-bank should only have scoped tools
|
||||
assert "retain" in single_bank_tools
|
||||
assert "recall" in single_bank_tools
|
||||
assert "list_mental_models" in single_bank_tools
|
||||
assert "create_mental_model" in single_bank_tools
|
||||
assert "list_banks" not in single_bank_tools
|
||||
assert "create_bank" not in single_bank_tools
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_logic_from_url_path():
|
||||
"""Test that routing correctly selects server based on URL structure."""
|
||||
"""Test that routing correctly selects server based on URL structure.
|
||||
|
||||
Simulates the path parsing logic from MCPMiddleware.__call__ after the
|
||||
prefix has been stripped. Any first path segment is treated as a bank_id.
|
||||
"""
|
||||
from hindsight_api.api.mcp import MCPMiddleware
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
# Mock memory
|
||||
mock_memory = MagicMock()
|
||||
@@ -263,28 +368,23 @@ async def test_routing_logic_from_url_path():
|
||||
middleware = MCPMiddleware(None, mock_memory)
|
||||
|
||||
# Simulate different URL patterns and verify routing
|
||||
# Path is what remains after stripping the /mcp prefix
|
||||
test_cases = [
|
||||
# (path_after_stripping_mcp, expected_bank_id_from_path, expected_bank_id, description)
|
||||
# (path_after_prefix_strip, expected_bank_id_from_path, expected_bank_id, description)
|
||||
("/alice/messages", True, "alice", "Bank ID in path with endpoint"),
|
||||
("/my-agent-123/", True, "my-agent-123", "Bank ID in path with trailing slash"),
|
||||
("ciccio/messages", True, "ciccio", "Bank ID without leading slash (after mount strip)"),
|
||||
("bob", True, "bob", "Bank ID only, no leading slash"),
|
||||
("/messages", False, None, "MCP endpoint, no bank ID"),
|
||||
("/sse/", True, "sse", "Bank named 'sse' routes to single-bank"),
|
||||
("/messages/", True, "messages", "Bank named 'messages' routes to single-bank"),
|
||||
("/", False, None, "Root path, no bank ID"),
|
||||
]
|
||||
|
||||
for path, expected_bank_from_path, expected_bank_id, description in test_cases:
|
||||
# Simulate the path parsing logic with leading slash normalization
|
||||
if path and not path.startswith("/"):
|
||||
path = "/" + path
|
||||
|
||||
bank_id = None
|
||||
bank_id_from_path = False
|
||||
MCP_ENDPOINTS = {"sse", "messages"}
|
||||
|
||||
if path.startswith("/") and len(path) > 1:
|
||||
parts = path[1:].split("/", 1)
|
||||
if parts[0] and parts[0] not in MCP_ENDPOINTS:
|
||||
if parts[0]:
|
||||
bank_id = parts[0]
|
||||
bank_id_from_path = True
|
||||
|
||||
|
||||
@@ -1,10 +1,17 @@
|
||||
"""Tests for the shared MCP tools module."""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.mcp_tools import build_content_dict, parse_timestamp
|
||||
from hindsight_api.mcp_tools import (
|
||||
MCPToolsConfig,
|
||||
_validate_mental_model_inputs,
|
||||
build_content_dict,
|
||||
parse_timestamp,
|
||||
register_mcp_tools,
|
||||
)
|
||||
|
||||
|
||||
class TestParseTimestamp:
|
||||
@@ -61,3 +68,579 @@ class TestBuildContentDict:
|
||||
result, error = build_content_dict("test content", "test_context", None)
|
||||
assert error is None
|
||||
assert "event_date" not in result
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Mental Model MCP Tool Tests
|
||||
# =========================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_memory():
|
||||
"""Create a mock MemoryEngine with mental model methods."""
|
||||
memory = MagicMock()
|
||||
memory.list_mental_models = AsyncMock(
|
||||
return_value=[
|
||||
{"id": "mm-1", "name": "Coding Prefs", "source_query": "coding preferences?", "content": "Prefers Python"},
|
||||
{"id": "mm-2", "name": "Goals", "source_query": "current goals?", "content": "Ship v2"},
|
||||
]
|
||||
)
|
||||
memory.get_mental_model = AsyncMock(
|
||||
return_value={
|
||||
"id": "mm-1",
|
||||
"name": "Coding Prefs",
|
||||
"source_query": "coding preferences?",
|
||||
"content": "Prefers Python",
|
||||
}
|
||||
)
|
||||
memory.create_mental_model = AsyncMock(return_value={"id": "mm-new"})
|
||||
memory.submit_async_refresh_mental_model = AsyncMock(return_value={"operation_id": "op-123"})
|
||||
memory.update_mental_model = AsyncMock(
|
||||
return_value={
|
||||
"id": "mm-1",
|
||||
"name": "Updated Name",
|
||||
"source_query": "new query?",
|
||||
"content": "Updated",
|
||||
}
|
||||
)
|
||||
memory.delete_mental_model = AsyncMock(return_value=True)
|
||||
return memory
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mcp_server_with_mental_models(mock_memory):
|
||||
"""Create a FastMCP server with mental model tools registered (multi-bank mode)."""
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP("test", stateless_http=True)
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=lambda: "test-bank",
|
||||
include_bank_id_param=True,
|
||||
tools={
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"create_mental_model",
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
},
|
||||
)
|
||||
register_mcp_tools(mcp, mock_memory, config)
|
||||
return mcp
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mcp_server_single_bank(mock_memory):
|
||||
"""Create a FastMCP server with mental model tools registered (single-bank mode)."""
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP("test")
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=lambda: "fixed-bank",
|
||||
include_bank_id_param=False,
|
||||
tools={
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"create_mental_model",
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
},
|
||||
)
|
||||
register_mcp_tools(mcp, mock_memory, config)
|
||||
return mcp
|
||||
|
||||
|
||||
class TestMentalModelToolRegistration:
|
||||
"""Test that mental model tools are registered correctly."""
|
||||
|
||||
def test_tools_registered_multi_bank(self, mcp_server_with_mental_models):
|
||||
tools = mcp_server_with_mental_models._tool_manager._tools
|
||||
expected = {
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"create_mental_model",
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
}
|
||||
assert expected == set(tools.keys())
|
||||
|
||||
def test_tools_registered_single_bank(self, mcp_server_single_bank):
|
||||
tools = mcp_server_single_bank._tool_manager._tools
|
||||
expected = {
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"create_mental_model",
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
}
|
||||
assert expected == set(tools.keys())
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_mental_models_propagates_request_context(self, mock_memory):
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP("test", stateless_http=True)
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=lambda: "test-bank",
|
||||
api_key_resolver=lambda: "test-api-key",
|
||||
include_bank_id_param=True,
|
||||
tools={"list_mental_models"},
|
||||
)
|
||||
register_mcp_tools(mcp, mock_memory, config)
|
||||
await _tools(mcp)["list_mental_models"].fn()
|
||||
request_context = mock_memory.list_mental_models.call_args.kwargs["request_context"]
|
||||
assert request_context.api_key == "test-api-key"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_mental_model_propagates_request_context(self, mock_memory):
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP("test", stateless_http=True)
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=lambda: "test-bank",
|
||||
api_key_resolver=lambda: "test-api-key",
|
||||
include_bank_id_param=True,
|
||||
tools={"create_mental_model"},
|
||||
)
|
||||
register_mcp_tools(mcp, mock_memory, config)
|
||||
await _tools(mcp)["create_mental_model"].fn(name="Test", source_query="query")
|
||||
request_context = mock_memory.create_mental_model.call_args.kwargs["request_context"]
|
||||
assert request_context.api_key == "test-api-key"
|
||||
|
||||
def test_mental_model_tools_in_default_set(self):
|
||||
"""Mental model tools should be in the default tools set when config.tools is None."""
|
||||
from fastmcp import FastMCP
|
||||
|
||||
memory = MagicMock()
|
||||
# Mock all engine methods that tools reference
|
||||
memory.retain_batch_async = AsyncMock()
|
||||
memory.submit_async_retain = AsyncMock(return_value={"operation_id": "op"})
|
||||
memory.recall_async = AsyncMock(return_value=MagicMock(results=[]))
|
||||
memory.reflect_async = AsyncMock()
|
||||
memory.list_banks = AsyncMock(return_value=[])
|
||||
memory.get_bank_profile = AsyncMock(return_value={})
|
||||
memory.update_bank = AsyncMock()
|
||||
memory.list_mental_models = AsyncMock(return_value=[])
|
||||
memory.get_mental_model = AsyncMock()
|
||||
memory.create_mental_model = AsyncMock()
|
||||
memory.submit_async_refresh_mental_model = AsyncMock()
|
||||
memory.update_mental_model = AsyncMock()
|
||||
memory.delete_mental_model = AsyncMock()
|
||||
|
||||
mcp = FastMCP("test", stateless_http=True)
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=lambda: "bank",
|
||||
include_bank_id_param=True,
|
||||
tools=None, # Default - all tools
|
||||
)
|
||||
register_mcp_tools(mcp, memory, config)
|
||||
tools = mcp._tool_manager._tools
|
||||
assert "list_mental_models" in tools
|
||||
assert "create_mental_model" in tools
|
||||
assert "refresh_mental_model" in tools
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_bank_mcp_server(mock_memory):
|
||||
"""Create a multi-bank MCP server where bank_id_resolver returns None."""
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP("test", stateless_http=True)
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=lambda: None,
|
||||
include_bank_id_param=True,
|
||||
tools={
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"create_mental_model",
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
},
|
||||
)
|
||||
register_mcp_tools(mcp, mock_memory, config)
|
||||
return mcp
|
||||
|
||||
|
||||
def _tools(mcp_server):
|
||||
"""Helper to get tools dict from MCP server."""
|
||||
return mcp_server._tool_manager._tools
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestListMentalModels:
|
||||
async def test_list_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
result = await _tools(mcp_server_with_mental_models)["list_mental_models"].fn()
|
||||
assert '"mm-1"' in result
|
||||
assert '"mm-2"' in result
|
||||
mock_memory.list_mental_models.assert_called_once()
|
||||
assert mock_memory.list_mental_models.call_args.kwargs["bank_id"] == "test-bank"
|
||||
|
||||
async def test_list_with_bank_id_override(self, mcp_server_with_mental_models, mock_memory):
|
||||
"""Explicit bank_id should override the resolver."""
|
||||
await _tools(mcp_server_with_mental_models)["list_mental_models"].fn(bank_id="other-bank")
|
||||
assert mock_memory.list_mental_models.call_args.kwargs["bank_id"] == "other-bank"
|
||||
|
||||
async def test_list_with_tags(self, mcp_server_with_mental_models, mock_memory):
|
||||
await _tools(mcp_server_with_mental_models)["list_mental_models"].fn(tags=["work"])
|
||||
assert mock_memory.list_mental_models.call_args.kwargs["tags"] == ["work"]
|
||||
|
||||
async def test_list_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
result = await _tools(mcp_server_single_bank)["list_mental_models"].fn()
|
||||
assert isinstance(result, dict)
|
||||
assert len(result["items"]) == 2
|
||||
assert mock_memory.list_mental_models.call_args.kwargs["bank_id"] == "fixed-bank"
|
||||
|
||||
async def test_list_no_bank_returns_error(self, no_bank_mcp_server):
|
||||
result = await _tools(no_bank_mcp_server)["list_mental_models"].fn()
|
||||
assert "error" in result
|
||||
|
||||
async def test_list_engine_error_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
mock_memory.list_mental_models.side_effect = RuntimeError("DB connection lost")
|
||||
result = await _tools(mcp_server_with_mental_models)["list_mental_models"].fn()
|
||||
assert "error" in result
|
||||
assert "DB connection lost" in result
|
||||
|
||||
async def test_list_engine_error_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
mock_memory.list_mental_models.side_effect = RuntimeError("DB connection lost")
|
||||
result = await _tools(mcp_server_single_bank)["list_mental_models"].fn()
|
||||
assert isinstance(result, dict)
|
||||
assert "error" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestGetMentalModel:
|
||||
async def test_get_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
result = await _tools(mcp_server_with_mental_models)["get_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert '"mm-1"' in result
|
||||
assert mock_memory.get_mental_model.call_args.kwargs["mental_model_id"] == "mm-1"
|
||||
|
||||
async def test_get_with_bank_id_override(self, mcp_server_with_mental_models, mock_memory):
|
||||
await _tools(mcp_server_with_mental_models)["get_mental_model"].fn(mental_model_id="mm-1", bank_id="other-bank")
|
||||
assert mock_memory.get_mental_model.call_args.kwargs["bank_id"] == "other-bank"
|
||||
|
||||
async def test_get_not_found_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
mock_memory.get_mental_model.return_value = None
|
||||
result = await _tools(mcp_server_with_mental_models)["get_mental_model"].fn(mental_model_id="missing")
|
||||
assert "not found" in result
|
||||
|
||||
async def test_get_not_found_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
mock_memory.get_mental_model.return_value = None
|
||||
result = await _tools(mcp_server_single_bank)["get_mental_model"].fn(mental_model_id="missing")
|
||||
assert isinstance(result, dict)
|
||||
assert "not found" in result["error"]
|
||||
|
||||
async def test_get_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
result = await _tools(mcp_server_single_bank)["get_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert isinstance(result, dict)
|
||||
assert result["id"] == "mm-1"
|
||||
|
||||
async def test_get_no_bank_returns_error(self, no_bank_mcp_server):
|
||||
result = await _tools(no_bank_mcp_server)["get_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert "error" in result
|
||||
|
||||
async def test_get_engine_error(self, mcp_server_with_mental_models, mock_memory):
|
||||
mock_memory.get_mental_model.side_effect = RuntimeError("DB error")
|
||||
result = await _tools(mcp_server_with_mental_models)["get_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert "error" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestCreateMentalModel:
|
||||
async def test_create_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
result = await _tools(mcp_server_with_mental_models)["create_mental_model"].fn(
|
||||
name="Test Model",
|
||||
source_query="What are the user's preferences?",
|
||||
)
|
||||
assert '"mm-new"' in result
|
||||
assert '"op-123"' in result
|
||||
mock_memory.create_mental_model.assert_called_once()
|
||||
call_kwargs = mock_memory.create_mental_model.call_args.kwargs
|
||||
assert call_kwargs["name"] == "Test Model"
|
||||
assert call_kwargs["source_query"] == "What are the user's preferences?"
|
||||
assert call_kwargs["content"] == "Generating content..."
|
||||
# Verify async refresh was scheduled
|
||||
mock_memory.submit_async_refresh_mental_model.assert_called_once()
|
||||
assert mock_memory.submit_async_refresh_mental_model.call_args.kwargs["mental_model_id"] == "mm-new"
|
||||
|
||||
async def test_create_with_custom_id(self, mcp_server_with_mental_models, mock_memory):
|
||||
await _tools(mcp_server_with_mental_models)["create_mental_model"].fn(
|
||||
name="Test", source_query="query", mental_model_id="custom-id"
|
||||
)
|
||||
assert mock_memory.create_mental_model.call_args.kwargs["mental_model_id"] == "custom-id"
|
||||
|
||||
async def test_create_with_tags_and_max_tokens(self, mcp_server_with_mental_models, mock_memory):
|
||||
await _tools(mcp_server_with_mental_models)["create_mental_model"].fn(
|
||||
name="Test", source_query="query", tags=["work", "coding"], max_tokens=4096
|
||||
)
|
||||
call_kwargs = mock_memory.create_mental_model.call_args.kwargs
|
||||
assert call_kwargs["tags"] == ["work", "coding"]
|
||||
assert call_kwargs["max_tokens"] == 4096
|
||||
|
||||
async def test_create_with_bank_id_override(self, mcp_server_with_mental_models, mock_memory):
|
||||
await _tools(mcp_server_with_mental_models)["create_mental_model"].fn(
|
||||
name="Test", source_query="query", bank_id="other-bank"
|
||||
)
|
||||
assert mock_memory.create_mental_model.call_args.kwargs["bank_id"] == "other-bank"
|
||||
assert mock_memory.submit_async_refresh_mental_model.call_args.kwargs["bank_id"] == "other-bank"
|
||||
|
||||
async def test_create_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
result = await _tools(mcp_server_single_bank)["create_mental_model"].fn(name="Test", source_query="query")
|
||||
assert isinstance(result, dict)
|
||||
assert result["mental_model_id"] == "mm-new"
|
||||
assert result["operation_id"] == "op-123"
|
||||
|
||||
async def test_create_no_bank_returns_error(self, no_bank_mcp_server):
|
||||
result = await _tools(no_bank_mcp_server)["create_mental_model"].fn(name="Test", source_query="query")
|
||||
assert "error" in result
|
||||
|
||||
async def test_create_value_error_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
"""ValueError from engine (e.g. invalid ID format) should return error, not crash."""
|
||||
mock_memory.create_mental_model.side_effect = ValueError("ID must be alphanumeric lowercase")
|
||||
result = await _tools(mcp_server_with_mental_models)["create_mental_model"].fn(
|
||||
name="Test", source_query="query", mental_model_id="INVALID!!"
|
||||
)
|
||||
assert "alphanumeric" in result
|
||||
|
||||
async def test_create_value_error_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
mock_memory.create_mental_model.side_effect = ValueError("ID must be alphanumeric lowercase")
|
||||
result = await _tools(mcp_server_single_bank)["create_mental_model"].fn(
|
||||
name="Test", source_query="query", mental_model_id="INVALID!!"
|
||||
)
|
||||
assert isinstance(result, dict)
|
||||
assert "alphanumeric" in result["error"]
|
||||
|
||||
async def test_create_engine_error(self, mcp_server_with_mental_models, mock_memory):
|
||||
mock_memory.create_mental_model.side_effect = RuntimeError("DB error")
|
||||
result = await _tools(mcp_server_with_mental_models)["create_mental_model"].fn(
|
||||
name="Test", source_query="query"
|
||||
)
|
||||
assert "error" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestUpdateMentalModel:
|
||||
async def test_update_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
result = await _tools(mcp_server_with_mental_models)["update_mental_model"].fn(
|
||||
mental_model_id="mm-1", name="Updated Name"
|
||||
)
|
||||
assert '"Updated Name"' in result
|
||||
call_kwargs = mock_memory.update_mental_model.call_args.kwargs
|
||||
assert call_kwargs["name"] == "Updated Name"
|
||||
assert call_kwargs["source_query"] is None # Not updated
|
||||
|
||||
async def test_update_multiple_fields(self, mcp_server_with_mental_models, mock_memory):
|
||||
await _tools(mcp_server_with_mental_models)["update_mental_model"].fn(
|
||||
mental_model_id="mm-1", name="New Name", source_query="new query?", tags=["updated"], max_tokens=4096
|
||||
)
|
||||
call_kwargs = mock_memory.update_mental_model.call_args.kwargs
|
||||
assert call_kwargs["name"] == "New Name"
|
||||
assert call_kwargs["source_query"] == "new query?"
|
||||
assert call_kwargs["tags"] == ["updated"]
|
||||
assert call_kwargs["max_tokens"] == 4096
|
||||
|
||||
async def test_update_with_bank_id_override(self, mcp_server_with_mental_models, mock_memory):
|
||||
await _tools(mcp_server_with_mental_models)["update_mental_model"].fn(
|
||||
mental_model_id="mm-1", name="X", bank_id="other-bank"
|
||||
)
|
||||
assert mock_memory.update_mental_model.call_args.kwargs["bank_id"] == "other-bank"
|
||||
|
||||
async def test_update_not_found_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
mock_memory.update_mental_model.return_value = None
|
||||
result = await _tools(mcp_server_with_mental_models)["update_mental_model"].fn(
|
||||
mental_model_id="missing", name="X"
|
||||
)
|
||||
assert "not found" in result
|
||||
|
||||
async def test_update_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
result = await _tools(mcp_server_single_bank)["update_mental_model"].fn(mental_model_id="mm-1", name="Updated")
|
||||
assert isinstance(result, dict)
|
||||
assert mock_memory.update_mental_model.call_args.kwargs["bank_id"] == "fixed-bank"
|
||||
|
||||
async def test_update_not_found_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
mock_memory.update_mental_model.return_value = None
|
||||
result = await _tools(mcp_server_single_bank)["update_mental_model"].fn(mental_model_id="missing", name="X")
|
||||
assert isinstance(result, dict)
|
||||
assert "not found" in result["error"]
|
||||
|
||||
async def test_update_no_bank_returns_error(self, no_bank_mcp_server):
|
||||
result = await _tools(no_bank_mcp_server)["update_mental_model"].fn(mental_model_id="mm-1", name="X")
|
||||
assert "error" in result
|
||||
|
||||
async def test_update_engine_error(self, mcp_server_with_mental_models, mock_memory):
|
||||
mock_memory.update_mental_model.side_effect = RuntimeError("DB error")
|
||||
result = await _tools(mcp_server_with_mental_models)["update_mental_model"].fn(mental_model_id="mm-1", name="X")
|
||||
assert "error" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestDeleteMentalModel:
|
||||
async def test_delete_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
result = await _tools(mcp_server_with_mental_models)["delete_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert '"deleted"' in result
|
||||
assert mock_memory.delete_mental_model.call_args.kwargs["mental_model_id"] == "mm-1"
|
||||
|
||||
async def test_delete_with_bank_id_override(self, mcp_server_with_mental_models, mock_memory):
|
||||
await _tools(mcp_server_with_mental_models)["delete_mental_model"].fn(
|
||||
mental_model_id="mm-1", bank_id="other-bank"
|
||||
)
|
||||
assert mock_memory.delete_mental_model.call_args.kwargs["bank_id"] == "other-bank"
|
||||
|
||||
async def test_delete_not_found_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
mock_memory.delete_mental_model.return_value = False
|
||||
result = await _tools(mcp_server_with_mental_models)["delete_mental_model"].fn(mental_model_id="missing")
|
||||
assert "not found" in result
|
||||
|
||||
async def test_delete_not_found_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
mock_memory.delete_mental_model.return_value = False
|
||||
result = await _tools(mcp_server_single_bank)["delete_mental_model"].fn(mental_model_id="missing")
|
||||
assert isinstance(result, dict)
|
||||
assert "not found" in result["error"]
|
||||
|
||||
async def test_delete_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
result = await _tools(mcp_server_single_bank)["delete_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert isinstance(result, dict)
|
||||
assert result["status"] == "deleted"
|
||||
|
||||
async def test_delete_no_bank_returns_error(self, no_bank_mcp_server):
|
||||
result = await _tools(no_bank_mcp_server)["delete_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert "error" in result
|
||||
|
||||
async def test_delete_engine_error(self, mcp_server_with_mental_models, mock_memory):
|
||||
mock_memory.delete_mental_model.side_effect = RuntimeError("DB error")
|
||||
result = await _tools(mcp_server_with_mental_models)["delete_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert "error" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestRefreshMentalModel:
|
||||
async def test_refresh_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
result = await _tools(mcp_server_with_mental_models)["refresh_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert '"op-123"' in result
|
||||
assert '"queued"' in result
|
||||
|
||||
async def test_refresh_with_bank_id_override(self, mcp_server_with_mental_models, mock_memory):
|
||||
await _tools(mcp_server_with_mental_models)["refresh_mental_model"].fn(
|
||||
mental_model_id="mm-1", bank_id="other-bank"
|
||||
)
|
||||
assert mock_memory.submit_async_refresh_mental_model.call_args.kwargs["bank_id"] == "other-bank"
|
||||
|
||||
async def test_refresh_not_found_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
mock_memory.submit_async_refresh_mental_model.side_effect = ValueError("Mental model 'missing' not found")
|
||||
result = await _tools(mcp_server_with_mental_models)["refresh_mental_model"].fn(mental_model_id="missing")
|
||||
assert "not found" in result
|
||||
|
||||
async def test_refresh_not_found_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
mock_memory.submit_async_refresh_mental_model.side_effect = ValueError("not found")
|
||||
result = await _tools(mcp_server_single_bank)["refresh_mental_model"].fn(mental_model_id="missing")
|
||||
assert isinstance(result, dict)
|
||||
assert "not found" in result["error"]
|
||||
|
||||
async def test_refresh_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
result = await _tools(mcp_server_single_bank)["refresh_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert isinstance(result, dict)
|
||||
assert result["operation_id"] == "op-123"
|
||||
|
||||
async def test_refresh_no_bank_returns_error(self, no_bank_mcp_server):
|
||||
result = await _tools(no_bank_mcp_server)["refresh_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert "error" in result
|
||||
|
||||
async def test_refresh_engine_error(self, mcp_server_with_mental_models, mock_memory):
|
||||
mock_memory.submit_async_refresh_mental_model.side_effect = RuntimeError("DB error")
|
||||
result = await _tools(mcp_server_with_mental_models)["refresh_mental_model"].fn(mental_model_id="mm-1")
|
||||
assert "error" in result
|
||||
|
||||
|
||||
class TestValidateMentalModelInputs:
|
||||
"""Tests for the _validate_mental_model_inputs helper."""
|
||||
|
||||
def test_valid_inputs(self):
|
||||
assert _validate_mental_model_inputs(name="Test", source_query="query", max_tokens=2048) is None
|
||||
|
||||
def test_none_inputs(self):
|
||||
assert _validate_mental_model_inputs() is None
|
||||
|
||||
def test_empty_name(self):
|
||||
result = _validate_mental_model_inputs(name="")
|
||||
assert result == "name cannot be empty"
|
||||
|
||||
def test_whitespace_name(self):
|
||||
result = _validate_mental_model_inputs(name=" ")
|
||||
assert result == "name cannot be empty"
|
||||
|
||||
def test_empty_source_query(self):
|
||||
result = _validate_mental_model_inputs(source_query="")
|
||||
assert result == "source_query cannot be empty"
|
||||
|
||||
def test_whitespace_source_query(self):
|
||||
result = _validate_mental_model_inputs(source_query=" \t ")
|
||||
assert result == "source_query cannot be empty"
|
||||
|
||||
def test_max_tokens_too_low(self):
|
||||
result = _validate_mental_model_inputs(max_tokens=0)
|
||||
assert "max_tokens must be between 256 and 8192" in result
|
||||
|
||||
def test_max_tokens_too_high(self):
|
||||
result = _validate_mental_model_inputs(max_tokens=10000)
|
||||
assert "max_tokens must be between 256 and 8192" in result
|
||||
|
||||
def test_max_tokens_at_lower_bound(self):
|
||||
assert _validate_mental_model_inputs(max_tokens=256) is None
|
||||
|
||||
def test_max_tokens_at_upper_bound(self):
|
||||
assert _validate_mental_model_inputs(max_tokens=8192) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestMentalModelInputValidation:
|
||||
"""Tests that validation is applied in create/update tools before engine calls."""
|
||||
|
||||
async def test_create_empty_name_returns_error_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
result = await _tools(mcp_server_with_mental_models)["create_mental_model"].fn(name="", source_query="query")
|
||||
assert "name cannot be empty" in result
|
||||
mock_memory.create_mental_model.assert_not_called()
|
||||
|
||||
async def test_create_empty_source_query_returns_error_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
result = await _tools(mcp_server_with_mental_models)["create_mental_model"].fn(name="Test", source_query="")
|
||||
assert "source_query cannot be empty" in result
|
||||
mock_memory.create_mental_model.assert_not_called()
|
||||
|
||||
async def test_create_max_tokens_too_low_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
result = await _tools(mcp_server_with_mental_models)["create_mental_model"].fn(
|
||||
name="Test", source_query="query", max_tokens=0
|
||||
)
|
||||
assert "max_tokens must be between 256 and 8192" in result
|
||||
mock_memory.create_mental_model.assert_not_called()
|
||||
|
||||
async def test_create_max_tokens_too_high_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
result = await _tools(mcp_server_single_bank)["create_mental_model"].fn(
|
||||
name="Test", source_query="query", max_tokens=10000
|
||||
)
|
||||
assert isinstance(result, dict)
|
||||
assert "max_tokens must be between 256 and 8192" in result["error"]
|
||||
mock_memory.create_mental_model.assert_not_called()
|
||||
|
||||
async def test_update_empty_name_returns_error_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
result = await _tools(mcp_server_with_mental_models)["update_mental_model"].fn(mental_model_id="mm-1", name="")
|
||||
assert "name cannot be empty" in result
|
||||
mock_memory.update_mental_model.assert_not_called()
|
||||
|
||||
async def test_update_empty_name_returns_error_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
result = await _tools(mcp_server_single_bank)["update_mental_model"].fn(mental_model_id="mm-1", name=" ")
|
||||
assert isinstance(result, dict)
|
||||
assert "name cannot be empty" in result["error"]
|
||||
mock_memory.update_mental_model.assert_not_called()
|
||||
|
||||
async def test_not_found_error_includes_bank_id_multi_bank(self, mcp_server_with_mental_models, mock_memory):
|
||||
mock_memory.get_mental_model.return_value = None
|
||||
result = await _tools(mcp_server_with_mental_models)["get_mental_model"].fn(mental_model_id="missing")
|
||||
assert "test-bank" in result
|
||||
|
||||
async def test_not_found_error_includes_bank_id_single_bank(self, mcp_server_single_bank, mock_memory):
|
||||
mock_memory.get_mental_model.return_value = None
|
||||
result = await _tools(mcp_server_single_bank)["get_mental_model"].fn(mental_model_id="missing")
|
||||
assert isinstance(result, dict)
|
||||
assert "fixed-bank" in result["error"]
|
||||
|
||||
@@ -852,3 +852,61 @@ class TestMentalModelRefreshTagSecurity:
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_refresh_mental_model_with_directives(self, memory: MemoryEngine, request_context):
|
||||
"""Test that refreshing a mental model with directives works correctly."""
|
||||
bank_id = f"test-refresh-directives-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Create a directive
|
||||
directive = await memory.create_directive(
|
||||
bank_id=bank_id,
|
||||
name="Response Style",
|
||||
content="Always be concise and professional",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Create a concept mental model to refresh
|
||||
concept = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Team Info",
|
||||
source_query="Team information summary",
|
||||
content="Initial team information",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Add some memories
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "Alice is the team lead and handles project planning."},
|
||||
{"content": "Bob is a senior engineer who mentors junior developers."},
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Wait for retain to complete
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Refresh the concept mental model (this should include directive in based_on)
|
||||
refreshed = await memory.refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=concept["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Wait for background tasks to complete
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
# Verify the refresh completed without errors
|
||||
assert refreshed is not None
|
||||
assert refreshed["content"] is not None
|
||||
|
||||
# Get the updated mental model
|
||||
updated = await memory.get_mental_model(bank_id, concept["id"], request_context=request_context)
|
||||
assert updated["content"] != "Initial team information"
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
"""
|
||||
Test reflect endpoint with empty based_on (no memories scenario).
|
||||
|
||||
This test verifies that the API returns the correct based_on format:
|
||||
- v0.3.0 (old): returned based_on as list []
|
||||
- v0.4.0+ (current): returns based_on as object {"memories": [], "mental_models": [], "directives": []}
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import httpx
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
"""Create an async test client for the FastAPI app."""
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_with_no_memories_empty_bank(api_client):
|
||||
"""Test reflect on an empty bank (no memories) with include.facts enabled."""
|
||||
bank_id = "test_empty_bank"
|
||||
|
||||
# Reflect on empty bank with facts requested
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/reflect",
|
||||
json={
|
||||
"query": "What do you know about machine learning?",
|
||||
"budget": "low",
|
||||
"include": {
|
||||
"facts": {} # Request facts but bank is empty
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# DEBUG: Print what the API actually returned
|
||||
import json
|
||||
print("\n" + "="*80)
|
||||
print("API Response:")
|
||||
print(json.dumps(data, indent=2))
|
||||
print("="*80 + "\n")
|
||||
|
||||
# Verify response structure
|
||||
assert "text" in data
|
||||
assert "based_on" in data
|
||||
|
||||
# The API should return based_on as either:
|
||||
# 1. null/None (if include.facts not set)
|
||||
# 2. {"memories": [], "mental_models": [], "directives": []} (if include.facts set but empty)
|
||||
# It should NEVER return based_on: []
|
||||
|
||||
based_on = data.get("based_on")
|
||||
if based_on is not None:
|
||||
assert isinstance(based_on, dict), f"based_on should be dict or null, got {type(based_on)}: {based_on}"
|
||||
assert not isinstance(based_on, list), f"based_on should NEVER be a list! Got: {based_on}"
|
||||
assert "memories" in based_on
|
||||
assert "mental_models" in based_on
|
||||
assert "directives" in based_on
|
||||
# All should be empty lists
|
||||
assert based_on["memories"] == []
|
||||
assert based_on["mental_models"] == []
|
||||
assert based_on["directives"] == []
|
||||
|
||||
# Verify the structure is parseable as proper types
|
||||
assert isinstance(data["text"], str)
|
||||
if based_on is not None:
|
||||
# Verify it's the v0.4.0+ format (object with arrays)
|
||||
assert isinstance(based_on["memories"], list)
|
||||
assert isinstance(based_on["mental_models"], list)
|
||||
assert isinstance(based_on["directives"], list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_without_include_facts(api_client):
|
||||
"""Test reflect without requesting facts (based_on should be None)."""
|
||||
bank_id = "test_no_facts"
|
||||
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/reflect",
|
||||
json={
|
||||
"query": "Hello world",
|
||||
"budget": "low"
|
||||
# No include.facts
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# When include.facts is not set, based_on should not be in response (or be null)
|
||||
based_on = data.get("based_on")
|
||||
assert based_on is None, f"based_on should be None when not requested, got {type(based_on)}: {based_on}"
|
||||
|
||||
# Verify structure
|
||||
assert isinstance(data["text"], str)
|
||||
@@ -0,0 +1,45 @@
|
||||
"""
|
||||
Test to verify reflect operation creates proper span hierarchy.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_creates_child_spans(memory, request_context):
|
||||
"""Test that reflect operation creates child LLM spans."""
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api.tracing import initialize_tracing, get_span_recorder, create_span_recorder
|
||||
|
||||
# Initialize tracing with a mock endpoint
|
||||
initialize_tracing(
|
||||
service_name="test-hindsight",
|
||||
endpoint="http://localhost:4318",
|
||||
deployment_environment="test"
|
||||
)
|
||||
|
||||
# Create span recorder
|
||||
recorder = create_span_recorder()
|
||||
|
||||
bank_id = f"test-reflect-hierarchy-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Add some memories
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Paris is the capital of France",
|
||||
context="Geography",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Run reflect
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What is the capital of France?",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"Reflect result: {result.text[:100]}")
|
||||
print(f"Usage: {result.usage}")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -0,0 +1,407 @@
|
||||
"""
|
||||
Unit tests for OpenTelemetry tracing instrumentation.
|
||||
|
||||
Tests the tracing module's ability to record LLM calls with GenAI semantic conventions.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.tracing import (
|
||||
PROVIDER_NAME_MAPPING,
|
||||
GenAIAttributes,
|
||||
LLMSpanRecorder,
|
||||
NoOpLLMSpanRecorder,
|
||||
_truncate_content,
|
||||
create_operation_span,
|
||||
initialize_tracing,
|
||||
is_tracing_enabled,
|
||||
)
|
||||
|
||||
|
||||
def test_provider_name_mapping():
|
||||
"""Test that provider names are correctly mapped to GenAI conventions."""
|
||||
assert PROVIDER_NAME_MAPPING["openai"] == "openai"
|
||||
assert PROVIDER_NAME_MAPPING["anthropic"] == "anthropic"
|
||||
assert PROVIDER_NAME_MAPPING["gemini"] == "google"
|
||||
assert PROVIDER_NAME_MAPPING["vertexai"] == "google"
|
||||
assert PROVIDER_NAME_MAPPING["groq"] == "groq"
|
||||
assert PROVIDER_NAME_MAPPING["ollama"] == "ollama"
|
||||
assert PROVIDER_NAME_MAPPING["openai-codex"] == "openai"
|
||||
assert PROVIDER_NAME_MAPPING["claude-code"] == "anthropic"
|
||||
|
||||
|
||||
def test_truncate_content_short():
|
||||
"""Test that short content is not truncated."""
|
||||
content = "This is a short message"
|
||||
result = _truncate_content(content)
|
||||
assert result == content
|
||||
|
||||
|
||||
def test_truncate_content_long():
|
||||
"""Test that long content is truncated."""
|
||||
content = "x" * 150000 # Exceeds MAX_CONTENT_LENGTH
|
||||
result = _truncate_content(content)
|
||||
assert len(result) < len(content)
|
||||
assert "[TRUNCATED:" in result
|
||||
assert result.startswith("x" * 100)
|
||||
|
||||
|
||||
def test_noop_span_recorder():
|
||||
"""Test that NoOpLLMSpanRecorder doesn't raise errors."""
|
||||
recorder = NoOpLLMSpanRecorder()
|
||||
# Should not raise any errors
|
||||
recorder.record_llm_call(
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
scope="test",
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
response_content="test response",
|
||||
input_tokens=10,
|
||||
output_tokens=5,
|
||||
duration=1.0,
|
||||
)
|
||||
|
||||
|
||||
def test_llm_span_recorder_format_messages():
|
||||
"""Test message formatting to GenAI convention."""
|
||||
mock_tracer = MagicMock()
|
||||
recorder = LLMSpanRecorder(mock_tracer)
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": "You are helpful"},
|
||||
{"role": "user", "content": "Hello"},
|
||||
]
|
||||
|
||||
result = recorder._format_messages(messages)
|
||||
parsed = json.loads(result)
|
||||
|
||||
assert len(parsed) == 2
|
||||
assert parsed[0]["role"] == "system"
|
||||
assert parsed[0]["content"] == "You are helpful"
|
||||
assert parsed[1]["role"] == "user"
|
||||
assert parsed[1]["content"] == "Hello"
|
||||
|
||||
|
||||
def test_llm_span_recorder_format_output():
|
||||
"""Test output formatting to GenAI convention."""
|
||||
mock_tracer = MagicMock()
|
||||
recorder = LLMSpanRecorder(mock_tracer)
|
||||
|
||||
result = recorder._format_output("Hello world", "stop")
|
||||
parsed = json.loads(result)
|
||||
|
||||
assert len(parsed) == 1
|
||||
assert parsed[0]["role"] == "assistant"
|
||||
assert parsed[0]["content"] == "Hello world"
|
||||
|
||||
|
||||
def test_llm_span_recorder_format_output_none():
|
||||
"""Test output formatting with None content."""
|
||||
mock_tracer = MagicMock()
|
||||
recorder = LLMSpanRecorder(mock_tracer)
|
||||
|
||||
result = recorder._format_output(None, None)
|
||||
parsed = json.loads(result)
|
||||
|
||||
assert parsed == []
|
||||
|
||||
|
||||
def test_llm_span_recorder_extract_system_instructions():
|
||||
"""Test system instruction extraction."""
|
||||
mock_tracer = MagicMock()
|
||||
recorder = LLMSpanRecorder(mock_tracer)
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": "You are helpful"},
|
||||
{"role": "user", "content": "Hello"},
|
||||
]
|
||||
|
||||
result = recorder._extract_system_instructions(messages)
|
||||
assert result == "You are helpful"
|
||||
|
||||
|
||||
def test_llm_span_recorder_extract_system_instructions_none():
|
||||
"""Test system instruction extraction with no system message."""
|
||||
mock_tracer = MagicMock()
|
||||
recorder = LLMSpanRecorder(mock_tracer)
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
]
|
||||
|
||||
result = recorder._extract_system_instructions(messages)
|
||||
assert result is None
|
||||
|
||||
|
||||
@patch("hindsight_api.tracing.time")
|
||||
def test_llm_span_recorder_record_success(mock_time):
|
||||
"""Test successful LLM call recording."""
|
||||
# Mock time
|
||||
mock_time.time_ns.return_value = 1000000000000 # 1 second in nanoseconds
|
||||
|
||||
# Create mock tracer and span
|
||||
mock_span = MagicMock()
|
||||
mock_tracer = MagicMock()
|
||||
mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
|
||||
|
||||
recorder = LLMSpanRecorder(mock_tracer)
|
||||
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
response_content = "Hi there!"
|
||||
|
||||
recorder.record_llm_call(
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
scope="test",
|
||||
messages=messages,
|
||||
response_content=response_content,
|
||||
input_tokens=10,
|
||||
output_tokens=5,
|
||||
duration=1.5,
|
||||
finish_reason="stop",
|
||||
error=None,
|
||||
)
|
||||
|
||||
# Verify span was created with correct name (hindsight.{scope})
|
||||
mock_tracer.start_as_current_span.assert_called_once()
|
||||
call_args = mock_tracer.start_as_current_span.call_args
|
||||
assert call_args[0][0] == "hindsight.test"
|
||||
|
||||
# Verify attributes were set
|
||||
assert mock_span.set_attribute.called
|
||||
attribute_calls = {call[0][0]: call[0][1] for call in mock_span.set_attribute.call_args_list}
|
||||
|
||||
assert attribute_calls[GenAIAttributes.OPERATION_NAME] == "chat"
|
||||
assert attribute_calls[GenAIAttributes.PROVIDER_NAME] == "openai"
|
||||
assert attribute_calls[GenAIAttributes.REQUEST_MODEL] == "gpt-4"
|
||||
assert attribute_calls[GenAIAttributes.RESPONSE_MODEL] == "gpt-4"
|
||||
assert attribute_calls[GenAIAttributes.USAGE_INPUT_TOKENS] == 10
|
||||
assert attribute_calls[GenAIAttributes.USAGE_OUTPUT_TOKENS] == 5
|
||||
assert attribute_calls["hindsight.scope"] == "test"
|
||||
|
||||
# Verify event was added
|
||||
mock_span.add_event.assert_called_once()
|
||||
event_call = mock_span.add_event.call_args
|
||||
assert event_call[0][0] == "gen_ai.client.inference.operation.details"
|
||||
|
||||
# Verify status was set to OK
|
||||
mock_span.set_status.assert_called()
|
||||
|
||||
# Verify span was ended
|
||||
mock_span.end.assert_called_once()
|
||||
|
||||
|
||||
@patch("hindsight_api.tracing.time")
|
||||
def test_llm_span_recorder_record_error(mock_time):
|
||||
"""Test error LLM call recording."""
|
||||
# Mock time
|
||||
mock_time.time_ns.return_value = 1000000000000
|
||||
|
||||
# Create mock tracer and span
|
||||
mock_span = MagicMock()
|
||||
mock_tracer = MagicMock()
|
||||
mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
|
||||
|
||||
recorder = LLMSpanRecorder(mock_tracer)
|
||||
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
error = ValueError("Test error")
|
||||
|
||||
recorder.record_llm_call(
|
||||
provider="anthropic",
|
||||
model="claude-3",
|
||||
scope="test",
|
||||
messages=messages,
|
||||
response_content=None,
|
||||
input_tokens=10,
|
||||
output_tokens=0,
|
||||
duration=0.5,
|
||||
finish_reason=None,
|
||||
error=error,
|
||||
)
|
||||
|
||||
# Verify error status was set
|
||||
mock_span.set_status.assert_called()
|
||||
status_call = mock_span.set_status.call_args[0][0]
|
||||
assert status_call.status_code.name == "ERROR"
|
||||
|
||||
# Verify error type attribute was set
|
||||
attribute_calls = {call[0][0]: call[0][1] for call in mock_span.set_attribute.call_args_list}
|
||||
assert attribute_calls[GenAIAttributes.ERROR_TYPE] == "ValueError"
|
||||
|
||||
# Verify exception was recorded
|
||||
mock_span.record_exception.assert_called_once_with(error)
|
||||
|
||||
|
||||
@patch("hindsight_api.tracing.time")
|
||||
def test_llm_span_recorder_provider_mapping(mock_time):
|
||||
"""Test that provider names are mapped correctly."""
|
||||
mock_time.time_ns.return_value = 1000000000000
|
||||
|
||||
mock_span = MagicMock()
|
||||
mock_tracer = MagicMock()
|
||||
mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
|
||||
|
||||
recorder = LLMSpanRecorder(mock_tracer)
|
||||
|
||||
# Test gemini -> google mapping
|
||||
recorder.record_llm_call(
|
||||
provider="gemini",
|
||||
model="gemini-pro",
|
||||
scope="test",
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
response_content="test",
|
||||
input_tokens=5,
|
||||
output_tokens=3,
|
||||
duration=1.0,
|
||||
)
|
||||
|
||||
attribute_calls = {call[0][0]: call[0][1] for call in mock_span.set_attribute.call_args_list}
|
||||
assert attribute_calls[GenAIAttributes.PROVIDER_NAME] == "google"
|
||||
|
||||
|
||||
# ==================== Parent Span Tests ====================
|
||||
|
||||
|
||||
def test_create_operation_span_disabled():
|
||||
"""Test that create_operation_span returns no-op when tracing is disabled."""
|
||||
# Tracing should be disabled by default
|
||||
assert not is_tracing_enabled()
|
||||
|
||||
# Should return a no-op context manager
|
||||
span = create_operation_span("test_operation", "test_bank_id")
|
||||
|
||||
# Should be usable as context manager without errors
|
||||
with span:
|
||||
pass
|
||||
|
||||
|
||||
@patch("hindsight_api.tracing._tracer")
|
||||
@patch("hindsight_api.tracing._tracing_enabled", True)
|
||||
def test_create_operation_span_enabled(mock_tracer):
|
||||
"""Test that create_operation_span creates a span when tracing is enabled."""
|
||||
# Mock the tracer
|
||||
mock_span = MagicMock()
|
||||
mock_tracer.start_as_current_span.return_value = mock_span
|
||||
|
||||
# Create operation span
|
||||
span = create_operation_span("retain", "bank123")
|
||||
|
||||
# Verify span was created with correct name
|
||||
mock_tracer.start_as_current_span.assert_called_once_with("hindsight.retain")
|
||||
|
||||
# Verify attributes were set
|
||||
mock_span.set_attribute.assert_any_call("hindsight.operation", "retain")
|
||||
mock_span.set_attribute.assert_any_call("hindsight.bank_id", "bank123")
|
||||
|
||||
|
||||
@patch("hindsight_api.tracing._tracer")
|
||||
@patch("hindsight_api.tracing._tracing_enabled", True)
|
||||
def test_create_operation_span_no_bank_id(mock_tracer):
|
||||
"""Test that create_operation_span works without bank_id."""
|
||||
mock_span = MagicMock()
|
||||
mock_tracer.start_as_current_span.return_value = mock_span
|
||||
|
||||
# Create operation span without bank_id
|
||||
span = create_operation_span("consolidation")
|
||||
|
||||
# Verify span was created
|
||||
mock_tracer.start_as_current_span.assert_called_once_with("hindsight.consolidation")
|
||||
|
||||
# Verify only operation attribute was set (not bank_id)
|
||||
assert mock_span.set_attribute.call_count == 1
|
||||
mock_span.set_attribute.assert_called_once_with("hindsight.operation", "consolidation")
|
||||
|
||||
|
||||
@patch("hindsight_api.tracing._tracer")
|
||||
@patch("hindsight_api.tracing._tracing_enabled", True)
|
||||
def test_create_operation_span_all_operations(mock_tracer):
|
||||
"""Test that all 4 operations can create parent spans."""
|
||||
mock_span = MagicMock()
|
||||
mock_tracer.start_as_current_span.return_value = mock_span
|
||||
|
||||
operations = ["retain", "consolidation", "reflect", "mental_model_refresh"]
|
||||
|
||||
for operation in operations:
|
||||
mock_tracer.reset_mock()
|
||||
mock_span.reset_mock()
|
||||
|
||||
span = create_operation_span(operation, "test_bank")
|
||||
|
||||
# Verify span was created with correct name
|
||||
mock_tracer.start_as_current_span.assert_called_once_with(f"hindsight.{operation}")
|
||||
|
||||
# Verify attributes
|
||||
mock_span.set_attribute.assert_any_call("hindsight.operation", operation)
|
||||
mock_span.set_attribute.assert_any_call("hindsight.bank_id", "test_bank")
|
||||
|
||||
|
||||
@patch("hindsight_api.tracing.time")
|
||||
@patch("hindsight_api.tracing._tracer")
|
||||
@patch("hindsight_api.tracing._tracing_enabled", True)
|
||||
def test_parent_child_span_hierarchy(mock_tracer, mock_time):
|
||||
"""Test that child LLM spans are created under parent operation spans."""
|
||||
mock_time.time_ns.return_value = 1000000000000
|
||||
|
||||
# Create mock parent span
|
||||
mock_parent_span = MagicMock()
|
||||
mock_parent_span.__enter__ = MagicMock(return_value=mock_parent_span)
|
||||
mock_parent_span.__exit__ = MagicMock(return_value=False)
|
||||
|
||||
# Create mock child span
|
||||
mock_child_span = MagicMock()
|
||||
|
||||
# Mock tracer to return parent span first, then child span
|
||||
mock_tracer.start_as_current_span.side_effect = [
|
||||
mock_parent_span, # Parent span
|
||||
MagicMock(__enter__=MagicMock(return_value=mock_child_span), __exit__=MagicMock(return_value=False)), # Child
|
||||
]
|
||||
|
||||
# Create parent operation span
|
||||
with create_operation_span("retain", "bank123"):
|
||||
# Simulate creating a child LLM span
|
||||
recorder = LLMSpanRecorder(mock_tracer)
|
||||
recorder.record_llm_call(
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
scope="retain_extract_facts",
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
response_content="response",
|
||||
input_tokens=10,
|
||||
output_tokens=5,
|
||||
duration=1.0,
|
||||
)
|
||||
|
||||
# Verify both parent and child spans were created
|
||||
assert mock_tracer.start_as_current_span.call_count == 2
|
||||
|
||||
# Verify parent span was created first
|
||||
first_call = mock_tracer.start_as_current_span.call_args_list[0]
|
||||
assert first_call[0][0] == "hindsight.retain"
|
||||
|
||||
# Verify child span was created second (hindsight.{scope})
|
||||
second_call = mock_tracer.start_as_current_span.call_args_list[1]
|
||||
assert second_call[0][0] == "hindsight.retain_extract_facts"
|
||||
|
||||
|
||||
@patch("hindsight_api.tracing._tracer")
|
||||
@patch("hindsight_api.tracing._tracing_enabled", True)
|
||||
def test_operation_span_context_manager(mock_tracer):
|
||||
"""Test that operation spans work as context managers."""
|
||||
mock_span = MagicMock()
|
||||
mock_span.__enter__ = MagicMock(return_value=mock_span)
|
||||
mock_span.__exit__ = MagicMock(return_value=False)
|
||||
mock_tracer.start_as_current_span.return_value = mock_span
|
||||
|
||||
# Use span as context manager
|
||||
with create_operation_span("reflect", "bank456"):
|
||||
# Do some work
|
||||
pass
|
||||
|
||||
# Verify span lifecycle
|
||||
mock_tracer.start_as_current_span.assert_called_once()
|
||||
mock_span.__enter__.assert_called_once()
|
||||
mock_span.__exit__.assert_called_once()
|
||||
@@ -0,0 +1,196 @@
|
||||
"""
|
||||
Integration tests for OpenTelemetry tracing with memory engine operations.
|
||||
|
||||
Tests that parent spans are correctly created for retain, consolidation, reflect,
|
||||
and mental_model_refresh operations.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("hindsight_api.engine.memory_engine.create_operation_span")
|
||||
async def test_retain_creates_parent_span(mock_create_span, memory, request_context):
|
||||
"""Test that retain operation creates a parent span."""
|
||||
# Setup
|
||||
mock_span = MagicMock()
|
||||
mock_span.__enter__ = MagicMock(return_value=mock_span)
|
||||
mock_span.__exit__ = MagicMock(return_value=False)
|
||||
mock_create_span.return_value = mock_span
|
||||
|
||||
bank_id = f"test-retain-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Execute retain (automatically creates bank if needed)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Test memory for tracing",
|
||||
context="Test context",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify parent span was created
|
||||
mock_create_span.assert_called()
|
||||
call_args = mock_create_span.call_args
|
||||
assert call_args[0][0] == "retain" # operation name
|
||||
assert call_args[0][1] == bank_id # bank_id
|
||||
|
||||
# Verify span was used as context manager
|
||||
mock_span.__enter__.assert_called()
|
||||
mock_span.__exit__.assert_called()
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("hindsight_api.engine.memory_engine.create_operation_span")
|
||||
async def test_consolidation_creates_parent_span(mock_create_span, memory, request_context):
|
||||
"""Test that consolidation operation creates a parent span."""
|
||||
# Setup
|
||||
mock_span = MagicMock()
|
||||
mock_span.__enter__ = MagicMock(return_value=mock_span)
|
||||
mock_span.__exit__ = MagicMock(return_value=False)
|
||||
mock_create_span.return_value = mock_span
|
||||
|
||||
bank_id = f"test-consolidation-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Execute consolidation (bank will be created automatically)
|
||||
await memory.run_consolidation(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify parent span was created
|
||||
mock_create_span.assert_called()
|
||||
call_args = mock_create_span.call_args
|
||||
assert call_args[0][0] == "consolidation"
|
||||
assert call_args[0][1] == bank_id
|
||||
|
||||
# Verify span was used as context manager
|
||||
mock_span.__enter__.assert_called()
|
||||
mock_span.__exit__.assert_called()
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("hindsight_api.engine.memory_engine.create_operation_span")
|
||||
async def test_reflect_creates_parent_span(mock_create_span, memory, request_context):
|
||||
"""Test that reflect operation creates a parent span."""
|
||||
# Setup
|
||||
mock_span = MagicMock()
|
||||
mock_span.__enter__ = MagicMock(return_value=mock_span)
|
||||
mock_span.__exit__ = MagicMock(return_value=False)
|
||||
mock_create_span.return_value = mock_span
|
||||
|
||||
bank_id = f"test-reflect-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Add some memories first
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Paris is the capital of France",
|
||||
context="Geography fact",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Reset mock to clear retain call
|
||||
mock_create_span.reset_mock()
|
||||
|
||||
# Execute reflect
|
||||
await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What is the capital of France?",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify parent span was created
|
||||
mock_create_span.assert_called()
|
||||
call_args = mock_create_span.call_args
|
||||
assert call_args[0][0] == "reflect"
|
||||
assert call_args[0][1] == bank_id
|
||||
|
||||
# Verify span was used as context manager
|
||||
mock_span.__enter__.assert_called()
|
||||
mock_span.__exit__.assert_called()
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("hindsight_api.engine.memory_engine.create_operation_span")
|
||||
async def test_retain_batch_creates_single_parent_span(mock_create_span, memory, request_context):
|
||||
"""Test that batch retain creates one parent span for the entire batch."""
|
||||
# Setup
|
||||
mock_span = MagicMock()
|
||||
mock_span.__enter__ = MagicMock(return_value=mock_span)
|
||||
mock_span.__exit__ = MagicMock(return_value=False)
|
||||
mock_create_span.return_value = mock_span
|
||||
|
||||
bank_id = f"test-batch-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Execute batch retain with multiple items
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "Memory 1", "context": "Context 1"},
|
||||
{"content": "Memory 2", "context": "Context 2"},
|
||||
{"content": "Memory 3", "context": "Context 3"},
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify parent span was created only once for the entire batch
|
||||
assert mock_create_span.call_count == 1
|
||||
call_args = mock_create_span.call_args
|
||||
assert call_args[0][0] == "retain"
|
||||
assert call_args[0][1] == bank_id
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("hindsight_api.tracing._tracing_enabled", False)
|
||||
@patch("hindsight_api.engine.memory_engine.create_operation_span")
|
||||
async def test_operations_work_when_tracing_disabled(mock_create_span, memory, request_context):
|
||||
"""Test that operations work correctly when tracing is disabled."""
|
||||
# Setup - create_operation_span should return a no-op context manager
|
||||
from contextlib import nullcontext
|
||||
|
||||
mock_create_span.return_value = nullcontext()
|
||||
|
||||
bank_id = f"test-no-trace-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# All operations should work without errors
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Test memory",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.run_consolidation(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="Test query",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify no errors occurred and spans were attempted to be created
|
||||
assert mock_create_span.call_count >= 3
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -0,0 +1,273 @@
|
||||
"""
|
||||
Comprehensive tracing span verification tests.
|
||||
|
||||
Verifies that all memory engine operations create correct parent and child spans
|
||||
with proper attributes and hierarchy.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skip(reason="Background consolidation causes StopIteration - need to investigate separately")
|
||||
@patch("hindsight_api.tracing._tracing_enabled", True)
|
||||
@patch("hindsight_api.tracing._tracer")
|
||||
async def test_recall_span_hierarchy(mock_tracer, memory, request_context):
|
||||
"""Test that recall creates proper parent and child spans."""
|
||||
# Setup mock spans
|
||||
mock_recall_span = MagicMock()
|
||||
mock_recall_span.__enter__ = MagicMock(return_value=mock_recall_span)
|
||||
mock_recall_span.__exit__ = MagicMock(return_value=False)
|
||||
|
||||
mock_embedding_span = MagicMock()
|
||||
mock_retrieval_span = MagicMock()
|
||||
mock_fusion_span = MagicMock()
|
||||
mock_rerank_span = MagicMock()
|
||||
|
||||
# Mock tracer to return spans in sequence
|
||||
mock_tracer.start_as_current_span.side_effect = [mock_recall_span]
|
||||
mock_tracer.start_span.side_effect = [
|
||||
mock_embedding_span,
|
||||
mock_retrieval_span,
|
||||
mock_fusion_span,
|
||||
mock_rerank_span,
|
||||
]
|
||||
|
||||
bank_id = f"test-recall-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Add some memories first
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Paris is the capital of France",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Wait a bit for any background tasks to settle
|
||||
import asyncio
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Reset mocks after retain
|
||||
mock_tracer.reset_mock()
|
||||
mock_recall_span.reset_mock()
|
||||
|
||||
# Execute recall
|
||||
await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="What is the capital of France?",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify parent span was created with start_as_current_span
|
||||
assert mock_tracer.start_as_current_span.called
|
||||
parent_call = mock_tracer.start_as_current_span.call_args
|
||||
assert parent_call[0][0] == "hindsight.recall"
|
||||
|
||||
# Verify parent span attributes were set
|
||||
recall_attrs = {call[0][0]: call[0][1] for call in mock_recall_span.set_attribute.call_args_list}
|
||||
assert "hindsight.bank_id" in recall_attrs
|
||||
assert recall_attrs["hindsight.bank_id"] == bank_id
|
||||
assert "hindsight.query" in recall_attrs
|
||||
assert "hindsight.fact_types" in recall_attrs
|
||||
assert "hindsight.thinking_budget" in recall_attrs
|
||||
assert "hindsight.max_tokens" in recall_attrs
|
||||
|
||||
# Verify child spans were created (if tracing is enabled)
|
||||
if mock_tracer.start_span.called:
|
||||
child_spans = [call[0][0] for call in mock_tracer.start_span.call_args_list]
|
||||
assert "hindsight.recall_embedding" in child_spans
|
||||
assert "hindsight.recall_retrieval" in child_spans
|
||||
assert "hindsight.recall_fusion" in child_spans
|
||||
assert "hindsight.recall_rerank" in child_spans
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mental_model_refresh_span_exists(memory, request_context):
|
||||
"""Test that mental model refresh functionality exists (span creation tested via unit tests)."""
|
||||
# This test verifies that refresh_mental_model method exists and can be called
|
||||
# The actual span creation is tested in unit tests with proper mocking
|
||||
bank_id = f"test-mmr-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Just verify the method exists - it will return None if no mental model found
|
||||
result = await memory.refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id="non-existent-id",
|
||||
request_context=request_context,
|
||||
)
|
||||
# Result will be None since mental model doesn't exist
|
||||
assert result is None
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consolidation_child_spans(memory, request_context):
|
||||
"""Test that consolidation creates child spans for its operations."""
|
||||
bank_id = f"test-cons-child-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Add memories to consolidate
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="The Eiffel Tower is in Paris",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Paris is the capital of France",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Run consolidation (this will create parent + child spans)
|
||||
await memory.run_consolidation(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Note: We can't easily verify the child spans without mocking the tracer,
|
||||
# but we can verify that consolidation completes successfully
|
||||
# The actual span creation is tested in unit tests
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_tool_call_spans(memory, request_context):
|
||||
"""Test that reflect creates tool call spans (not reflect_generation)."""
|
||||
bank_id = f"test-reflect-tools-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Add some memories
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Machine learning is a subset of AI",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Execute reflect (will create reflect_tool_call spans)
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What is machine learning?",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify reflect completed successfully
|
||||
assert result.text
|
||||
assert len(result.text) > 0
|
||||
|
||||
# The span names are verified via unit tests with mocked tracers
|
||||
# This integration test ensures the operation completes successfully
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_operations_create_spans(memory, request_context):
|
||||
"""Comprehensive test that all operations create their respective spans."""
|
||||
bank_id = f"test-all-ops-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# 1. Retain operation
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Test memory for comprehensive span test",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# 2. Recall operation
|
||||
await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="test memory",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# 3. Reflect operation
|
||||
await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What can you tell me about the test?",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# 4. Consolidation operation
|
||||
await memory.run_consolidation(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# All operations completed successfully
|
||||
# Span hierarchy verification is done in unit tests with mocked tracers
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("hindsight_api.tracing._tracing_enabled", True)
|
||||
@patch("hindsight_api.tracing._tracer")
|
||||
async def test_recall_span_attributes(mock_tracer, memory, request_context):
|
||||
"""Verify that recall spans have all required attributes."""
|
||||
# Setup mock span
|
||||
mock_span = MagicMock()
|
||||
mock_span.__enter__ = MagicMock(return_value=mock_span)
|
||||
mock_span.__exit__ = MagicMock(return_value=False)
|
||||
mock_tracer.start_as_current_span.return_value = mock_span
|
||||
|
||||
bank_id = f"test-attrs-{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Add memory
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Test content for attributes",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Reset mock
|
||||
mock_span.reset_mock()
|
||||
|
||||
# Execute recall with specific parameters
|
||||
await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="test query for attributes",
|
||||
fact_type=["world", "experience"],
|
||||
max_tokens=2048,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Collect all attributes set on the span
|
||||
attrs = {call[0][0]: call[0][1] for call in mock_span.set_attribute.call_args_list}
|
||||
|
||||
# Verify required attributes
|
||||
assert "hindsight.bank_id" in attrs
|
||||
assert "hindsight.query" in attrs
|
||||
assert "hindsight.fact_types" in attrs
|
||||
assert "hindsight.max_tokens" in attrs
|
||||
assert "hindsight.thinking_budget" in attrs
|
||||
|
||||
# Verify attribute values
|
||||
assert attrs["hindsight.bank_id"] == bank_id
|
||||
assert "test query" in attrs["hindsight.query"]
|
||||
assert attrs["hindsight.max_tokens"] == 2048
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "hindsight-cli"
|
||||
version = "0.4.9"
|
||||
version = "0.4.10"
|
||||
edition = "2021"
|
||||
authors = ["Hindsight Team"]
|
||||
description = "A beautiful CLI for Hindsight - semantic memory system"
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
@@ -489,7 +489,7 @@ class Configuration:
|
||||
return "Python SDK Debug Report:\n"\
|
||||
"OS: {env}\n"\
|
||||
"Python Version: {pyversion}\n"\
|
||||
"Version of the API: 0.4.9\n"\
|
||||
"Version of the API: 0.4.10\n"\
|
||||
"SDK Package Version: 0.0.7".\
|
||||
format(env=sys.platform, pyversion=sys.version)
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
+1
-1
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
+1
-1
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.9
|
||||
The version of the OpenAPI document: 0.4.10
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user