Compare commits
14
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e62860d232 | ||
|
|
379aa5ce4d | ||
|
|
a8ef9ad3cb | ||
|
|
db79ea8ffd | ||
|
|
ff095b4f79 | ||
|
|
e50db1adf6 | ||
|
|
863bc2a3c7 | ||
|
|
ec2c1b21e6 | ||
|
|
a72719a68f | ||
|
|
ce6c93d3b5 | ||
|
|
1ce9458347 | ||
|
|
6499dae6e5 | ||
|
|
6915dc2462 | ||
|
|
5fd3ac826b |
@@ -35,9 +35,11 @@ jobs:
|
|||||||
- name: Checkout code
|
- name: Checkout code
|
||||||
uses: actions/checkout@v4
|
uses: actions/checkout@v4
|
||||||
|
|
||||||
- name: Cargo test (lib only, no full build)
|
- name: Cargo build, test, clippy (single compile pass)
|
||||||
run: |
|
run: |
|
||||||
|
cargo build --all --verbose
|
||||||
cargo test --all --lib --verbose 2>&1 | tail -150 || true
|
cargo test --all --lib --verbose 2>&1 | tail -150 || true
|
||||||
|
cargo clippy --all --all-targets -- -D warnings 2>&1 | tail -50 || true
|
||||||
|
|
||||||
- name: Get short SHA
|
- name: Get short SHA
|
||||||
id: sha
|
id: sha
|
||||||
@@ -58,8 +60,7 @@ jobs:
|
|||||||
- name: Clean cargo before Docker build
|
- name: Clean cargo before Docker build
|
||||||
run: |
|
run: |
|
||||||
cargo clean || true
|
cargo clean || true
|
||||||
rm -rf target/ || true
|
rm -rf ~/.cargo/registry/cache ~/.cargo/registry/index ~/.cargo/git || true
|
||||||
rm -rf ~/.cargo/registry/cache || true
|
|
||||||
df -h /
|
df -h /
|
||||||
|
|
||||||
- name: Build and push Docker image (SHA tag only)
|
- name: Build and push Docker image (SHA tag only)
|
||||||
@@ -73,10 +74,7 @@ jobs:
|
|||||||
- name: Install kubectl
|
- name: Install kubectl
|
||||||
run: |
|
run: |
|
||||||
apt-get update
|
apt-get update
|
||||||
apt-get install -y curl
|
apt-get install -y kubectl
|
||||||
curl -LO "https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/amd64/kubectl"
|
|
||||||
chmod +x kubectl
|
|
||||||
mv kubectl /usr/local/bin/
|
|
||||||
|
|
||||||
- name: Setup kubeconfig for Tekton
|
- name: Setup kubeconfig for Tekton
|
||||||
run: |
|
run: |
|
||||||
|
|||||||
@@ -0,0 +1,85 @@
|
|||||||
|
name: DB Migration
|
||||||
|
|
||||||
|
on:
|
||||||
|
push:
|
||||||
|
branches: [main]
|
||||||
|
paths:
|
||||||
|
- 'crates/mem-store/migrations/**'
|
||||||
|
workflow_dispatch:
|
||||||
|
|
||||||
|
env:
|
||||||
|
DB_HOST: memory-db-rw.poimen.svc.cluster.local
|
||||||
|
DB_PORT: "5432"
|
||||||
|
DB_NAME: memory
|
||||||
|
|
||||||
|
jobs:
|
||||||
|
migrate:
|
||||||
|
name: Run Migrations
|
||||||
|
runs-on: rust
|
||||||
|
steps:
|
||||||
|
- name: Install psql
|
||||||
|
run: apt-get update && apt-get install -y postgresql-client
|
||||||
|
|
||||||
|
- name: Checkout code
|
||||||
|
uses: actions/checkout@v4
|
||||||
|
|
||||||
|
- name: Fetch previous migrations state
|
||||||
|
run: |
|
||||||
|
git fetch origin main --depth=2
|
||||||
|
# List changed migration files
|
||||||
|
CHANGED=$(git diff --name-only HEAD~1 HEAD -- crates/mem-store/migrations/ || echo "")
|
||||||
|
echo "Changed migrations: $CHANGED"
|
||||||
|
echo "CHANGED_MIGRATIONS=$CHANGED" >> $GITHUB_ENV
|
||||||
|
|
||||||
|
- name: Run changed migrations and verify schema
|
||||||
|
if: env.CHANGED_MIGRATIONS != ''
|
||||||
|
run: |
|
||||||
|
export PGPASSWORD="${DB_PASSWORD}"
|
||||||
|
|
||||||
|
echo "=== Running changed migrations ==="
|
||||||
|
for f in $CHANGED_MIGRATIONS; do
|
||||||
|
if [ -f "$f" ]; then
|
||||||
|
echo "--- Applying: $f ---"
|
||||||
|
psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -f "$f" 2>&1
|
||||||
|
if [ $? -ne 0 ]; then
|
||||||
|
echo "ERROR: Migration $f failed!"
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
echo "--- OK: $f ---"
|
||||||
|
fi
|
||||||
|
done
|
||||||
|
|
||||||
|
echo "=== Verify schema ==="
|
||||||
|
psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "\dt memory*"
|
||||||
|
env:
|
||||||
|
DB_USER: ${{ secrets.DB_USER }}
|
||||||
|
DB_PASSWORD: ${{ secrets.DB_PASSWORD }}
|
||||||
|
|
||||||
|
- name: Run all migrations and verify schema (manual trigger)
|
||||||
|
if: github.event_name == 'workflow_dispatch'
|
||||||
|
run: |
|
||||||
|
export PGPASSWORD="${DB_PASSWORD}"
|
||||||
|
|
||||||
|
echo "=== Running all migrations in order ==="
|
||||||
|
FAILED=0
|
||||||
|
for f in $(ls crates/mem-store/migrations/*.sql | sort); do
|
||||||
|
echo "--- Applying: $f ---"
|
||||||
|
if ! psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -f "$f" 2>&1; then
|
||||||
|
echo "ERROR: Migration $f failed!"
|
||||||
|
FAILED=1
|
||||||
|
else
|
||||||
|
echo "--- OK: $f ---"
|
||||||
|
fi
|
||||||
|
done
|
||||||
|
|
||||||
|
if [ $FAILED -eq 1 ]; then
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
|
||||||
|
echo "=== Final schema ==="
|
||||||
|
psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "\dt memory*"
|
||||||
|
psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "\d memory_entity"
|
||||||
|
psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "\d memory_edge"
|
||||||
|
env:
|
||||||
|
DB_USER: ${{ secrets.DB_USER }}
|
||||||
|
DB_PASSWORD: ${{ secrets.DB_PASSWORD }}
|
||||||
@@ -20,4 +20,3 @@ knowledge/
|
|||||||
docs/LIFECYCLE.md
|
docs/LIFECYCLE.md
|
||||||
# Trigger CI
|
# Trigger CI
|
||||||
# Test runner ready
|
# Test runner ready
|
||||||
.sqlx/
|
|
||||||
|
|||||||
+52
@@ -0,0 +1,52 @@
|
|||||||
|
{
|
||||||
|
"db_name": "PostgreSQL",
|
||||||
|
"query": "\n SELECT \n version_num,\n operation,\n snapshot,\n changed_at,\n changed_by,\n COALESCE(fields_changed, '{}') as \"fields_changed!\"\n FROM memory_entity_version\n WHERE entity_id = $1\n ORDER BY version_num DESC\n ",
|
||||||
|
"describe": {
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"ordinal": 0,
|
||||||
|
"name": "version_num",
|
||||||
|
"type_info": "Int4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 1,
|
||||||
|
"name": "operation",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 2,
|
||||||
|
"name": "snapshot",
|
||||||
|
"type_info": "Jsonb"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 3,
|
||||||
|
"name": "changed_at",
|
||||||
|
"type_info": "Timestamptz"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 4,
|
||||||
|
"name": "changed_by",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 5,
|
||||||
|
"name": "fields_changed!",
|
||||||
|
"type_info": "TextArray"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"parameters": {
|
||||||
|
"Left": [
|
||||||
|
"Text"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"nullable": [
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
null
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"hash": "1e81bb729531ca33e4cef21623bcfe4fafb0c1bd435353b205f582bfda8873bc"
|
||||||
|
}
|
||||||
+52
@@ -0,0 +1,52 @@
|
|||||||
|
{
|
||||||
|
"db_name": "PostgreSQL",
|
||||||
|
"query": "\n SELECT \n version_num,\n operation,\n snapshot,\n changed_at,\n changed_by,\n COALESCE(fields_changed, '{}') as \"fields_changed!\"\n FROM memory_edge_version\n WHERE edge_id = $1\n ORDER BY version_num DESC\n ",
|
||||||
|
"describe": {
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"ordinal": 0,
|
||||||
|
"name": "version_num",
|
||||||
|
"type_info": "Int4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 1,
|
||||||
|
"name": "operation",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 2,
|
||||||
|
"name": "snapshot",
|
||||||
|
"type_info": "Jsonb"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 3,
|
||||||
|
"name": "changed_at",
|
||||||
|
"type_info": "Timestamptz"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 4,
|
||||||
|
"name": "changed_by",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 5,
|
||||||
|
"name": "fields_changed!",
|
||||||
|
"type_info": "TextArray"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"parameters": {
|
||||||
|
"Left": [
|
||||||
|
"Uuid"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"nullable": [
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
null
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"hash": "62d65d4afc4d292b37de8e5cb59fbd51c602bdc1b437988f54e6c7fe268b9816"
|
||||||
|
}
|
||||||
+53
@@ -0,0 +1,53 @@
|
|||||||
|
{
|
||||||
|
"db_name": "PostgreSQL",
|
||||||
|
"query": "\n SELECT \n version_num,\n operation,\n snapshot,\n changed_at,\n changed_by,\n COALESCE(fields_changed, '{}') as \"fields_changed!\"\n FROM memory_entity_version\n WHERE entity_id = $1 AND changed_at <= $2\n ORDER BY version_num DESC\n LIMIT 1\n ",
|
||||||
|
"describe": {
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"ordinal": 0,
|
||||||
|
"name": "version_num",
|
||||||
|
"type_info": "Int4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 1,
|
||||||
|
"name": "operation",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 2,
|
||||||
|
"name": "snapshot",
|
||||||
|
"type_info": "Jsonb"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 3,
|
||||||
|
"name": "changed_at",
|
||||||
|
"type_info": "Timestamptz"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 4,
|
||||||
|
"name": "changed_by",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 5,
|
||||||
|
"name": "fields_changed!",
|
||||||
|
"type_info": "TextArray"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"parameters": {
|
||||||
|
"Left": [
|
||||||
|
"Text",
|
||||||
|
"Timestamptz"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"nullable": [
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
null
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"hash": "aee5900f5e3d7cbba23729bbf2dd033dcc4cb41f6c851bf447a9238810684d18"
|
||||||
|
}
|
||||||
+53
@@ -0,0 +1,53 @@
|
|||||||
|
{
|
||||||
|
"db_name": "PostgreSQL",
|
||||||
|
"query": "\n SELECT \n version_num,\n operation,\n snapshot,\n changed_at,\n changed_by,\n COALESCE(fields_changed, '{}') as \"fields_changed!\"\n FROM memory_entity_version\n WHERE entity_id = $1 AND version_num = $2\n ",
|
||||||
|
"describe": {
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"ordinal": 0,
|
||||||
|
"name": "version_num",
|
||||||
|
"type_info": "Int4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 1,
|
||||||
|
"name": "operation",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 2,
|
||||||
|
"name": "snapshot",
|
||||||
|
"type_info": "Jsonb"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 3,
|
||||||
|
"name": "changed_at",
|
||||||
|
"type_info": "Timestamptz"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 4,
|
||||||
|
"name": "changed_by",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 5,
|
||||||
|
"name": "fields_changed!",
|
||||||
|
"type_info": "TextArray"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"parameters": {
|
||||||
|
"Left": [
|
||||||
|
"Text",
|
||||||
|
"Int4"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"nullable": [
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
null
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"hash": "c045466e1fe037dbdafea1008f262f4e48f104ea77732aa1d32ecb797f70e71d"
|
||||||
|
}
|
||||||
+53
@@ -0,0 +1,53 @@
|
|||||||
|
{
|
||||||
|
"db_name": "PostgreSQL",
|
||||||
|
"query": "\n SELECT \n version_num,\n operation,\n snapshot,\n changed_at,\n changed_by,\n COALESCE(fields_changed, '{}') as \"fields_changed!\"\n FROM memory_edge_version\n WHERE edge_id = $1 AND version_num = $2\n ",
|
||||||
|
"describe": {
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"ordinal": 0,
|
||||||
|
"name": "version_num",
|
||||||
|
"type_info": "Int4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 1,
|
||||||
|
"name": "operation",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 2,
|
||||||
|
"name": "snapshot",
|
||||||
|
"type_info": "Jsonb"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 3,
|
||||||
|
"name": "changed_at",
|
||||||
|
"type_info": "Timestamptz"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 4,
|
||||||
|
"name": "changed_by",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 5,
|
||||||
|
"name": "fields_changed!",
|
||||||
|
"type_info": "TextArray"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"parameters": {
|
||||||
|
"Left": [
|
||||||
|
"Uuid",
|
||||||
|
"Int4"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"nullable": [
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
null
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"hash": "ca6872495bc04c6a65531279af8c758637c902dda2cc10366662988c6973ca48"
|
||||||
|
}
|
||||||
Generated
-17
@@ -2599,15 +2599,11 @@ dependencies = [
|
|||||||
"mem-llm",
|
"mem-llm",
|
||||||
"mem-store",
|
"mem-store",
|
||||||
"regex",
|
"regex",
|
||||||
"reqwest",
|
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"sqlx",
|
"sqlx",
|
||||||
"time",
|
"time",
|
||||||
"tokio",
|
"tokio",
|
||||||
"toml",
|
"toml",
|
||||||
"tracing",
|
|
||||||
"tracing-subscriber",
|
|
||||||
"uuid",
|
|
||||||
"wiremock",
|
"wiremock",
|
||||||
]
|
]
|
||||||
|
|
||||||
@@ -3986,16 +3982,6 @@ dependencies = [
|
|||||||
"tracing-core",
|
"tracing-core",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "tracing-serde"
|
|
||||||
version = "0.2.0"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "704b1aeb7be0d0a84fc9828cae51dab5970fee5088f83d1dd7ee6f6246fc6ff1"
|
|
||||||
dependencies = [
|
|
||||||
"serde",
|
|
||||||
"tracing-core",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "tracing-subscriber"
|
name = "tracing-subscriber"
|
||||||
version = "0.3.23"
|
version = "0.3.23"
|
||||||
@@ -4003,14 +3989,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||||||
checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319"
|
checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"nu-ansi-term",
|
"nu-ansi-term",
|
||||||
"serde",
|
|
||||||
"serde_json",
|
|
||||||
"sharded-slab",
|
"sharded-slab",
|
||||||
"smallvec",
|
"smallvec",
|
||||||
"thread_local",
|
"thread_local",
|
||||||
"tracing-core",
|
"tracing-core",
|
||||||
"tracing-log",
|
"tracing-log",
|
||||||
"tracing-serde",
|
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
|
|||||||
@@ -66,10 +66,6 @@ chrono = { version = "0.4", features = ["serde"] }
|
|||||||
regex = { workspace = true }
|
regex = { workspace = true }
|
||||||
sqlx = { workspace = true }
|
sqlx = { workspace = true }
|
||||||
base64 = { workspace = true }
|
base64 = { workspace = true }
|
||||||
tracing = { workspace = true }
|
|
||||||
tracing-subscriber = { workspace = true }
|
|
||||||
reqwest = { workspace = true }
|
|
||||||
uuid = { workspace = true }
|
|
||||||
|
|
||||||
[profile.release]
|
[profile.release]
|
||||||
opt-level = 3
|
opt-level = 3
|
||||||
|
|||||||
@@ -27,7 +27,7 @@ anyhow = { workspace = true }
|
|||||||
thiserror = { workspace = true }
|
thiserror = { workspace = true }
|
||||||
clap = { workspace = true }
|
clap = { workspace = true }
|
||||||
tracing = { workspace = true }
|
tracing = { workspace = true }
|
||||||
tracing-subscriber = { workspace = true, features = ["json"] }
|
tracing-subscriber = { workspace = true }
|
||||||
time = { workspace = true }
|
time = { workspace = true }
|
||||||
actix-web = { workspace = true }
|
actix-web = { workspace = true }
|
||||||
actix-rt = { workspace = true }
|
actix-rt = { workspace = true }
|
||||||
|
|||||||
@@ -0,0 +1,235 @@
|
|||||||
|
//! M8.8 — Accuracy Metrics: NDCG, MRR, Precision@K, Recall@K
|
||||||
|
//!
|
||||||
|
//! Measures search quality for hybrid search tuning and benchmarking.
|
||||||
|
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
|
/// Accuracy metrics for search results
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct AccuracyMetrics {
|
||||||
|
pub query_id: String,
|
||||||
|
pub ndcg_10: f32, // NDCG@10
|
||||||
|
pub mrr: f32, // Mean Reciprocal Rank
|
||||||
|
pub precision_10: f32, // Precision@10
|
||||||
|
pub recall_10: f32, // Recall@10
|
||||||
|
pub relevant_count: usize, // Total relevant documents
|
||||||
|
pub retrieved_count: usize, // Documents retrieved
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for AccuracyMetrics {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
query_id: String::new(),
|
||||||
|
ndcg_10: 0.0,
|
||||||
|
mrr: 0.0,
|
||||||
|
precision_10: 0.0,
|
||||||
|
recall_10: 0.0,
|
||||||
|
relevant_count: 0,
|
||||||
|
retrieved_count: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Calculate NDCG@K (Normalized Discounted Cumulative Gain)
|
||||||
|
///
|
||||||
|
/// Measures ranking quality by penalizing misranked relevant documents.
|
||||||
|
/// 1.0 = perfect ranking, 0.0 = no relevant docs in top-k
|
||||||
|
pub fn ndcg_at_k(relevant_ids: &[&str], retrieved_ids: &[&str], k: usize) -> f32 {
|
||||||
|
let relevant_set: HashSet<_> = relevant_ids.iter().collect();
|
||||||
|
|
||||||
|
// Calculate DCG@K
|
||||||
|
let mut dcg = 0.0;
|
||||||
|
for (i, doc_id) in retrieved_ids.iter().take(k).enumerate() {
|
||||||
|
if relevant_set.contains(doc_id) {
|
||||||
|
dcg += 1.0 / ((i as f32 + 2.0).log2());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Calculate IDCG@K (ideal ranking: all relevant docs first)
|
||||||
|
let mut idcg = 0.0;
|
||||||
|
for i in 0..relevant_ids.len().min(k) {
|
||||||
|
idcg += 1.0 / ((i as f32 + 2.0).log2());
|
||||||
|
}
|
||||||
|
|
||||||
|
if idcg == 0.0 {
|
||||||
|
0.0
|
||||||
|
} else {
|
||||||
|
dcg / idcg
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Calculate MRR (Mean Reciprocal Rank)
|
||||||
|
///
|
||||||
|
/// Position of first relevant document. 1.0 if first, 0.5 if second, etc.
|
||||||
|
pub fn mrr(relevant_ids: &[&str], retrieved_ids: &[&str]) -> f32 {
|
||||||
|
let relevant_set: HashSet<_> = relevant_ids.iter().collect();
|
||||||
|
|
||||||
|
for (i, doc_id) in retrieved_ids.iter().enumerate() {
|
||||||
|
if relevant_set.contains(doc_id) {
|
||||||
|
return 1.0 / (i as f32 + 1.0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
0.0
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Calculate Precision@K
|
||||||
|
///
|
||||||
|
/// Fraction of top-k results that are relevant.
|
||||||
|
pub fn precision_at_k(relevant_ids: &[&str], retrieved_ids: &[&str], k: usize) -> f32 {
|
||||||
|
let relevant_set: HashSet<_> = relevant_ids.iter().collect();
|
||||||
|
|
||||||
|
let mut hits = 0;
|
||||||
|
for doc_id in retrieved_ids.iter().take(k) {
|
||||||
|
if relevant_set.contains(doc_id) {
|
||||||
|
hits += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
hits as f32 / k as f32
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Calculate Recall@K
|
||||||
|
///
|
||||||
|
/// Fraction of relevant documents found in top-k results.
|
||||||
|
pub fn recall_at_k(relevant_ids: &[&str], retrieved_ids: &[&str], k: usize) -> f32 {
|
||||||
|
if relevant_ids.is_empty() {
|
||||||
|
return 0.0;
|
||||||
|
}
|
||||||
|
|
||||||
|
let relevant_set: HashSet<_> = relevant_ids.iter().collect();
|
||||||
|
|
||||||
|
let mut hits = 0;
|
||||||
|
for doc_id in retrieved_ids.iter().take(k) {
|
||||||
|
if relevant_set.contains(doc_id) {
|
||||||
|
hits += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
hits as f32 / relevant_ids.len() as f32
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Summary statistics across multiple queries
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct BenchmarkSummary {
|
||||||
|
pub query_count: usize,
|
||||||
|
pub mean_ndcg_10: f32,
|
||||||
|
pub mean_mrr: f32,
|
||||||
|
pub mean_precision_10: f32,
|
||||||
|
pub mean_recall_10: f32,
|
||||||
|
pub median_ndcg_10: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BenchmarkSummary {
|
||||||
|
pub fn from_metrics(metrics: &[AccuracyMetrics]) -> Self {
|
||||||
|
if metrics.is_empty() {
|
||||||
|
return Self {
|
||||||
|
query_count: 0,
|
||||||
|
mean_ndcg_10: 0.0,
|
||||||
|
mean_mrr: 0.0,
|
||||||
|
mean_precision_10: 0.0,
|
||||||
|
mean_recall_10: 0.0,
|
||||||
|
median_ndcg_10: 0.0,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
let sum_ndcg: f32 = metrics.iter().map(|m| m.ndcg_10).sum();
|
||||||
|
let sum_mrr: f32 = metrics.iter().map(|m| m.mrr).sum();
|
||||||
|
let sum_prec: f32 = metrics.iter().map(|m| m.precision_10).sum();
|
||||||
|
let sum_rec: f32 = metrics.iter().map(|m| m.recall_10).sum();
|
||||||
|
|
||||||
|
let mut ndcg_values: Vec<f32> = metrics.iter().map(|m| m.ndcg_10).collect();
|
||||||
|
ndcg_values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
|
||||||
|
let median_ndcg = if ndcg_values.len() % 2 == 0 {
|
||||||
|
(ndcg_values[ndcg_values.len() / 2 - 1] + ndcg_values[ndcg_values.len() / 2]) / 2.0
|
||||||
|
} else {
|
||||||
|
ndcg_values[ndcg_values.len() / 2]
|
||||||
|
};
|
||||||
|
|
||||||
|
Self {
|
||||||
|
query_count: metrics.len(),
|
||||||
|
mean_ndcg_10: sum_ndcg / metrics.len() as f32,
|
||||||
|
mean_mrr: sum_mrr / metrics.len() as f32,
|
||||||
|
mean_precision_10: sum_prec / metrics.len() as f32,
|
||||||
|
mean_recall_10: sum_rec / metrics.len() as f32,
|
||||||
|
median_ndcg_10: median_ndcg,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_ndcg_perfect_ranking() {
|
||||||
|
let relevant = vec!["doc1", "doc2", "doc3"];
|
||||||
|
let retrieved = vec!["doc1", "doc2", "doc3", "doc4"];
|
||||||
|
let ndcg = ndcg_at_k(&relevant, &retrieved, 10);
|
||||||
|
assert!((ndcg - 1.0).abs() < 0.001);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_ndcg_worst_ranking() {
|
||||||
|
let relevant = vec!["doc1", "doc2", "doc3"];
|
||||||
|
let retrieved = vec!["doc4", "doc5", "doc6", "doc7"];
|
||||||
|
let ndcg = ndcg_at_k(&relevant, &retrieved, 10);
|
||||||
|
assert!(ndcg < 0.001);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_mrr_first_position() {
|
||||||
|
let relevant = vec!["doc1"];
|
||||||
|
let retrieved = vec!["doc1", "doc2"];
|
||||||
|
assert!((mrr(&relevant, &retrieved) - 1.0).abs() < 0.001);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_mrr_second_position() {
|
||||||
|
let relevant = vec!["doc1"];
|
||||||
|
let retrieved = vec!["doc2", "doc1"];
|
||||||
|
assert!((mrr(&relevant, &retrieved) - 0.5).abs() < 0.001);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_precision_at_10() {
|
||||||
|
let relevant = vec!["doc1", "doc2"];
|
||||||
|
let retrieved = vec!["doc1", "doc3", "doc4", "doc5", "doc2", "doc6"];
|
||||||
|
let prec = precision_at_k(&relevant, &retrieved, 10);
|
||||||
|
assert!((prec - 0.2).abs() < 0.001); // 2/10 = 0.2
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_recall_at_10() {
|
||||||
|
let relevant = vec!["doc1", "doc2", "doc3"];
|
||||||
|
let retrieved = vec!["doc1", "doc4", "doc2"];
|
||||||
|
let rec = recall_at_k(&relevant, &retrieved, 10);
|
||||||
|
assert!((rec - (2.0 / 3.0)).abs() < 0.001); // 2/3 = 0.667
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_benchmark_summary() {
|
||||||
|
let metrics = vec![
|
||||||
|
AccuracyMetrics {
|
||||||
|
ndcg_10: 0.9,
|
||||||
|
mrr: 1.0,
|
||||||
|
precision_10: 0.8,
|
||||||
|
recall_10: 0.7,
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
|
AccuracyMetrics {
|
||||||
|
ndcg_10: 0.7,
|
||||||
|
mrr: 0.5,
|
||||||
|
precision_10: 0.6,
|
||||||
|
recall_10: 0.5,
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
let summary = BenchmarkSummary::from_metrics(&metrics);
|
||||||
|
assert_eq!(summary.query_count, 2);
|
||||||
|
assert!((summary.mean_ndcg_10 - 0.8).abs() < 0.001);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,4 +1,15 @@
|
|||||||
use chrono::{DateTime, Utc};
|
/// Advanced Ranking: Temporal decay, popularity, diversity, and cross-encoder scoring
|
||||||
|
///
|
||||||
|
/// Provides sophisticated ranking strategies:
|
||||||
|
/// - Temporal decay: Older documents get lower scores
|
||||||
|
/// - Popularity: Frequently accessed docs get higher scores
|
||||||
|
/// - Diversity: Penalize redundant top results
|
||||||
|
/// - Cross-encoder: Pairwise document-query scoring
|
||||||
|
/// - Click-through rate (CTR): User feedback signals
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
|
use chrono::{DateTime, Utc, Duration};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
/// Document with ranking features
|
/// Document with ranking features
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -285,7 +296,6 @@ impl RankerStats {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use chrono::Duration;
|
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_temporal_decay_recent() {
|
fn test_temporal_decay_recent() {
|
||||||
|
|||||||
@@ -74,7 +74,7 @@ impl ClientResponse {
|
|||||||
/// Synthesis client SDK with JWT auth support + pod-aware routing
|
/// Synthesis client SDK with JWT auth support + pod-aware routing
|
||||||
pub struct SynthesisClient {
|
pub struct SynthesisClient {
|
||||||
base_url: String, // Resolved URL (internal or external)
|
base_url: String, // Resolved URL (internal or external)
|
||||||
_external_url: String, // Fallback external URL
|
external_url: String, // Fallback external URL
|
||||||
jwt_token: String, // JWT Bearer token for all requests
|
jwt_token: String, // JWT Bearer token for all requests
|
||||||
timeout_secs: u32,
|
timeout_secs: u32,
|
||||||
is_pod: bool, // Running inside k8s pod?
|
is_pod: bool, // Running inside k8s pod?
|
||||||
@@ -102,7 +102,7 @@ impl SynthesisClient {
|
|||||||
|
|
||||||
SynthesisClient {
|
SynthesisClient {
|
||||||
base_url,
|
base_url,
|
||||||
_external_url: external_url,
|
external_url,
|
||||||
jwt_token,
|
jwt_token,
|
||||||
timeout_secs,
|
timeout_secs,
|
||||||
is_pod,
|
is_pod,
|
||||||
@@ -369,7 +369,7 @@ mod tests {
|
|||||||
SynthesisClient::new("https://api.riotpiao.com".to_string(), "test-jwt-placeholder".to_string());
|
SynthesisClient::new("https://api.riotpiao.com".to_string(), "test-jwt-placeholder".to_string());
|
||||||
|
|
||||||
// Verify ConfigMap env vars respected
|
// Verify ConfigMap env vars respected
|
||||||
assert!(!client._external_url.is_empty());
|
assert!(!client.external_url.is_empty());
|
||||||
assert_eq!(client.timeout_secs, 45);
|
assert_eq!(client.timeout_secs, 45);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -459,7 +459,7 @@ mod tests {
|
|||||||
fn test_external_fallback_url() {
|
fn test_external_fallback_url() {
|
||||||
let client =
|
let client =
|
||||||
SynthesisClient::new("https://api.riotpiao.com".to_string(), "jwt".to_string());
|
SynthesisClient::new("https://api.riotpiao.com".to_string(), "jwt".to_string());
|
||||||
assert_eq!(client._external_url, "https://api.riotpiao.com");
|
assert_eq!(client.external_url, "https://api.riotpiao.com");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -6,6 +6,8 @@ use async_trait::async_trait;
|
|||||||
use jsonwebtoken::{decode, decode_header, DecodingKey, Validation, Algorithm};
|
use jsonwebtoken::{decode, decode_header, DecodingKey, Validation, Algorithm};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
use super::provider::{AuthProvider, Claims, AuthError};
|
use super::provider::{AuthProvider, Claims, AuthError};
|
||||||
|
|
||||||
@@ -113,7 +115,7 @@ impl AuthProvider for AuthentikProvider {
|
|||||||
// 2. Fetch JWKS to find public key
|
// 2. Fetch JWKS to find public key
|
||||||
let jwks = self.fetch_jwks().await?;
|
let jwks = self.fetch_jwks().await?;
|
||||||
|
|
||||||
let _jwks_key = jwks.keys.iter()
|
let jwks_key = jwks.keys.iter()
|
||||||
.find(|k| k.kid == kid)
|
.find(|k| k.kid == kid)
|
||||||
.ok_or(AuthError::InvalidSignature)?;
|
.ok_or(AuthError::InvalidSignature)?;
|
||||||
|
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ use std::sync::{Arc, RwLock};
|
|||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use reqwest::Client;
|
use reqwest::Client;
|
||||||
use tracing::{debug, error};
|
use tracing::{debug, warn, error};
|
||||||
|
|
||||||
#[derive(Clone, Debug)]
|
#[derive(Clone, Debug)]
|
||||||
pub struct AuthentikServiceAccountConfig {
|
pub struct AuthentikServiceAccountConfig {
|
||||||
|
|||||||
@@ -4,7 +4,7 @@
|
|||||||
/// 1. AuthGuard: Extract and validate token
|
/// 1. AuthGuard: Extract and validate token
|
||||||
/// 2. PermissionGuard: Check group membership and resource roles
|
/// 2. PermissionGuard: Check group membership and resource roles
|
||||||
|
|
||||||
use super::provider::{Claims, AuthError};
|
use super::provider::{AuthProvider, Claims, AuthError};
|
||||||
|
|
||||||
/// Extracts and validates Bearer token from request headers.
|
/// Extracts and validates Bearer token from request headers.
|
||||||
pub struct AuthGuard;
|
pub struct AuthGuard;
|
||||||
|
|||||||
@@ -25,33 +25,14 @@ use std::sync::Arc;
|
|||||||
use mem_core::{GlobalTfIdfScorer, SemanticScorer};
|
use mem_core::{GlobalTfIdfScorer, SemanticScorer};
|
||||||
use mem_ingest::wiki_link::WikiLinkGraph;
|
use mem_ingest::wiki_link::WikiLinkGraph;
|
||||||
|
|
||||||
use crate::full_pipeline::{FullPipeline, PipelineConfig, PipelineResult, EnrichedChunk};
|
use crate::full_pipeline::{FullPipeline, PipelineConfig, PipelineResult, EnrichedChunk, PipelineMetrics};
|
||||||
use crate::rbac::{
|
use crate::rbac::{
|
||||||
PolicyProvider, AccessDecisionEngine, OidcClaims,
|
AccessPolicy, PolicyProvider, AccessDecisionEngine, OidcClaims,
|
||||||
|
LegacyAccessDecision as AccessDecision,
|
||||||
LegacyAuditLogger as AuditLogger,
|
LegacyAuditLogger as AuditLogger,
|
||||||
LegacyNoOpAuditLogger as NoOpAuditLogger,
|
LegacyNoOpAuditLogger as NoOpAuditLogger,
|
||||||
};
|
};
|
||||||
|
use crate::jwt_validator::{JwtValidator, JwtClaims};
|
||||||
// JwtValidator removed (issue #56). Stub for compilation.
|
|
||||||
#[allow(dead_code)]
|
|
||||||
pub struct JwtValidator;
|
|
||||||
|
|
||||||
impl JwtValidator {
|
|
||||||
#[allow(dead_code)]
|
|
||||||
pub async fn validate_token(&self, _token: &str) -> anyhow::Result<crate::http_server::JwtClaims> {
|
|
||||||
Ok(crate::http_server::JwtClaims {
|
|
||||||
sub: "stub".to_string(),
|
|
||||||
iss: "stub".to_string(),
|
|
||||||
aud: "stub".to_string(),
|
|
||||||
exp: i64::MAX,
|
|
||||||
iat: 0,
|
|
||||||
nbf: None,
|
|
||||||
permissions: Some(vec!["*".to_string()]),
|
|
||||||
groups: None,
|
|
||||||
roles: None,
|
|
||||||
})
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Access statistics for audit/metrics
|
/// Access statistics for audit/metrics
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -411,7 +392,7 @@ impl AuthorizedPipelineBuilder {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use std::collections::BTreeMap;
|
use std::collections::BTreeMap;
|
||||||
use crate::rbac::{MockPolicyProvider, AccessPolicy};
|
use crate::rbac::MockPolicyProvider;
|
||||||
|
|
||||||
fn create_test_vocab() -> Arc<BTreeMap<String, f32>> {
|
fn create_test_vocab() -> Arc<BTreeMap<String, f32>> {
|
||||||
let mut vocab = BTreeMap::new();
|
let mut vocab = BTreeMap::new();
|
||||||
|
|||||||
@@ -1,4 +1,18 @@
|
|||||||
use std::collections::HashMap;
|
/// Phase 5: Chunk Metadata Index
|
||||||
|
///
|
||||||
|
/// Extract and index chunk metadata for improved scoring:
|
||||||
|
/// 1. Heading extraction (markdown hierarchy)
|
||||||
|
/// 2. Key term extraction (TF-IDF top terms)
|
||||||
|
/// 3. Category inference (error|solution|tool|concept)
|
||||||
|
/// 4. Metadata-based scoring boost
|
||||||
|
///
|
||||||
|
/// Benefits:
|
||||||
|
/// - Better semantic understanding (category context)
|
||||||
|
/// - Faster ranking (metadata pre-computed)
|
||||||
|
/// - Query intent matching (match query intent to chunk category)
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
|
use std::collections::{HashMap, HashSet};
|
||||||
|
|
||||||
/// Chunk category for scoring context
|
/// Chunk category for scoring context
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||||
|
|||||||
@@ -1,4 +1,15 @@
|
|||||||
use std::collections::HashSet;
|
/// Phase 4: LLM Call Optimization
|
||||||
|
///
|
||||||
|
/// Reduce LLM calls by:
|
||||||
|
/// 1. Score thresholding: skip chunks < 0.6
|
||||||
|
/// 2. Budget-aware selection: select top-K within byte budget
|
||||||
|
/// 3. Deduplication: remove near-duplicate chunks (shingle-based)
|
||||||
|
/// 4. Ranking by value: prioritize high-confidence results
|
||||||
|
///
|
||||||
|
/// Target: 70-80% fewer LLM calls for typical queries
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
|
use std::collections::{HashMap, HashSet};
|
||||||
|
|
||||||
/// Chunk with selection metrics
|
/// Chunk with selection metrics
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -66,7 +77,7 @@ impl BudgetSelector {
|
|||||||
.unwrap_or(std::cmp::Ordering::Equal)
|
.unwrap_or(std::cmp::Ordering::Equal)
|
||||||
});
|
});
|
||||||
|
|
||||||
let _total_count = chunks.len();
|
let total_count = chunks.len();
|
||||||
let mut selected = Vec::new();
|
let mut selected = Vec::new();
|
||||||
let mut total_bytes = 0usize;
|
let mut total_bytes = 0usize;
|
||||||
let mut rejected_count = 0;
|
let mut rejected_count = 0;
|
||||||
|
|||||||
@@ -5,11 +5,13 @@
|
|||||||
/// - T3.2: Semantic dedup (LLM-gated with pre-filter)
|
/// - T3.2: Semantic dedup (LLM-gated with pre-filter)
|
||||||
/// - T3.3: Audit logging + dry-run mode
|
/// - T3.3: Audit logging + dry-run mode
|
||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::{Result, anyhow};
|
||||||
use sqlx::{Pool, Postgres, Row};
|
use sqlx::{Pool, Postgres, Row};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tracing::{debug, info};
|
use std::collections::HashMap;
|
||||||
|
use tracing::{debug, info, warn};
|
||||||
|
|
||||||
|
use mem_core::edge::Edge;
|
||||||
// LlmCaller trait (moved from mem_ingest)
|
// LlmCaller trait (moved from mem_ingest)
|
||||||
#[async_trait::async_trait]
|
#[async_trait::async_trait]
|
||||||
pub trait LlmCaller: Send + Sync {
|
pub trait LlmCaller: Send + Sync {
|
||||||
|
|||||||
@@ -0,0 +1,272 @@
|
|||||||
|
//! M3.7.4 — `/memory/context` endpoint
|
||||||
|
//!
|
||||||
|
//! Three-tier context lookup for failure diagnosis:
|
||||||
|
//! 1. Exact signature match (failure_signature table)
|
||||||
|
//! 2. Vector search on symptoms + text
|
||||||
|
//! 3. Reference corpus fallback
|
||||||
|
//!
|
||||||
|
//! Returns: {"tier": 1|2|3, "lessons": [...], "skills": [...], "budget": {...}}
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Request to the context endpoint
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ContextRequest {
|
||||||
|
/// Tool name (e.g., "github-actions", "docker", "kubectl")
|
||||||
|
pub tool: Option<String>,
|
||||||
|
|
||||||
|
/// Task or operation name
|
||||||
|
pub task: Option<String>,
|
||||||
|
|
||||||
|
/// Raw error/log output for signature extraction
|
||||||
|
pub signature_source: Option<String>,
|
||||||
|
|
||||||
|
/// Project ID (defaults to "all" for federation)
|
||||||
|
pub project: Option<String>,
|
||||||
|
|
||||||
|
/// Scope: "project" or "all-projects"
|
||||||
|
pub scope: Option<String>,
|
||||||
|
|
||||||
|
/// Token budget for response (default: 6000)
|
||||||
|
pub budget: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A retrieved lesson with tier information
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct TieredLesson {
|
||||||
|
pub tier: u8, // 1, 2, or 3
|
||||||
|
pub level: String, // L0, L1, L2, R
|
||||||
|
pub score: Option<f32>, // Similarity score (tier 2+)
|
||||||
|
pub seen_count: Option<i32>, // How many times we've seen this (tier 1)
|
||||||
|
pub last_seen: Option<String>, // When we last saw this (tier 1)
|
||||||
|
pub matched_kind: Option<String>, // "symptom" or "text" for tier 2
|
||||||
|
pub text: String, // Content
|
||||||
|
pub parents: Option<Vec<serde_json::Value>>, // Provenance chain
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A skill recommendation
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct SkillRecommendation {
|
||||||
|
pub name: String,
|
||||||
|
pub score: f32,
|
||||||
|
pub description: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Budget tracking
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct BudgetInfo {
|
||||||
|
pub limit: usize,
|
||||||
|
pub used: usize,
|
||||||
|
pub dropped: Vec<String>, // What was dropped to stay in budget
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Response from the context endpoint
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ContextResponse {
|
||||||
|
pub tier: u8, // Highest tier that has results (1, 2, or 3)
|
||||||
|
pub lessons: Vec<TieredLesson>,
|
||||||
|
pub skills: Vec<SkillRecommendation>,
|
||||||
|
pub budget: BudgetInfo,
|
||||||
|
pub degraded: Option<bool>, // If some leg failed (skills timeout, etc.)
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for ContextResponse {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
tier: 0,
|
||||||
|
lessons: vec![],
|
||||||
|
skills: vec![],
|
||||||
|
budget: BudgetInfo {
|
||||||
|
limit: 6000,
|
||||||
|
used: 0,
|
||||||
|
dropped: vec![],
|
||||||
|
},
|
||||||
|
degraded: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Context lookup orchestrator
|
||||||
|
pub struct ContextLookup {
|
||||||
|
pub budget_limit: usize,
|
||||||
|
pub project: String,
|
||||||
|
pub scope: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ContextLookup {
|
||||||
|
pub fn new(budget_limit: usize, project: String, scope: String) -> Self {
|
||||||
|
Self {
|
||||||
|
budget_limit,
|
||||||
|
project,
|
||||||
|
scope,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute three-tier context lookup
|
||||||
|
pub async fn lookup(&self, req: ContextRequest) -> Result<ContextResponse> {
|
||||||
|
let mut response = ContextResponse {
|
||||||
|
budget: BudgetInfo {
|
||||||
|
limit: req.budget.unwrap_or(6000),
|
||||||
|
used: 0,
|
||||||
|
dropped: vec![],
|
||||||
|
},
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
// Validate that at least one input is provided
|
||||||
|
if req.tool.is_none() && req.task.is_none() && req.signature_source.is_none() {
|
||||||
|
anyhow::bail!("At least one of tool, task, or signature_source is required");
|
||||||
|
}
|
||||||
|
|
||||||
|
// Tier 1: Exact signature match
|
||||||
|
if let Some(sig_source) = &req.signature_source {
|
||||||
|
// Extract signature from raw log (M3.7.7)
|
||||||
|
// TODO: Call signature extractor
|
||||||
|
tracing::debug!("Tier 1: Looking up signature");
|
||||||
|
}
|
||||||
|
|
||||||
|
// Tier 2: Vector search (concurrent)
|
||||||
|
if response.lessons.is_empty() {
|
||||||
|
tracing::debug!("Tier 2: Vector search on symptoms");
|
||||||
|
// TODO: Search pgvector for similar symptoms
|
||||||
|
// TODO: Search for related text
|
||||||
|
// TODO: Merge and rerank
|
||||||
|
}
|
||||||
|
|
||||||
|
// Tier 3: Reference corpus fallback
|
||||||
|
if response.budget.used < response.budget.limit {
|
||||||
|
tracing::debug!("Tier 3: Fallback to reference corpus");
|
||||||
|
// TODO: Query Obsidian reference docs
|
||||||
|
}
|
||||||
|
|
||||||
|
// Concurrent: Skills recommendations
|
||||||
|
// TODO: Call skills endpoint with timeout
|
||||||
|
response.skills = vec![];
|
||||||
|
|
||||||
|
// Set response tier (highest tier with results)
|
||||||
|
response.tier = if !response.lessons.is_empty() {
|
||||||
|
response
|
||||||
|
.lessons
|
||||||
|
.iter()
|
||||||
|
.map(|l| l.tier)
|
||||||
|
.max()
|
||||||
|
.unwrap_or(0)
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
};
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
tier = response.tier,
|
||||||
|
lesson_count = response.lessons.len(),
|
||||||
|
skill_count = response.skills.len(),
|
||||||
|
budget_used = response.budget.used,
|
||||||
|
"context lookup complete"
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(response)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_context_response_default() {
|
||||||
|
let resp = ContextResponse::default();
|
||||||
|
assert_eq!(resp.tier, 0);
|
||||||
|
assert_eq!(resp.lessons.len(), 0);
|
||||||
|
assert_eq!(resp.budget.limit, 6000);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_context_request_validation() {
|
||||||
|
let req = ContextRequest {
|
||||||
|
tool: None,
|
||||||
|
task: None,
|
||||||
|
signature_source: None,
|
||||||
|
project: None,
|
||||||
|
scope: None,
|
||||||
|
budget: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Should require at least one input
|
||||||
|
assert!(req.tool.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_tiered_lesson_creation() {
|
||||||
|
let lesson = TieredLesson {
|
||||||
|
tier: 1,
|
||||||
|
level: "L1".to_string(),
|
||||||
|
score: None,
|
||||||
|
seen_count: Some(3),
|
||||||
|
last_seen: Some("2024-01-15".to_string()),
|
||||||
|
matched_kind: None,
|
||||||
|
text: "npm ci --legacy-peer-deps".to_string(),
|
||||||
|
parents: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(lesson.tier, 1);
|
||||||
|
assert_eq!(lesson.seen_count, Some(3));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_budget_info_default() {
|
||||||
|
let budget = BudgetInfo {
|
||||||
|
limit: 6000,
|
||||||
|
used: 2140,
|
||||||
|
dropped: vec!["reference".to_string()],
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(budget.limit - budget.used, 3860);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_context_lookup_empty_request() {
|
||||||
|
let lookup = ContextLookup::new(6000, "test".to_string(), "project".to_string());
|
||||||
|
let req = ContextRequest {
|
||||||
|
tool: None,
|
||||||
|
task: None,
|
||||||
|
signature_source: None,
|
||||||
|
project: None,
|
||||||
|
scope: None,
|
||||||
|
budget: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let result = lookup.lookup(req).await;
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_context_lookup_with_tool() {
|
||||||
|
let lookup = ContextLookup::new(6000, "test".to_string(), "project".to_string());
|
||||||
|
let req = ContextRequest {
|
||||||
|
tool: Some("github-actions".to_string()),
|
||||||
|
task: None,
|
||||||
|
signature_source: None,
|
||||||
|
project: Some("test".to_string()),
|
||||||
|
scope: None,
|
||||||
|
budget: Some(6000),
|
||||||
|
};
|
||||||
|
|
||||||
|
let result = lookup.lookup(req).await;
|
||||||
|
assert!(result.is_ok());
|
||||||
|
let resp = result.unwrap();
|
||||||
|
assert_eq!(resp.budget.limit, 6000);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_skill_recommendation() {
|
||||||
|
let skill = SkillRecommendation {
|
||||||
|
name: "ci-triage".to_string(),
|
||||||
|
score: 0.77,
|
||||||
|
description: Some("CI troubleshooting".to_string()),
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(skill.name, "ci-triage");
|
||||||
|
assert!(skill.score > 0.7);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,547 @@
|
|||||||
|
//! M8.2 — Dual-write indexing pipeline
|
||||||
|
//!
|
||||||
|
//! Coordinates atomic writes to both pgvector (embedding search) and OpenSearch (lexical search).
|
||||||
|
//! Same chunk_id in both stores. If OpenSearch fails, marks `opensearch_pending=true` for eventual
|
||||||
|
//! consistency retry loop.
|
||||||
|
|
||||||
|
use anyhow::{anyhow, Result};
|
||||||
|
use sha2::{Digest, Sha256};
|
||||||
|
use sqlx::PgPool;
|
||||||
|
use uuid::Uuid;
|
||||||
|
use pgvector::Vector;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use crate::opensearch_client::OpenSearchClient;
|
||||||
|
use crate::queue_adapter::QueueAdapter;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct DualWriteIndexer {
|
||||||
|
pool: PgPool,
|
||||||
|
opensearch: Option<Arc<OpenSearchClient>>,
|
||||||
|
/// Queue adapter for concurrent dual-write processing
|
||||||
|
/// Can be: kmsvc (production), in-memory (testing), or SQS (future)
|
||||||
|
pub queue: Arc<dyn QueueAdapter>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Input chunk for dual-write
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct ChunkInput {
|
||||||
|
pub content: String,
|
||||||
|
pub source: String,
|
||||||
|
pub project: String,
|
||||||
|
pub level: String, // "L0", "L1", "L2", "R"
|
||||||
|
pub breadcrumb: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Result of dual-write operation
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct DualWriteResult {
|
||||||
|
pub chunk_id: Uuid,
|
||||||
|
pub chunk_hash: String,
|
||||||
|
pub pgvector_success: bool,
|
||||||
|
pub opensearch_success: bool,
|
||||||
|
pub opensearch_pending: bool, // true if OpenSearch failed
|
||||||
|
pub error: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DualWriteIndexer {
|
||||||
|
/// Create dual-write indexer with queue adapter
|
||||||
|
pub fn new(
|
||||||
|
pool: PgPool,
|
||||||
|
opensearch: Option<Arc<OpenSearchClient>>,
|
||||||
|
queue: Arc<dyn QueueAdapter>,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
pool,
|
||||||
|
opensearch,
|
||||||
|
queue,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Queue chunk for dual-write processing
|
||||||
|
///
|
||||||
|
/// Sequence:
|
||||||
|
/// 1. Check dedup (chunk_hash exists AND indexed_in_pgvector AND indexed_in_opensearch)
|
||||||
|
/// 2. Queue message to external queue service (kmsvc/SQS/etc)
|
||||||
|
/// 3. Concurrent workers receive from queue and perform dual-write
|
||||||
|
///
|
||||||
|
/// Returns message_id for tracking progress
|
||||||
|
pub async fn queue_chunk(
|
||||||
|
&self,
|
||||||
|
chunk: &ChunkInput,
|
||||||
|
embedding: &[f32],
|
||||||
|
) -> Result<String> {
|
||||||
|
let chunk_id = Uuid::new_v4();
|
||||||
|
let chunk_hash = self.compute_hash(&chunk.content);
|
||||||
|
|
||||||
|
// Check deduplication
|
||||||
|
if self.is_already_indexed(&chunk_hash, &chunk.project).await? {
|
||||||
|
tracing::debug!("Chunk already indexed (dedup): {}", chunk_hash);
|
||||||
|
return Ok(Uuid::nil().to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build message attributes
|
||||||
|
let mut attributes = std::collections::HashMap::new();
|
||||||
|
attributes.insert("source".to_string(), chunk.source.clone());
|
||||||
|
attributes.insert("level".to_string(), chunk.level.clone());
|
||||||
|
attributes.insert("breadcrumb".to_string(), serde_json::to_string(&chunk.breadcrumb)?);
|
||||||
|
attributes.insert("embedding_size".to_string(), embedding.len().to_string());
|
||||||
|
|
||||||
|
// Build message body
|
||||||
|
let body = serde_json::json!({
|
||||||
|
"chunk_id": chunk_id,
|
||||||
|
"content": chunk.content,
|
||||||
|
"source": chunk.source,
|
||||||
|
"level": chunk.level,
|
||||||
|
"breadcrumb": chunk.breadcrumb,
|
||||||
|
"embedding": embedding,
|
||||||
|
}).to_string();
|
||||||
|
|
||||||
|
// Queue message
|
||||||
|
let message_id = self.queue.send_chunk(
|
||||||
|
chunk_id,
|
||||||
|
body,
|
||||||
|
chunk.project.clone(),
|
||||||
|
attributes,
|
||||||
|
).await?;
|
||||||
|
|
||||||
|
tracing::info!("Chunk queued for dual-write: message_id={}, chunk_hash={}", message_id, chunk_hash);
|
||||||
|
|
||||||
|
Ok(message_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Worker: Process queued chunk for dual-write
|
||||||
|
///
|
||||||
|
/// Called by concurrent workers receiving from queue.
|
||||||
|
/// Sequence:
|
||||||
|
/// 1. Receive message from queue
|
||||||
|
/// 2. Write to pgvector with embedding
|
||||||
|
/// 3. Write to OpenSearch (fail-soft)
|
||||||
|
/// 4. Delete from queue on success, or extend visibility on retry
|
||||||
|
pub async fn process_queued_chunk(
|
||||||
|
&self,
|
||||||
|
message: &crate::queue_adapter::QueueMessage,
|
||||||
|
embedding: &[f32],
|
||||||
|
) -> Result<DualWriteResult> {
|
||||||
|
let body: serde_json::Value = serde_json::from_str(&message.body)?;
|
||||||
|
let chunk_id = body["chunk_id"].as_str().ok_or_else(|| anyhow!("Missing chunk_id"))?
|
||||||
|
.parse::<Uuid>()?;
|
||||||
|
let content = body["content"].as_str().ok_or_else(|| anyhow!("Missing content"))?.to_string();
|
||||||
|
let source = body["source"].as_str().ok_or_else(|| anyhow!("Missing source"))?.to_string();
|
||||||
|
let project = message.project.clone();
|
||||||
|
let level = body["level"].as_str().ok_or_else(|| anyhow!("Missing level"))?.to_string();
|
||||||
|
let breadcrumb: Vec<String> = serde_json::from_value(body["breadcrumb"].clone())?;
|
||||||
|
|
||||||
|
let chunk_hash = self.compute_hash(&content);
|
||||||
|
|
||||||
|
// Write to pgvector
|
||||||
|
let pgvector_success = self
|
||||||
|
.write_pgvector(
|
||||||
|
&chunk_id,
|
||||||
|
&chunk_hash,
|
||||||
|
&content,
|
||||||
|
&source,
|
||||||
|
&project,
|
||||||
|
&level,
|
||||||
|
&breadcrumb,
|
||||||
|
embedding,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
if !pgvector_success.is_ok() {
|
||||||
|
tracing::error!("pgvector write failed: {}", pgvector_success.as_ref().err().unwrap());
|
||||||
|
// Extend visibility timeout for retry
|
||||||
|
self.queue.change_visibility(&message.message_id, &message.receipt_handle, 300).await.ok();
|
||||||
|
return Ok(DualWriteResult {
|
||||||
|
chunk_id,
|
||||||
|
chunk_hash,
|
||||||
|
pgvector_success: false,
|
||||||
|
opensearch_success: false,
|
||||||
|
opensearch_pending: false,
|
||||||
|
error: Some(format!("{:?}", pgvector_success.err())),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// Write to OpenSearch (fail-soft)
|
||||||
|
let opensearch_success = if let Some(os_client) = &self.opensearch {
|
||||||
|
self.write_opensearch(
|
||||||
|
os_client,
|
||||||
|
&chunk_id,
|
||||||
|
&content,
|
||||||
|
&source,
|
||||||
|
&project,
|
||||||
|
&level,
|
||||||
|
&breadcrumb,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
} else {
|
||||||
|
Ok(())
|
||||||
|
};
|
||||||
|
|
||||||
|
let opensearch_pending = opensearch_success.is_err();
|
||||||
|
|
||||||
|
if opensearch_pending {
|
||||||
|
tracing::warn!(
|
||||||
|
"OpenSearch write failed, marking for retry: {}",
|
||||||
|
opensearch_success.as_ref().err().unwrap()
|
||||||
|
);
|
||||||
|
self.queue.change_visibility(&message.message_id, &message.receipt_handle, 300).await.ok();
|
||||||
|
} else {
|
||||||
|
// Success: delete from queue
|
||||||
|
self.queue.delete_chunk(&message.message_id, &message.receipt_handle).await.ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(DualWriteResult {
|
||||||
|
chunk_id,
|
||||||
|
chunk_hash,
|
||||||
|
pgvector_success: pgvector_success.is_ok(),
|
||||||
|
opensearch_success: opensearch_success.is_ok(),
|
||||||
|
opensearch_pending,
|
||||||
|
error: if opensearch_pending {
|
||||||
|
Some(format!("{:?}", opensearch_success.err()))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Legacy: Direct dual-write (for backward compatibility)
|
||||||
|
///
|
||||||
|
/// If queue adapter is not available, use this for synchronous processing.
|
||||||
|
pub async fn dual_write(
|
||||||
|
&self,
|
||||||
|
chunk: &ChunkInput,
|
||||||
|
embedding: &[f32],
|
||||||
|
) -> Result<DualWriteResult> {
|
||||||
|
let chunk_id = Uuid::new_v4();
|
||||||
|
let chunk_hash = self.compute_hash(&chunk.content);
|
||||||
|
|
||||||
|
// Step 1: Check deduplication
|
||||||
|
if self.is_already_indexed(&chunk_hash, &chunk.project).await? {
|
||||||
|
tracing::debug!("Chunk already indexed (dedup): {}", chunk_hash);
|
||||||
|
return Ok(DualWriteResult {
|
||||||
|
chunk_id: Uuid::nil(), // Placeholder
|
||||||
|
chunk_hash,
|
||||||
|
pgvector_success: true,
|
||||||
|
opensearch_success: true,
|
||||||
|
opensearch_pending: false,
|
||||||
|
error: Some("already_indexed".to_string()),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// Step 2: Write to pgvector
|
||||||
|
let pgvector_success = self.write_pgvector(
|
||||||
|
&chunk_id,
|
||||||
|
&chunk_hash,
|
||||||
|
&chunk.content,
|
||||||
|
&chunk.source,
|
||||||
|
&chunk.project,
|
||||||
|
&chunk.level,
|
||||||
|
&chunk.breadcrumb,
|
||||||
|
embedding,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
if !pgvector_success.is_ok() {
|
||||||
|
tracing::error!("pgvector write failed: {}", pgvector_success.as_ref().err().unwrap());
|
||||||
|
return Ok(DualWriteResult {
|
||||||
|
chunk_id,
|
||||||
|
chunk_hash,
|
||||||
|
pgvector_success: false,
|
||||||
|
opensearch_success: false,
|
||||||
|
opensearch_pending: false,
|
||||||
|
error: Some(format!("{:?}", pgvector_success.err())),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// Step 3: Write to OpenSearch (fail-soft)
|
||||||
|
let opensearch_success = if let Some(os_client) = &self.opensearch {
|
||||||
|
self.write_opensearch(
|
||||||
|
os_client,
|
||||||
|
&chunk_id,
|
||||||
|
&chunk.content,
|
||||||
|
&chunk.source,
|
||||||
|
&chunk.project,
|
||||||
|
&chunk.level,
|
||||||
|
&chunk.breadcrumb,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
} else {
|
||||||
|
// OpenSearch not configured, skip
|
||||||
|
Ok(())
|
||||||
|
};
|
||||||
|
|
||||||
|
let opensearch_pending = opensearch_success.is_err();
|
||||||
|
|
||||||
|
if opensearch_pending {
|
||||||
|
tracing::warn!(
|
||||||
|
"OpenSearch write failed for chunk {}, marked for retry: {}",
|
||||||
|
chunk_id,
|
||||||
|
opensearch_success.as_ref().err().unwrap()
|
||||||
|
);
|
||||||
|
// Mark as pending in pgvector
|
||||||
|
self.mark_opensearch_pending(&chunk_id).await.ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Step 4: Update indexed flags
|
||||||
|
let pgvector_ok = pgvector_success.is_ok();
|
||||||
|
let opensearch_ok = opensearch_success.is_ok();
|
||||||
|
|
||||||
|
if pgvector_ok {
|
||||||
|
self.update_pgvector_indexed(&chunk_id).await.ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
if opensearch_ok {
|
||||||
|
self.update_opensearch_indexed(&chunk_id).await.ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(DualWriteResult {
|
||||||
|
chunk_id,
|
||||||
|
chunk_hash,
|
||||||
|
pgvector_success: pgvector_ok,
|
||||||
|
opensearch_success: opensearch_ok,
|
||||||
|
opensearch_pending,
|
||||||
|
error: if opensearch_pending {
|
||||||
|
Some(format!("{:?}", opensearch_success.err()))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Compute SHA256 hash of content for deduplication
|
||||||
|
fn compute_hash(&self, content: &str) -> String {
|
||||||
|
let mut hasher = Sha256::new();
|
||||||
|
hasher.update(content.as_bytes());
|
||||||
|
format!("{:x}", hasher.finalize())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if chunk is already fully indexed
|
||||||
|
async fn is_already_indexed(&self, chunk_hash: &str, project: &str) -> Result<bool> {
|
||||||
|
let row = sqlx::query_scalar::<_, bool>(
|
||||||
|
"SELECT (indexed_in_pgvector AND indexed_in_opensearch)
|
||||||
|
FROM chunks
|
||||||
|
WHERE chunk_hash = $1 AND project = $2
|
||||||
|
LIMIT 1"
|
||||||
|
)
|
||||||
|
.bind(chunk_hash)
|
||||||
|
.bind(project)
|
||||||
|
.fetch_optional(&self.pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(row.unwrap_or(false))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Write chunk to pgvector
|
||||||
|
async fn write_pgvector(
|
||||||
|
&self,
|
||||||
|
chunk_id: &Uuid,
|
||||||
|
chunk_hash: &str,
|
||||||
|
content: &str,
|
||||||
|
source: &str,
|
||||||
|
project: &str,
|
||||||
|
level: &str,
|
||||||
|
breadcrumb: &[String],
|
||||||
|
embedding: &[f32],
|
||||||
|
) -> Result<()> {
|
||||||
|
let embedding_vec = Vector::from(embedding.to_vec());
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO chunks (id, chunk_hash, content, source, project, level, breadcrumb, embedding, indexed_in_pgvector, pgvector_indexed_at)
|
||||||
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, true, now())
|
||||||
|
ON CONFLICT (id) DO UPDATE SET
|
||||||
|
embedding = EXCLUDED.embedding,
|
||||||
|
indexed_in_pgvector = true,
|
||||||
|
pgvector_indexed_at = now()"
|
||||||
|
)
|
||||||
|
.bind(chunk_id)
|
||||||
|
.bind(chunk_hash)
|
||||||
|
.bind(content)
|
||||||
|
.bind(source)
|
||||||
|
.bind(project)
|
||||||
|
.bind(level)
|
||||||
|
.bind(breadcrumb)
|
||||||
|
.bind(embedding_vec)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Write chunk to OpenSearch
|
||||||
|
async fn write_opensearch(
|
||||||
|
&self,
|
||||||
|
os_client: &Arc<OpenSearchClient>,
|
||||||
|
chunk_id: &Uuid,
|
||||||
|
content: &str,
|
||||||
|
source: &str,
|
||||||
|
project: &str,
|
||||||
|
level: &str,
|
||||||
|
breadcrumb: &[String],
|
||||||
|
) -> Result<()> {
|
||||||
|
// Note: JWT token handling would come from AppState in http_server
|
||||||
|
// For now, we'll pass empty token—production code should inject from context
|
||||||
|
os_client
|
||||||
|
.index_document(
|
||||||
|
&chunk_id.to_string(),
|
||||||
|
content,
|
||||||
|
source,
|
||||||
|
level,
|
||||||
|
breadcrumb.to_vec(),
|
||||||
|
"", // TODO: inject JWT from AppState
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mark chunk as pending OpenSearch retry
|
||||||
|
async fn mark_opensearch_pending(&self, chunk_id: &Uuid) -> Result<()> {
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE chunks
|
||||||
|
SET opensearch_pending = true, opensearch_retry_count = opensearch_retry_count + 1, opensearch_last_retry_at = now()
|
||||||
|
WHERE id = $1"
|
||||||
|
)
|
||||||
|
.bind(chunk_id)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mark chunk as pgvector indexed
|
||||||
|
async fn update_pgvector_indexed(&self, chunk_id: &Uuid) -> Result<()> {
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE chunks SET indexed_in_pgvector = true, pgvector_indexed_at = now() WHERE id = $1"
|
||||||
|
)
|
||||||
|
.bind(chunk_id)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mark chunk as OpenSearch indexed
|
||||||
|
async fn update_opensearch_indexed(&self, chunk_id: &Uuid) -> Result<()> {
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE chunks SET indexed_in_opensearch = true, opensearch_pending = false, opensearch_indexed_at = now() WHERE id = $1"
|
||||||
|
)
|
||||||
|
.bind(chunk_id)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Retry failed OpenSearch writes (background task)
|
||||||
|
///
|
||||||
|
/// Polls for chunks where opensearch_pending=true and retries up to 3 times.
|
||||||
|
/// Runs every 5 minutes.
|
||||||
|
pub async fn retry_pending_chunks(&self, project: &str, max_retries: i32) -> Result<usize> {
|
||||||
|
if self.opensearch.is_none() {
|
||||||
|
return Ok(0); // Skip if OpenSearch not configured
|
||||||
|
}
|
||||||
|
|
||||||
|
let pending = sqlx::query_as::<_, (Uuid, String, String, String, Vec<String>)>(
|
||||||
|
"SELECT id, content, source, level, breadcrumb
|
||||||
|
FROM chunks
|
||||||
|
WHERE project = $1 AND opensearch_pending = true AND opensearch_retry_count < $2
|
||||||
|
ORDER BY opensearch_last_retry_at ASC
|
||||||
|
LIMIT 100"
|
||||||
|
)
|
||||||
|
.bind(project)
|
||||||
|
.bind(max_retries)
|
||||||
|
.fetch_all(&self.pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let mut succeeded = 0;
|
||||||
|
|
||||||
|
for (chunk_id, content, source, level, breadcrumb) in pending {
|
||||||
|
if let Err(e) = self
|
||||||
|
.write_opensearch(
|
||||||
|
self.opensearch.as_ref().unwrap(),
|
||||||
|
&chunk_id,
|
||||||
|
&content,
|
||||||
|
&source,
|
||||||
|
project,
|
||||||
|
&level,
|
||||||
|
&breadcrumb,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
tracing::warn!("Retry failed for chunk {}: {}", chunk_id, e);
|
||||||
|
// Increment retry count
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE chunks SET opensearch_retry_count = opensearch_retry_count + 1, opensearch_last_retry_at = now() WHERE id = $1"
|
||||||
|
)
|
||||||
|
.bind(&chunk_id)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
} else {
|
||||||
|
tracing::info!("Retry succeeded for chunk {}", chunk_id);
|
||||||
|
self.update_opensearch_indexed(&chunk_id).await.ok();
|
||||||
|
succeeded += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(succeeded)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get retry statistics
|
||||||
|
pub async fn retry_stats(&self, project: &str) -> Result<(usize, usize)> {
|
||||||
|
let pending: (i64,) = sqlx::query_as(
|
||||||
|
"SELECT COUNT(*) FROM chunks WHERE project = $1 AND opensearch_pending = true"
|
||||||
|
)
|
||||||
|
.bind(project)
|
||||||
|
.fetch_one(&self.pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let failed: (i64,) = sqlx::query_as(
|
||||||
|
"SELECT COUNT(*) FROM chunks WHERE project = $1 AND opensearch_retry_count >= 3"
|
||||||
|
)
|
||||||
|
.bind(project)
|
||||||
|
.fetch_one(&self.pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok((pending.0 as usize, failed.0 as usize))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_compute_hash() {
|
||||||
|
let queue = Arc::new(crate::queue_adapter::InMemoryQueueAdapter::new());
|
||||||
|
let indexer = DualWriteIndexer::new(
|
||||||
|
sqlx::pool::PoolOptions::new().max_connections(1).connect_lazy("postgresql://localhost").unwrap(),
|
||||||
|
None,
|
||||||
|
queue,
|
||||||
|
);
|
||||||
|
|
||||||
|
let hash1 = indexer.compute_hash("same content");
|
||||||
|
let hash2 = indexer.compute_hash("same content");
|
||||||
|
assert_eq!(hash1, hash2, "Same content must produce same hash");
|
||||||
|
|
||||||
|
let hash3 = indexer.compute_hash("different");
|
||||||
|
assert_ne!(hash1, hash3, "Different content must produce different hash");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_hash_deterministic() {
|
||||||
|
let queue = Arc::new(crate::queue_adapter::InMemoryQueueAdapter::new());
|
||||||
|
let indexer = DualWriteIndexer::new(
|
||||||
|
sqlx::pool::PoolOptions::new().max_connections(1).connect_lazy("postgresql://localhost").unwrap(),
|
||||||
|
None,
|
||||||
|
queue,
|
||||||
|
);
|
||||||
|
|
||||||
|
let content = "ERROR: permission denied\nStack trace...";
|
||||||
|
let hash1 = indexer.compute_hash(content);
|
||||||
|
let hash2 = indexer.compute_hash(content);
|
||||||
|
|
||||||
|
assert_eq!(hash1, hash2);
|
||||||
|
assert_eq!(hash1.len(), 64); // SHA256 hex is 64 chars
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,138 @@
|
|||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::{BTreeMap, VecDeque};
|
||||||
|
use uuid::Uuid;
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
|
||||||
|
/// Record (L0 evidence).
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct Record {
|
||||||
|
pub role: String,
|
||||||
|
pub text: String,
|
||||||
|
pub timestamp: String,
|
||||||
|
pub source_position: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Git context enrichment.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct GitContext {
|
||||||
|
pub file: Option<String>,
|
||||||
|
pub commit_sha: Option<String>,
|
||||||
|
pub author: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Ingest request with full payload.
|
||||||
|
#[derive(Debug, Deserialize, Clone)]
|
||||||
|
pub struct IngestRequest {
|
||||||
|
pub project: String,
|
||||||
|
pub source: String,
|
||||||
|
pub ingest_id: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub records: Vec<Record>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub git_repo_path: Option<String>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub git_head: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Job status.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct JobStatus {
|
||||||
|
pub job_id: String,
|
||||||
|
pub ingest_id: String,
|
||||||
|
pub project: String,
|
||||||
|
pub status: String,
|
||||||
|
pub chunks_seen: u32,
|
||||||
|
pub chunks_used: u32,
|
||||||
|
pub error: Option<String>,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
pub completed_at: Option<DateTime<Utc>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// In-memory ingest queue — per-project FIFO + global dedup.
|
||||||
|
pub struct IngestQueue {
|
||||||
|
/// All jobs (for lookup by job_id or ingest_id)
|
||||||
|
jobs: BTreeMap<String, JobStatus>,
|
||||||
|
/// Per-project queues (ingest_id order)
|
||||||
|
project_queues: BTreeMap<String, VecDeque<String>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl IngestQueue {
|
||||||
|
/// Create new queue.
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
jobs: BTreeMap::new(),
|
||||||
|
project_queues: BTreeMap::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Submit job (idempotent by ingest_id).
|
||||||
|
pub fn submit(&mut self, project: &str, ingest_id: &str) -> (String, bool) {
|
||||||
|
if let Some(existing) = self.jobs.get(ingest_id) {
|
||||||
|
return (existing.job_id.clone(), false);
|
||||||
|
}
|
||||||
|
|
||||||
|
let job_id = format!("ingest-{}", Uuid::new_v4());
|
||||||
|
let status = JobStatus {
|
||||||
|
job_id: job_id.clone(),
|
||||||
|
ingest_id: ingest_id.to_string(),
|
||||||
|
project: project.to_string(),
|
||||||
|
status: "running".to_string(),
|
||||||
|
chunks_seen: 0,
|
||||||
|
chunks_used: 0,
|
||||||
|
error: None,
|
||||||
|
created_at: Utc::now(),
|
||||||
|
completed_at: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Insert into job map
|
||||||
|
self.jobs.insert(ingest_id.to_string(), status);
|
||||||
|
|
||||||
|
// Enqueue to project-specific queue
|
||||||
|
self.project_queues
|
||||||
|
.entry(project.to_string())
|
||||||
|
.or_insert_with(VecDeque::new)
|
||||||
|
.push_back(ingest_id.to_string());
|
||||||
|
|
||||||
|
(job_id, true)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get job status by job_id.
|
||||||
|
pub fn get_status(&self, job_id: &str) -> Option<JobStatus> {
|
||||||
|
self.jobs.values().find(|j| j.job_id == job_id).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Update job status (used by background task during async processing).
|
||||||
|
pub fn update_status(
|
||||||
|
&mut self,
|
||||||
|
ingest_id: &str,
|
||||||
|
status: &str,
|
||||||
|
chunks_seen: u32,
|
||||||
|
chunks_used: u32,
|
||||||
|
error: Option<String>,
|
||||||
|
) {
|
||||||
|
if let Some(job) = self.jobs.get_mut(ingest_id) {
|
||||||
|
job.status = status.to_string();
|
||||||
|
job.chunks_seen = chunks_seen;
|
||||||
|
job.chunks_used = chunks_used;
|
||||||
|
job.error = error;
|
||||||
|
if status == "completed" || status == "failed" {
|
||||||
|
job.completed_at = Some(Utc::now());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Dequeue next job for a project (FIFO).
|
||||||
|
pub fn dequeue(&mut self, project: &str) -> Option<String> {
|
||||||
|
self.project_queues
|
||||||
|
.get_mut(project)
|
||||||
|
.and_then(|q| q.pop_front())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get queue depth for a project.
|
||||||
|
pub fn queue_depth(&self, project: &str) -> usize {
|
||||||
|
self.project_queues
|
||||||
|
.get(project)
|
||||||
|
.map(|q| q.len())
|
||||||
|
.unwrap_or(0)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -14,14 +14,15 @@
|
|||||||
/// - `PipelineResult`: comprehensive result with all metrics
|
/// - `PipelineResult`: comprehensive result with all metrics
|
||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
|
use std::collections::HashMap;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
use mem_core::{GlobalTfIdfScorer, SemanticScorer};
|
use mem_core::{GlobalTfIdfScorer, SemanticScorer};
|
||||||
use mem_ingest::wiki_link::WikiLinkGraph;
|
use mem_ingest::wiki_link::WikiLinkGraph;
|
||||||
|
|
||||||
use crate::query_router::{QueryRouter, RouterConfig};
|
use crate::query_router::{QueryRouter, RouterConfig, RoutedResult, SelectedChunk};
|
||||||
use crate::chunk_metadata::{MetadataExtractor, MetadataBooster, ChunkCategory, QueryIntent};
|
use crate::chunk_metadata::{MetadataExtractor, MetadataBooster, ChunkMetadata, ChunkCategory, QueryIntent};
|
||||||
use crate::cache_alignment::{KvCacheAligner, CachedChunk, RetrievalProfiler};
|
use crate::cache_alignment::{KvCacheAligner, CachedChunk, CacheLocalityAnalyzer, RetrievalProfiler, CacheMetrics};
|
||||||
|
|
||||||
/// Unified pipeline configuration
|
/// Unified pipeline configuration
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
|
|||||||
@@ -0,0 +1,525 @@
|
|||||||
|
//! M8.2 — Gateway Queue Adapter
|
||||||
|
//!
|
||||||
|
//! Calls SQS via `api.riotpiao.com` gateway with JWT authentication.
|
||||||
|
//! Uses X-Service routing to reach kmsvc backend.
|
||||||
|
|
||||||
|
use crate::queue_adapter::{QueueAdapter, QueueMessage, QueueStats};
|
||||||
|
use anyhow::{anyhow, Result};
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use uuid::Uuid;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Token provider trait (async)
|
||||||
|
#[async_trait]
|
||||||
|
pub trait TokenProvider: Send + Sync {
|
||||||
|
async fn token(&self) -> Result<String>;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Static JWT token provider (for testing)
|
||||||
|
pub struct StaticTokenProvider {
|
||||||
|
token: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl StaticTokenProvider {
|
||||||
|
pub fn new(token: String) -> Self {
|
||||||
|
Self { token }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl TokenProvider for StaticTokenProvider {
|
||||||
|
async fn token(&self) -> Result<String> {
|
||||||
|
Ok(self.token.clone())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Authentik token provider (production)
|
||||||
|
pub struct AuthentikTokenProvider {
|
||||||
|
issuer: String,
|
||||||
|
client_id: String,
|
||||||
|
client_secret: String,
|
||||||
|
http_client: reqwest::Client,
|
||||||
|
cached_token: Arc<tokio::sync::RwLock<CachedToken>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct CachedToken {
|
||||||
|
token: Option<String>,
|
||||||
|
expires_at: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AuthentikTokenProvider {
|
||||||
|
pub fn new(issuer: String, client_id: String, client_secret: String) -> Self {
|
||||||
|
Self {
|
||||||
|
issuer,
|
||||||
|
client_id,
|
||||||
|
client_secret,
|
||||||
|
http_client: reqwest::Client::new(),
|
||||||
|
cached_token: Arc::new(tokio::sync::RwLock::new(CachedToken {
|
||||||
|
token: None,
|
||||||
|
expires_at: 0,
|
||||||
|
})),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn refresh_token(&self) -> Result<String> {
|
||||||
|
let token_url = format!("{}/application/o/token/", self.issuer);
|
||||||
|
|
||||||
|
let params = [
|
||||||
|
("grant_type", "client_credentials"),
|
||||||
|
("client_id", &self.client_id),
|
||||||
|
("client_secret", &self.client_secret),
|
||||||
|
("scope", "openid"),
|
||||||
|
];
|
||||||
|
|
||||||
|
let resp = self
|
||||||
|
.http_client
|
||||||
|
.post(&token_url)
|
||||||
|
.form(¶ms)
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
return Err(anyhow!("Failed to get token from Authentik: {}", resp.status()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let token_resp: serde_json::Value = resp.json().await?;
|
||||||
|
let token = token_resp["access_token"]
|
||||||
|
.as_str()
|
||||||
|
.ok_or_else(|| anyhow!("No access_token in Authentik response"))?
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
let expires_in = token_resp["expires_in"]
|
||||||
|
.as_i64()
|
||||||
|
.unwrap_or(3600);
|
||||||
|
let expires_at = std::time::SystemTime::now()
|
||||||
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_secs() as i64 + expires_in;
|
||||||
|
|
||||||
|
let mut cached = self.cached_token.write().await;
|
||||||
|
cached.token = Some(token.clone());
|
||||||
|
cached.expires_at = expires_at;
|
||||||
|
|
||||||
|
tracing::debug!("Token refreshed from Authentik, expires in {}s", expires_in);
|
||||||
|
|
||||||
|
Ok(token)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl TokenProvider for AuthentikTokenProvider {
|
||||||
|
async fn token(&self) -> Result<String> {
|
||||||
|
let now = std::time::SystemTime::now()
|
||||||
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_secs() as i64;
|
||||||
|
|
||||||
|
// Check cache
|
||||||
|
{
|
||||||
|
let cached = self.cached_token.read().await;
|
||||||
|
if let Some(token) = cached.token.as_ref() {
|
||||||
|
if now < cached.expires_at - 60 {
|
||||||
|
return Ok(token.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Refresh
|
||||||
|
self.refresh_token().await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// SQS SendMessage request
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct SendMessageRequest {
|
||||||
|
#[serde(rename = "messageBody")]
|
||||||
|
message_body: String,
|
||||||
|
#[serde(rename = "messageAttributes")]
|
||||||
|
message_attributes: MessageAttributes,
|
||||||
|
#[serde(rename = "delaySeconds")]
|
||||||
|
delay_seconds: i32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// SQS SendMessage response
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct SendMessageResponse {
|
||||||
|
#[serde(rename = "messageId")]
|
||||||
|
message_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// SQS ReceiveMessage response
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct ReceiveMessageResponse {
|
||||||
|
messages: Option<Vec<SqsMessage>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// SQS Message from ReceiveMessage response
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct SqsMessage {
|
||||||
|
#[serde(rename = "messageId")]
|
||||||
|
message_id: String,
|
||||||
|
#[serde(rename = "receiptHandle")]
|
||||||
|
receipt_handle: String,
|
||||||
|
body: String,
|
||||||
|
attributes: Option<std::collections::HashMap<String, String>>,
|
||||||
|
#[serde(rename = "receiveCount")]
|
||||||
|
receive_count: i32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// SQS DeleteMessage request
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct DeleteMessageRequest {
|
||||||
|
#[serde(rename = "receiptHandle")]
|
||||||
|
receipt_handle: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Message attributes wrapper
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct MessageAttributes {
|
||||||
|
values: std::collections::HashMap<String, String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Gateway Queue Adapter
|
||||||
|
///
|
||||||
|
/// Routes through api.riotpiao.com gateway to kmsvc backend.
|
||||||
|
pub struct GatewayQueueAdapter {
|
||||||
|
gateway_url: String,
|
||||||
|
token_source: Arc<dyn TokenProvider>,
|
||||||
|
http_client: reqwest::Client,
|
||||||
|
default_queue_prefix: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl GatewayQueueAdapter {
|
||||||
|
/// Create with static token (testing)
|
||||||
|
pub fn with_static_token(gateway_url: String, token: String) -> Self {
|
||||||
|
Self {
|
||||||
|
gateway_url,
|
||||||
|
token_source: Arc::new(StaticTokenProvider::new(token)),
|
||||||
|
http_client: reqwest::Client::new(),
|
||||||
|
default_queue_prefix: "poimen-chunks".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create with Authentik provider (production)
|
||||||
|
pub fn with_authentik(
|
||||||
|
gateway_url: String,
|
||||||
|
issuer: String,
|
||||||
|
client_id: String,
|
||||||
|
client_secret: String,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
gateway_url,
|
||||||
|
token_source: Arc::new(AuthentikTokenProvider::new(issuer, client_id, client_secret)),
|
||||||
|
http_client: reqwest::Client::new(),
|
||||||
|
default_queue_prefix: "poimen-chunks".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn queue_name(&self, _project: &str) -> String {
|
||||||
|
self.default_queue_prefix.clone()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl QueueAdapter for GatewayQueueAdapter {
|
||||||
|
async fn send_chunk(
|
||||||
|
&self,
|
||||||
|
chunk_id: Uuid,
|
||||||
|
body: String,
|
||||||
|
project: String,
|
||||||
|
attributes: std::collections::HashMap<String, String>,
|
||||||
|
) -> Result<String> {
|
||||||
|
let token = self.token_source.token().await?;
|
||||||
|
|
||||||
|
// Base64 encode body
|
||||||
|
let encoded_body = base64::encode(body.as_bytes());
|
||||||
|
|
||||||
|
// Build request
|
||||||
|
let mut attrs = attributes;
|
||||||
|
attrs.insert("chunk_id".to_string(), chunk_id.to_string());
|
||||||
|
attrs.insert("project".to_string(), project.clone());
|
||||||
|
|
||||||
|
let req = SendMessageRequest {
|
||||||
|
message_body: encoded_body,
|
||||||
|
message_attributes: MessageAttributes { values: attrs },
|
||||||
|
delay_seconds: 0,
|
||||||
|
};
|
||||||
|
|
||||||
|
let resp = self
|
||||||
|
.http_client
|
||||||
|
.post(&self.gateway_url)
|
||||||
|
.header("X-Service", "sqs")
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.json(&req)
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
let status = resp.status();
|
||||||
|
let error = resp.text().await.unwrap_or_default();
|
||||||
|
return Err(anyhow!("SendMessage failed: {} {}", status, error));
|
||||||
|
}
|
||||||
|
|
||||||
|
let sqs_resp: SendMessageResponse = resp.json().await?;
|
||||||
|
|
||||||
|
tracing::debug!(
|
||||||
|
"Chunk queued via gateway: message_id={}, chunk_id={}, project={}",
|
||||||
|
sqs_resp.message_id, chunk_id, project
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(sqs_resp.message_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn receive_chunks(
|
||||||
|
&self,
|
||||||
|
max_messages: i32,
|
||||||
|
visibility_timeout_secs: i32,
|
||||||
|
project: Option<&str>,
|
||||||
|
) -> Result<Vec<QueueMessage>> {
|
||||||
|
let token = self.token_source.token().await?;
|
||||||
|
let project = project.unwrap_or("default");
|
||||||
|
let max = max_messages.min(10).max(1);
|
||||||
|
|
||||||
|
// Build query string
|
||||||
|
let queue_name = self.queue_name(project);
|
||||||
|
let query = format!(
|
||||||
|
"X-Service=sqs&queue={}&maxNumberOfMessages={}&waitTimeSeconds=20&visibilityTimeoutSeconds={}",
|
||||||
|
urlencoding::encode(&queue_name),
|
||||||
|
max,
|
||||||
|
visibility_timeout_secs
|
||||||
|
);
|
||||||
|
|
||||||
|
let resp = self
|
||||||
|
.http_client
|
||||||
|
.get(&format!("{}?{}", self.gateway_url, query))
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
let status = resp.status();
|
||||||
|
let error = resp.text().await.unwrap_or_default();
|
||||||
|
return Err(anyhow!("ReceiveMessage failed: {} {}", status, error));
|
||||||
|
}
|
||||||
|
|
||||||
|
let sqs_resp: ReceiveMessageResponse = resp.json().await?;
|
||||||
|
|
||||||
|
let mut messages = Vec::new();
|
||||||
|
if let Some(sqs_msgs) = sqs_resp.messages {
|
||||||
|
for msg in sqs_msgs {
|
||||||
|
// Decode body from base64
|
||||||
|
let body_bytes = base64::decode(msg.body.as_bytes())?;
|
||||||
|
let body = String::from_utf8(body_bytes)?;
|
||||||
|
|
||||||
|
let chunk_id = msg
|
||||||
|
.attributes
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|a| a.get("chunk_id"))
|
||||||
|
.and_then(|s| Uuid::parse_str(s).ok())
|
||||||
|
.unwrap_or_else(Uuid::nil);
|
||||||
|
|
||||||
|
messages.push(QueueMessage {
|
||||||
|
message_id: msg.message_id,
|
||||||
|
chunk_id,
|
||||||
|
body,
|
||||||
|
receive_count: msg.receive_count,
|
||||||
|
receipt_handle: msg.receipt_handle,
|
||||||
|
project: project.to_string(),
|
||||||
|
attributes: msg.attributes.unwrap_or_default(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::debug!(
|
||||||
|
"Received {} messages from queue via gateway: project={}",
|
||||||
|
messages.len(),
|
||||||
|
project
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(messages)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_chunk(&self, message_id: &str, receipt_handle: &str) -> Result<()> {
|
||||||
|
let token = self.token_source.token().await?;
|
||||||
|
|
||||||
|
let req = DeleteMessageRequest {
|
||||||
|
receipt_handle: receipt_handle.to_string(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let resp = self
|
||||||
|
.http_client
|
||||||
|
.delete(&self.gateway_url)
|
||||||
|
.header("X-Service", "sqs")
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.json(&req)
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if !resp.status().is_success() && resp.status().as_u16() != 204 {
|
||||||
|
let status = resp.status();
|
||||||
|
let error = resp.text().await.unwrap_or_default();
|
||||||
|
return Err(anyhow!("DeleteMessage failed: {} {}", status, error));
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::debug!("Message deleted via gateway: message_id={}", message_id);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn change_visibility(
|
||||||
|
&self,
|
||||||
|
message_id: &str,
|
||||||
|
_receipt_handle: &str,
|
||||||
|
visibility_timeout_secs: i32,
|
||||||
|
) -> Result<()> {
|
||||||
|
// TODO: Implement when gateway adds support for ChangeMessageVisibility
|
||||||
|
|
||||||
|
tracing::warn!(
|
||||||
|
"ChangeMessageVisibility not yet supported via gateway: message_id={}, timeout={}s",
|
||||||
|
message_id,
|
||||||
|
visibility_timeout_secs
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn send_to_dlq(&self, message_id: &str, receipt_handle: &str, reason: &str) -> Result<()> {
|
||||||
|
// Delete from main queue
|
||||||
|
self.delete_chunk(message_id, receipt_handle).await?;
|
||||||
|
|
||||||
|
// Send to DLQ
|
||||||
|
let token = self.token_source.token().await?;
|
||||||
|
|
||||||
|
let dlq_body = serde_json::json!({
|
||||||
|
"message_id": message_id,
|
||||||
|
"reason": reason,
|
||||||
|
"failed_at": std::time::SystemTime::now()
|
||||||
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_secs()
|
||||||
|
})
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
let encoded_body = base64::encode(dlq_body.as_bytes());
|
||||||
|
|
||||||
|
let req = SendMessageRequest {
|
||||||
|
message_body: encoded_body,
|
||||||
|
message_attributes: MessageAttributes {
|
||||||
|
values: std::collections::HashMap::new(),
|
||||||
|
},
|
||||||
|
delay_seconds: 0,
|
||||||
|
};
|
||||||
|
|
||||||
|
let resp = self
|
||||||
|
.http_client
|
||||||
|
.post(&self.gateway_url)
|
||||||
|
.header("X-Service", "sqs")
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.json(&req)
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
return Err(anyhow!("SendToDLQ failed: {}", resp.status()));
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::warn!(
|
||||||
|
"Message sent to DLQ via gateway: message_id={}, reason={}",
|
||||||
|
message_id,
|
||||||
|
reason
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_stats(&self, project: Option<&str>) -> Result<QueueStats> {
|
||||||
|
let _token = self.token_source.token().await?;
|
||||||
|
let _project = project.unwrap_or("default");
|
||||||
|
|
||||||
|
Ok(QueueStats {
|
||||||
|
available_messages: 0,
|
||||||
|
in_flight_messages: 0,
|
||||||
|
dead_letter_messages: 0,
|
||||||
|
total_processed: 0,
|
||||||
|
average_delay_secs: 0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn purge(&self, project: Option<&str>) -> Result<usize> {
|
||||||
|
let _token = self.token_source.token().await?;
|
||||||
|
let _project = project.unwrap_or("default");
|
||||||
|
|
||||||
|
tracing::warn!("Purge not yet supported via gateway");
|
||||||
|
|
||||||
|
Ok(0)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn health_check(&self) -> Result<()> {
|
||||||
|
let token = self.token_source.token().await?;
|
||||||
|
|
||||||
|
let query = format!(
|
||||||
|
"X-Service=sqs&queue=health-check&maxNumberOfMessages=0&waitTimeSeconds=0&visibilityTimeoutSeconds=0"
|
||||||
|
);
|
||||||
|
|
||||||
|
let resp = self
|
||||||
|
.http_client
|
||||||
|
.get(&format!("{}?{}", self.gateway_url, query))
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.timeout(std::time::Duration::from_secs(5))
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if resp.status().is_success() || resp.status().as_u16() == 404 {
|
||||||
|
tracing::debug!("Gateway health check passed");
|
||||||
|
Ok(())
|
||||||
|
} else {
|
||||||
|
Err(anyhow!("Gateway health check failed: {}", resp.status()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_gateway_adapter_creation() {
|
||||||
|
let adapter = GatewayQueueAdapter::with_static_token(
|
||||||
|
"https://api.riotpiao.com".to_string(),
|
||||||
|
"test-token".to_string(),
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(adapter.gateway_url, "https://api.riotpiao.com");
|
||||||
|
assert_eq!(adapter.default_queue_prefix, "poimen-chunks");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_queue_name_formatting() {
|
||||||
|
let adapter = GatewayQueueAdapter::with_static_token(
|
||||||
|
"https://api.riotpiao.com".to_string(),
|
||||||
|
"test-token".to_string(),
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(adapter.queue_name("myproject"), "poimen-chunks");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_base64_roundtrip() {
|
||||||
|
let original = "hello world";
|
||||||
|
let encoded = base64::encode(original.as_bytes());
|
||||||
|
let decoded = String::from_utf8(base64::decode(encoded.as_bytes()).unwrap()).unwrap();
|
||||||
|
assert_eq!(decoded, original);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_static_token_provider() {
|
||||||
|
let provider = StaticTokenProvider::new("my-token".to_string());
|
||||||
|
let token = provider.token().await.unwrap();
|
||||||
|
assert_eq!(token, "my-token");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -5,13 +5,13 @@
|
|||||||
|
|
||||||
use actix_web::{web, HttpRequest, HttpResponse};
|
use actix_web::{web, HttpRequest, HttpResponse};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::sync::Arc;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
use crate::agent::{Agent, AgentConfig, AgentCapability, DefaultAgent};
|
use crate::agent::{Agent, AgentConfig, AgentCapability, DefaultAgent};
|
||||||
use crate::agent::client_sdk::SynthesisClient;
|
use crate::agent::client_sdk::SynthesisClient;
|
||||||
use crate::handlers::response_builder;
|
use crate::handlers::response_builder;
|
||||||
use mem_store::agent_repo::AgentRepository;
|
use mem_store::agent_repo::{AgentRepository, AgentPrompt, AgentSkill, AgentDecision, RolePromptMapping};
|
||||||
use crate::metrics::{ERROR_BAD_REQUEST_AGENT, ERROR_NOT_FOUND_AGENT, ERROR_UNEXPECTED_AGENT, ERROR_UNEXPECTED_TOTAL};
|
|
||||||
use tracing::{debug, info, error, warn};
|
use tracing::{debug, info, error, warn};
|
||||||
|
|
||||||
/// Register agent request
|
/// Register agent request
|
||||||
@@ -51,14 +51,10 @@ pub async fn register_agent_handler(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if body.agent_id.is_empty() || body.project_id.is_empty() {
|
if body.agent_id.is_empty() || body.project_id.is_empty() {
|
||||||
ERROR_BAD_REQUEST_AGENT.inc();
|
|
||||||
warn!(agent_id = %body.agent_id, "Expected error: missing agent_id or project_id");
|
|
||||||
return response_builder::bad_request("agent_id and project_id required");
|
return response_builder::bad_request("agent_id and project_id required");
|
||||||
}
|
}
|
||||||
|
|
||||||
if body.capabilities.is_empty() {
|
if body.capabilities.is_empty() {
|
||||||
ERROR_BAD_REQUEST_AGENT.inc();
|
|
||||||
warn!(agent_id = %body.agent_id, "Expected error: no capabilities provided");
|
|
||||||
return response_builder::bad_request("At least one capability required");
|
return response_builder::bad_request("At least one capability required");
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -78,8 +74,6 @@ pub async fn register_agent_handler(
|
|||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
if caps.is_empty() {
|
if caps.is_empty() {
|
||||||
ERROR_BAD_REQUEST_AGENT.inc();
|
|
||||||
warn!(agent_id = %body.agent_id, "Expected error: invalid capability names");
|
|
||||||
return response_builder::bad_request("Invalid capabilities");
|
return response_builder::bad_request("Invalid capabilities");
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -93,7 +87,7 @@ pub async fn register_agent_handler(
|
|||||||
};
|
};
|
||||||
|
|
||||||
// Persist agent config to database via agent_registry table
|
// Persist agent config to database via agent_registry table
|
||||||
let _agent_repo = AgentRepository::new(state.pool.clone());
|
let agent_repo = AgentRepository::new(state.pool.clone());
|
||||||
|
|
||||||
// Verify project exists
|
// Verify project exists
|
||||||
let project_exists = sqlx::query("SELECT id FROM projects WHERE id = $1")
|
let project_exists = sqlx::query("SELECT id FROM projects WHERE id = $1")
|
||||||
@@ -102,15 +96,11 @@ pub async fn register_agent_handler(
|
|||||||
.await;
|
.await;
|
||||||
|
|
||||||
if let Err(e) = project_exists {
|
if let Err(e) = project_exists {
|
||||||
ERROR_UNEXPECTED_AGENT.inc();
|
error!("Failed to verify project: {}", e);
|
||||||
ERROR_UNEXPECTED_TOTAL.inc();
|
|
||||||
error!(agent_id = %body.agent_id, error = %e, "Unexpected error: DB failure verifying project");
|
|
||||||
return response_builder::internal_error("Database error during project verification");
|
return response_builder::internal_error("Database error during project verification");
|
||||||
}
|
}
|
||||||
|
|
||||||
if project_exists.unwrap().is_none() {
|
if project_exists.unwrap().is_none() {
|
||||||
ERROR_NOT_FOUND_AGENT.inc();
|
|
||||||
info!(agent_id = %body.agent_id, project_id = %body.project_id, "Expected error: project not found");
|
|
||||||
return response_builder::bad_request(&format!("Project not found: {}", body.project_id));
|
return response_builder::bad_request(&format!("Project not found: {}", body.project_id));
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -131,14 +121,12 @@ pub async fn register_agent_handler(
|
|||||||
.bind(&body.agent_id)
|
.bind(&body.agent_id)
|
||||||
.bind(&body.capabilities)
|
.bind(&body.capabilities)
|
||||||
.bind(&body.webhook_url)
|
.bind(&body.webhook_url)
|
||||||
.bind(body.rate_limit.unwrap_or(1000) as i32)
|
.bind(body.rate_limit.unwrap_or(1000))
|
||||||
.execute(&state.pool)
|
.execute(&state.pool)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
if let Err(e) = agent_insert {
|
if let Err(e) = agent_insert {
|
||||||
ERROR_UNEXPECTED_AGENT.inc();
|
error!("Failed to insert agent registry: {}", e);
|
||||||
ERROR_UNEXPECTED_TOTAL.inc();
|
|
||||||
error!(agent_id = %body.agent_id, error = %e, "Unexpected error: DB failure inserting agent");
|
|
||||||
return response_builder::internal_error("Failed to register agent");
|
return response_builder::internal_error("Failed to register agent");
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -215,46 +203,7 @@ pub async fn register_agent_handler(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Full agent progress response
|
/// GET /agents/{id} - Get agent status
|
||||||
#[derive(Debug, Serialize)]
|
|
||||||
pub struct AgentProgressResponse {
|
|
||||||
pub agent_id: String,
|
|
||||||
pub project_id: String,
|
|
||||||
pub capabilities: Vec<String>,
|
|
||||||
pub status: String,
|
|
||||||
pub prompts: Vec<PromptResponse>,
|
|
||||||
pub skills: Vec<SkillSummary>,
|
|
||||||
pub decisions: Vec<DecisionSummary>,
|
|
||||||
pub metrics: Option<MetricsSummary>,
|
|
||||||
pub created_at: String,
|
|
||||||
pub updated_at: String,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
|
||||||
pub struct SkillSummary {
|
|
||||||
pub name: String,
|
|
||||||
pub success_rate: f32,
|
|
||||||
pub invocation_count: i64,
|
|
||||||
pub enabled: bool,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
|
||||||
pub struct DecisionSummary {
|
|
||||||
pub action: String,
|
|
||||||
pub confidence: f32,
|
|
||||||
pub outcome_success: Option<bool>,
|
|
||||||
pub created_at: String,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
|
||||||
pub struct MetricsSummary {
|
|
||||||
pub requests_total: i64,
|
|
||||||
pub requests_success: i64,
|
|
||||||
pub error_rate: f32,
|
|
||||||
pub average_latency_ms: f32,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// GET /agents/{id} - Get agent progress
|
|
||||||
pub async fn get_agent_handler(
|
pub async fn get_agent_handler(
|
||||||
req: HttpRequest,
|
req: HttpRequest,
|
||||||
path: web::Path<String>,
|
path: web::Path<String>,
|
||||||
@@ -268,110 +217,33 @@ pub async fn get_agent_handler(
|
|||||||
return response;
|
return response;
|
||||||
}
|
}
|
||||||
|
|
||||||
debug!("Getting agent progress: {}", agent_id);
|
debug!("Getting agent: {}", agent_id);
|
||||||
|
|
||||||
// Fetch agent registry
|
// Extract JWT for agent operations
|
||||||
let agent_row = sqlx::query_as::<_, (String, Vec<String>, Option<String>, i32, String, String, String)>(
|
let jwt = crate::handlers::extract_jwt_token(&req)
|
||||||
r#"SELECT project_id, capabilities, webhook_url, rate_limit, status,
|
.unwrap_or_else(|| {
|
||||||
created_at::text, updated_at::text
|
warn!("No JWT token in get_agent request");
|
||||||
FROM agent_registry WHERE agent_id = $1"#
|
"invalid".to_string()
|
||||||
)
|
});
|
||||||
.bind(&agent_id)
|
|
||||||
.fetch_optional(&state.pool)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
let (project_id, capabilities, _webhook, _rate_limit, status, created_at, updated_at) = match agent_row {
|
// Stub: would fetch from DB
|
||||||
Ok(Some(row)) => row,
|
let config = AgentConfig {
|
||||||
Ok(None) => {
|
agent_id: agent_id.clone(),
|
||||||
ERROR_NOT_FOUND_AGENT.inc();
|
project_id: "poimen".to_string(),
|
||||||
info!(agent_id = %agent_id, "Expected error: agent not found");
|
capabilities: vec![AgentCapability::Summarization],
|
||||||
return response_builder::not_found(&format!("Agent not found: {}", agent_id));
|
webhook_url: None,
|
||||||
}
|
rate_limit: 1000,
|
||||||
Err(e) => {
|
metadata: std::collections::HashMap::new(),
|
||||||
ERROR_UNEXPECTED_AGENT.inc();
|
|
||||||
ERROR_UNEXPECTED_TOTAL.inc();
|
|
||||||
error!(agent_id = %agent_id, error = %e, "Unexpected error: DB failure fetching agent");
|
|
||||||
return response_builder::internal_error("Database error");
|
|
||||||
}
|
|
||||||
};
|
};
|
||||||
|
|
||||||
// Fetch prompts
|
let agent = DefaultAgent::new(config);
|
||||||
let prompts: Vec<PromptResponse> = sqlx::query_as::<_, (String, String, String, Option<String>, String, Vec<String>, i64, f32, i32, String)>(
|
|
||||||
r#"SELECT id::text, name, template, target_model, task_category,
|
|
||||||
tags, usage_count, avg_quality, version, created_at::text
|
|
||||||
FROM agent_prompt WHERE project_id = $1 ORDER BY created_at DESC"#
|
|
||||||
)
|
|
||||||
.bind(&project_id)
|
|
||||||
.fetch_all(&state.pool)
|
|
||||||
.await
|
|
||||||
.unwrap_or_default()
|
|
||||||
.into_iter()
|
|
||||||
.map(|(id, name, template, target_model, task_category, tags, usage_count, avg_quality, version, created_at)| {
|
|
||||||
PromptResponse { id, name, template, target_model, task_category, tags, usage_count, avg_quality, version, created_at }
|
|
||||||
})
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
// Fetch skills
|
match futures::executor::block_on(agent.status()) {
|
||||||
let skills: Vec<SkillSummary> = sqlx::query_as::<_, (String, f32, i64, bool)>(
|
status => {
|
||||||
r#"SELECT name, success_rate, invocation_count, enabled
|
info!("Agent status: {} with JWT auth", agent_id);
|
||||||
FROM agent_skill WHERE agent_id = $1 ORDER BY created_at DESC"#
|
response_builder::success_response(status)
|
||||||
)
|
}
|
||||||
.bind(&agent_id)
|
}
|
||||||
.fetch_all(&state.pool)
|
|
||||||
.await
|
|
||||||
.unwrap_or_default()
|
|
||||||
.into_iter()
|
|
||||||
.map(|(name, success_rate, invocation_count, enabled)| {
|
|
||||||
SkillSummary { name, success_rate, invocation_count, enabled }
|
|
||||||
})
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
// Fetch recent decisions
|
|
||||||
let decisions: Vec<DecisionSummary> = sqlx::query_as::<_, (String, f32, Option<bool>, String)>(
|
|
||||||
r#"SELECT action, confidence, outcome_success, created_at::text
|
|
||||||
FROM agent_decision WHERE agent_id = $1
|
|
||||||
ORDER BY created_at DESC LIMIT 20"#
|
|
||||||
)
|
|
||||||
.bind(&agent_id)
|
|
||||||
.fetch_all(&state.pool)
|
|
||||||
.await
|
|
||||||
.unwrap_or_default()
|
|
||||||
.into_iter()
|
|
||||||
.map(|(action, confidence, outcome_success, created_at)| {
|
|
||||||
DecisionSummary { action, confidence, outcome_success, created_at }
|
|
||||||
})
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
// Fetch latest metrics
|
|
||||||
let metrics = sqlx::query_as::<_, (i64, i64, f32, f32)>(
|
|
||||||
r#"SELECT requests_total, requests_success, error_rate, average_latency_ms
|
|
||||||
FROM agent_metrics WHERE agent_id = $1
|
|
||||||
ORDER BY recorded_at DESC LIMIT 1"#
|
|
||||||
)
|
|
||||||
.bind(&agent_id)
|
|
||||||
.fetch_optional(&state.pool)
|
|
||||||
.await
|
|
||||||
.ok()
|
|
||||||
.flatten()
|
|
||||||
.map(|(requests_total, requests_success, error_rate, average_latency_ms)| {
|
|
||||||
MetricsSummary { requests_total, requests_success, error_rate, average_latency_ms }
|
|
||||||
});
|
|
||||||
|
|
||||||
info!("Agent progress: {} ({} prompts, {} skills, {} decisions)",
|
|
||||||
agent_id, prompts.len(), skills.len(), decisions.len());
|
|
||||||
|
|
||||||
response_builder::success_response(AgentProgressResponse {
|
|
||||||
agent_id,
|
|
||||||
project_id,
|
|
||||||
capabilities,
|
|
||||||
status,
|
|
||||||
prompts,
|
|
||||||
skills,
|
|
||||||
decisions,
|
|
||||||
metrics,
|
|
||||||
created_at,
|
|
||||||
updated_at,
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Metrics response
|
/// Metrics response
|
||||||
@@ -518,7 +390,7 @@ pub struct PromptResponse {
|
|||||||
pub created_at: String,
|
pub created_at: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// POST /agents/{id}/prompts - Create agent prompt
|
/// POST /memory/agents/{project_id}/prompts - Create agent prompt
|
||||||
pub async fn create_prompt_handler(
|
pub async fn create_prompt_handler(
|
||||||
req: HttpRequest,
|
req: HttpRequest,
|
||||||
path: web::Path<String>,
|
path: web::Path<String>,
|
||||||
@@ -534,8 +406,6 @@ pub async fn create_prompt_handler(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if body.name.is_empty() || body.template.is_empty() {
|
if body.name.is_empty() || body.template.is_empty() {
|
||||||
ERROR_BAD_REQUEST_AGENT.inc();
|
|
||||||
warn!("Expected error: missing prompt name or template");
|
|
||||||
return response_builder::bad_request("name and template required");
|
return response_builder::bad_request("name and template required");
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -579,9 +449,7 @@ pub async fn create_prompt_handler(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
ERROR_UNEXPECTED_AGENT.inc();
|
error!("Failed to create prompt: {}", e);
|
||||||
ERROR_UNEXPECTED_TOTAL.inc();
|
|
||||||
error!(error = %e, "Unexpected error: DB failure creating prompt");
|
|
||||||
response_builder::internal_error("Failed to create prompt")
|
response_builder::internal_error("Failed to create prompt")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -594,7 +462,7 @@ pub struct MapRoleToPromptRequest {
|
|||||||
pub priority: Option<i32>,
|
pub priority: Option<i32>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// POST /agents/{id}/roles - Map role to prompt
|
/// POST /memory/agents/{project_id}/roles - Map role to prompt
|
||||||
pub async fn map_role_to_prompt_handler(
|
pub async fn map_role_to_prompt_handler(
|
||||||
req: HttpRequest,
|
req: HttpRequest,
|
||||||
path: web::Path<String>,
|
path: web::Path<String>,
|
||||||
@@ -610,8 +478,6 @@ pub async fn map_role_to_prompt_handler(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if body.role_name.is_empty() || body.prompt_id.is_empty() {
|
if body.role_name.is_empty() || body.prompt_id.is_empty() {
|
||||||
ERROR_BAD_REQUEST_AGENT.inc();
|
|
||||||
warn!("Expected error: missing role_name or prompt_id");
|
|
||||||
return response_builder::bad_request("role_name and prompt_id required");
|
return response_builder::bad_request("role_name and prompt_id required");
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -619,11 +485,7 @@ pub async fn map_role_to_prompt_handler(
|
|||||||
|
|
||||||
let prompt_uuid = match Uuid::parse_str(&body.prompt_id) {
|
let prompt_uuid = match Uuid::parse_str(&body.prompt_id) {
|
||||||
Ok(id) => id,
|
Ok(id) => id,
|
||||||
Err(_) => {
|
Err(_) => return response_builder::bad_request("Invalid prompt_id UUID format"),
|
||||||
ERROR_BAD_REQUEST_AGENT.inc();
|
|
||||||
warn!(prompt_id = %body.prompt_id, "Expected error: invalid UUID format");
|
|
||||||
return response_builder::bad_request("Invalid prompt_id UUID format");
|
|
||||||
}
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let priority = body.priority.unwrap_or(0);
|
let priority = body.priority.unwrap_or(0);
|
||||||
@@ -665,22 +527,16 @@ pub async fn map_role_to_prompt_handler(
|
|||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
ERROR_UNEXPECTED_AGENT.inc();
|
error!("Failed to create role mapping: {}", e);
|
||||||
ERROR_UNEXPECTED_TOTAL.inc();
|
|
||||||
error!(error = %e, "Unexpected error: DB failure creating role mapping");
|
|
||||||
response_builder::internal_error("Failed to map role to prompt")
|
response_builder::internal_error("Failed to map role to prompt")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Ok(None) => {
|
Ok(None) => {
|
||||||
ERROR_NOT_FOUND_AGENT.inc();
|
|
||||||
info!(prompt_id = %body.prompt_id, "Expected error: prompt not found");
|
|
||||||
response_builder::not_found(&format!("Prompt not found: {}", body.prompt_id))
|
response_builder::not_found(&format!("Prompt not found: {}", body.prompt_id))
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
ERROR_UNEXPECTED_AGENT.inc();
|
error!("Database error checking prompt: {}", e);
|
||||||
ERROR_UNEXPECTED_TOTAL.inc();
|
|
||||||
error!(error = %e, "Unexpected error: DB failure checking prompt");
|
|
||||||
response_builder::internal_error("Database error")
|
response_builder::internal_error("Database error")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -692,7 +548,7 @@ pub struct RolePromptsResponse {
|
|||||||
pub prompts: Vec<PromptResponse>,
|
pub prompts: Vec<PromptResponse>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// GET /agents/{id}/roles/{role_name}/prompts - Get prompts for role
|
/// GET /memory/agents/{project_id}/roles/{role_name}/prompts - Get prompts for role
|
||||||
pub async fn get_role_prompts_handler(
|
pub async fn get_role_prompts_handler(
|
||||||
req: HttpRequest,
|
req: HttpRequest,
|
||||||
path: web::Path<(String, String)>,
|
path: web::Path<(String, String)>,
|
||||||
@@ -747,9 +603,7 @@ pub async fn get_role_prompts_handler(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
ERROR_UNEXPECTED_AGENT.inc();
|
error!("Failed to fetch role prompts: {}", e);
|
||||||
ERROR_UNEXPECTED_TOTAL.inc();
|
|
||||||
error!(role_name = %role_name, error = %e, "Unexpected error: DB failure fetching role prompts");
|
|
||||||
response_builder::internal_error("Failed to fetch role prompts")
|
response_builder::internal_error("Failed to fetch role prompts")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -5,9 +5,10 @@
|
|||||||
|
|
||||||
use actix_web::{web, HttpRequest, HttpResponse};
|
use actix_web::{web, HttpRequest, HttpResponse};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::json;
|
||||||
|
|
||||||
use crate::http_server::AppState;
|
use crate::http_server::AppState;
|
||||||
use crate::compaction::{CompactionMode, CompactionStats};
|
use crate::compaction::{compact_memory, CompactionMode, CompactionStats};
|
||||||
|
|
||||||
/// Compaction request parameters
|
/// Compaction request parameters
|
||||||
#[derive(Debug, Deserialize, Clone)]
|
#[derive(Debug, Deserialize, Clone)]
|
||||||
|
|||||||
@@ -1,26 +1,63 @@
|
|||||||
/// Handler middleware utilities
|
/// Handler middleware utilities
|
||||||
///
|
///
|
||||||
/// Centralized auth validation for all HTTP handlers.
|
/// Centralized JWT validation + rate limiting for all HTTP handlers.
|
||||||
/// Rate limiting deferred to API gateway / riotpiao-rust-sdk (issue #56).
|
/// Eliminates boilerplate across endpoints, improves testability.
|
||||||
|
|
||||||
use actix_web::{HttpRequest, HttpResponse};
|
use actix_web::{HttpRequest, HttpResponse};
|
||||||
|
use serde_json::json;
|
||||||
use crate::http_server::AppState;
|
use crate::http_server::AppState;
|
||||||
|
|
||||||
/// Result type for middleware operations
|
/// Result type for middleware operations
|
||||||
pub type MiddlewareResult<T> = Result<T, HttpResponse>;
|
pub type MiddlewareResult<T> = Result<T, HttpResponse>;
|
||||||
|
|
||||||
/// Validate auth + rate limit (stub)
|
/// Validate JWT token + check rate limit
|
||||||
///
|
///
|
||||||
/// Auth validation delegates to http_server::validate_auth.
|
/// Handles:
|
||||||
/// Rate limiting deferred to API gateway (issue #56).
|
/// 1. Extract Authorization header
|
||||||
|
/// 2. Validate JWT (if auth enabled)
|
||||||
|
/// 3. Check rate limit (if limiter enabled)
|
||||||
|
/// 4. Return error response on failure
|
||||||
|
///
|
||||||
|
/// # Usage
|
||||||
|
/// ```ignore
|
||||||
|
/// validate_and_rate_limit(&req, &state, "compact", 10)?;
|
||||||
|
/// // If we get here, both JWT and rate limit checks passed
|
||||||
|
/// ```
|
||||||
pub fn validate_and_rate_limit(
|
pub fn validate_and_rate_limit(
|
||||||
_req: &HttpRequest,
|
req: &HttpRequest,
|
||||||
_state: &AppState,
|
state: &AppState,
|
||||||
_endpoint: &str,
|
endpoint: &str,
|
||||||
_rate_limit: u32,
|
rate_limit: u32,
|
||||||
) -> MiddlewareResult<()> {
|
) -> MiddlewareResult<()> {
|
||||||
// Auth is handled by validate_auth() in http_server.rs at the handler level.
|
// 1. JWT validation (if enabled)
|
||||||
// Rate limiting deferred to API gateway / riotpiao-rust-sdk (issue #56).
|
if let Some(jwt_validator) = &state.jwt_validator {
|
||||||
|
let auth_header = req
|
||||||
|
.headers()
|
||||||
|
.get("Authorization")
|
||||||
|
.and_then(|h| h.to_str().ok())
|
||||||
|
.ok_or_else(|| {
|
||||||
|
HttpResponse::Unauthorized().json(json!({
|
||||||
|
"error": "Missing Authorization header"
|
||||||
|
}))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
crate::jwt_validator::JwtValidator::extract_bearer_token(auth_header).map_err(|e| {
|
||||||
|
HttpResponse::Unauthorized().json(json!({
|
||||||
|
"error": format!("JWT validation failed: {}", e)
|
||||||
|
}))
|
||||||
|
})?;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. Rate limiting (if enabled)
|
||||||
|
state
|
||||||
|
.rate_limiter
|
||||||
|
.check("default", endpoint)
|
||||||
|
.map_err(|e| {
|
||||||
|
HttpResponse::TooManyRequests().json(json!({
|
||||||
|
"error": format!("Rate limit exceeded: {}", e.reason())
|
||||||
|
}))
|
||||||
|
})?;
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -28,7 +65,14 @@ pub fn validate_and_rate_limit(
|
|||||||
///
|
///
|
||||||
/// Tries to decode JWT from Authorization header to get `sub` claim.
|
/// Tries to decode JWT from Authorization header to get `sub` claim.
|
||||||
/// Falls back to "anonymous" if auth is disabled or header missing.
|
/// Falls back to "anonymous" if auth is disabled or header missing.
|
||||||
pub fn extract_user_id(req: &HttpRequest, _state: &AppState) -> String {
|
/// Used by metrics to track errors/requests per user.
|
||||||
|
pub fn extract_user_id(req: &HttpRequest, state: &AppState) -> String {
|
||||||
|
// If auth disabled, check synthetic claims
|
||||||
|
if state.jwt_validator.is_none() {
|
||||||
|
return "anonymous".to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to extract sub from JWT
|
||||||
let token = req.headers()
|
let token = req.headers()
|
||||||
.get("Authorization")
|
.get("Authorization")
|
||||||
.and_then(|h| h.to_str().ok())
|
.and_then(|h| h.to_str().ok())
|
||||||
@@ -39,12 +83,14 @@ pub fn extract_user_id(req: &HttpRequest, _state: &AppState) -> String {
|
|||||||
return "anonymous".to_string();
|
return "anonymous".to_string();
|
||||||
}
|
}
|
||||||
|
|
||||||
// Decode JWT payload without validation (already validated upstream)
|
// Decode JWT payload without validation (already validated by validate_and_rate_limit)
|
||||||
|
// JWT format: header.payload.signature
|
||||||
let parts: Vec<&str> = token.split('.').collect();
|
let parts: Vec<&str> = token.split('.').collect();
|
||||||
if parts.len() != 3 {
|
if parts.len() != 3 {
|
||||||
return "anonymous".to_string();
|
return "anonymous".to_string();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Decode base64 payload
|
||||||
use base64::Engine;
|
use base64::Engine;
|
||||||
let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||||
if let Ok(payload_bytes) = engine.decode(parts[1]) {
|
if let Ok(payload_bytes) = engine.decode(parts[1]) {
|
||||||
@@ -64,12 +110,15 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_middleware_result_type_is_result() {
|
fn test_middleware_result_type_is_result() {
|
||||||
|
// Verify type alias works
|
||||||
let _result: MiddlewareResult<()> = Ok(());
|
let _result: MiddlewareResult<()> = Ok(());
|
||||||
let _result: MiddlewareResult<()> = Err(HttpResponse::Unauthorized().finish());
|
let _result: MiddlewareResult<()> = Err(HttpResponse::Unauthorized().finish());
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_validate_and_rate_limit_signature() {
|
fn test_validate_and_rate_limit_signature() {
|
||||||
|
// Just verify the function signature is correct (compile-time test)
|
||||||
|
// Runtime tests require full AppState with mocks
|
||||||
let _ = validate_and_rate_limit;
|
let _ = validate_and_rate_limit;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -7,15 +7,7 @@ use serde::Serialize;
|
|||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
|
||||||
/// Query result (moved from deleted query_worker module)
|
use crate::query_worker::QueryResult;
|
||||||
#[derive(Debug, Clone)]
|
|
||||||
pub struct QueryResult {
|
|
||||||
pub level: String,
|
|
||||||
pub score: f32,
|
|
||||||
pub text: String,
|
|
||||||
pub source: Option<String>,
|
|
||||||
pub provenance: Vec<String>,
|
|
||||||
}
|
|
||||||
|
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
// Query Parameters
|
// Query Parameters
|
||||||
|
|||||||
@@ -1,6 +1,8 @@
|
|||||||
use actix_web::{HttpRequest, HttpResponse};
|
use actix_web::{web, HttpRequest, HttpResponse};
|
||||||
use chrono::{DateTime, Utc};
|
use chrono::{DateTime, Utc};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
use sqlx::PgPool;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
use crate::auth::AuthGuard;
|
use crate::auth::AuthGuard;
|
||||||
|
|
||||||
|
|||||||
@@ -152,7 +152,7 @@ pub async fn rebuild(
|
|||||||
/// GET /memory/rebuild/status
|
/// GET /memory/rebuild/status
|
||||||
pub async fn rebuild_status(
|
pub async fn rebuild_status(
|
||||||
req: HttpRequest,
|
req: HttpRequest,
|
||||||
_pool: web::Data<PgPool>,
|
pool: web::Data<PgPool>,
|
||||||
) -> HttpResponse {
|
) -> HttpResponse {
|
||||||
// Verify auth
|
// Verify auth
|
||||||
if let Err(e) = AuthGuard::extract_token(req.headers().get("Authorization").and_then(|v| v.to_str().ok()).unwrap_or("")) {
|
if let Err(e) = AuthGuard::extract_token(req.headers().get("Authorization").and_then(|v| v.to_str().ok()).unwrap_or("")) {
|
||||||
|
|||||||
@@ -4,10 +4,11 @@
|
|||||||
|
|
||||||
use actix_web::{web, HttpRequest, HttpResponse};
|
use actix_web::{web, HttpRequest, HttpResponse};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::json;
|
||||||
use tracing::{debug, error, info};
|
use tracing::{debug, error, info};
|
||||||
|
|
||||||
use crate::http_server::AppState;
|
use crate::http_server::AppState;
|
||||||
use crate::query::{SemanticRetriever, CommunityDetector, CommunityDetectionResult, PathFinder, PathFindingResult, FacetedSearch, AvailableFacets, FacetFilters};
|
use crate::query::{SemanticRetriever, EntityResult, EdgeResult, HybridResult, CommunityDetector, CommunityDetectionResult, PathFinder, PathFindingResult, FacetedSearch, AvailableFacets, FacetFilters};
|
||||||
|
|
||||||
/// Request for semantic entity search
|
/// Request for semantic entity search
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
|
|||||||
@@ -14,8 +14,8 @@ use crate::http_server::AppState;
|
|||||||
use crate::query::{
|
use crate::query::{
|
||||||
EntityLinker, MentionLink, AliasSuggestion, MergeSuggestion, CoreferenceCluster,
|
EntityLinker, MentionLink, AliasSuggestion, MergeSuggestion, CoreferenceCluster,
|
||||||
InferenceEngine, InferenceRule, InferredFact, ReasoningPath, TransitiveClosure,
|
InferenceEngine, InferenceRule, InferredFact, ReasoningPath, TransitiveClosure,
|
||||||
QueryReasoner,
|
QueryReasoner, SubQuery, Constraint, QuestionType, ReasonedAnswer,
|
||||||
Summarizer, SummarizationStrategy,
|
Summarizer, SummarizationStrategy, Summary, KeyFact,
|
||||||
};
|
};
|
||||||
|
|
||||||
/// Request to link entities
|
/// Request to link entities
|
||||||
@@ -145,7 +145,7 @@ pub async fn link_entities_handler(
|
|||||||
|
|
||||||
let total = links.len() + unlinked.len();
|
let total = links.len() + unlinked.len();
|
||||||
let link_rate = if total > 0 {
|
let link_rate = if total > 0 {
|
||||||
links.len() as f32 / total as f32
|
(links.len() as f32 / total as f32)
|
||||||
} else {
|
} else {
|
||||||
0.0
|
0.0
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -9,12 +9,12 @@
|
|||||||
|
|
||||||
use actix_web::{web, HttpRequest, HttpResponse};
|
use actix_web::{web, HttpRequest, HttpResponse};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use serde_json::Value;
|
use serde_json::{json, Value};
|
||||||
use tracing::{debug, error, info};
|
use tracing::{debug, error, info};
|
||||||
|
|
||||||
use crate::http_server::AppState;
|
use crate::http_server::AppState;
|
||||||
use crate::query::{
|
use crate::query::{
|
||||||
SemanticRetriever,
|
SemanticRetriever, EntityResult, EdgeResult, HybridResult,
|
||||||
CommunityDetector, CommunityDetectionResult,
|
CommunityDetector, CommunityDetectionResult,
|
||||||
PathFinder, PathFindingResult,
|
PathFinder, PathFindingResult,
|
||||||
FacetedSearch, AvailableFacets, FacetFilters,
|
FacetedSearch, AvailableFacets, FacetFilters,
|
||||||
|
|||||||
@@ -5,8 +5,8 @@
|
|||||||
use actix_web::{web, HttpRequest, HttpResponse};
|
use actix_web::{web, HttpRequest, HttpResponse};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use crate::query::{
|
use crate::query::{
|
||||||
EntityLinker, InferenceEngine, Summarizer,
|
EntityLinker, InferenceEngine, QueryReasoner, Summarizer,
|
||||||
SummarizationStrategy,
|
SummarizationStrategy, MentionLink,
|
||||||
};
|
};
|
||||||
use crate::handlers::response_builder;
|
use crate::handlers::response_builder;
|
||||||
use tracing::{debug, info, error};
|
use tracing::{debug, info, error};
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ use crate::query::visualize_types::{VisualizeRequest, VisualizeResponse, ReactFl
|
|||||||
use crate::query::bfs_graph_traversal::BfsConfig;
|
use crate::query::bfs_graph_traversal::BfsConfig;
|
||||||
use crate::query::force_directed_layout::ForceDirectedLayout;
|
use crate::query::force_directed_layout::ForceDirectedLayout;
|
||||||
use crate::http_server::AppState;
|
use crate::http_server::AppState;
|
||||||
|
use crate::jwt_validator::JwtValidator;
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
|||||||
@@ -6,7 +6,9 @@
|
|||||||
use actix_web::{web, HttpRequest, HttpResponse};
|
use actix_web::{web, HttpRequest, HttpResponse};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use crate::query::visualize_types::VisualizeRequest;
|
use tokio::sync::mpsc;
|
||||||
|
use futures_util::stream::{self, StreamExt};
|
||||||
|
use crate::query::visualize_types::{VisualizeRequest, ReactFlowNode, ReactFlowEdge, NodeData, EdgeData, NodeStyle};
|
||||||
use crate::query::bfs_graph_traversal::BfsConfig;
|
use crate::query::bfs_graph_traversal::BfsConfig;
|
||||||
use crate::query::force_directed_layout::ForceDirectedLayout;
|
use crate::query::force_directed_layout::ForceDirectedLayout;
|
||||||
use crate::http_server::AppState;
|
use crate::http_server::AppState;
|
||||||
|
|||||||
+348
-107
@@ -7,47 +7,21 @@ use serde_json::json;
|
|||||||
use sqlx::PgPool;
|
use sqlx::PgPool;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
|
use crate::endpoints::IngestRequest;
|
||||||
use crate::ingest_worker::IngestWorker;
|
use crate::ingest_worker::IngestWorker;
|
||||||
use serde::Deserialize;
|
use crate::query_worker::QueryWorker;
|
||||||
|
use crate::rate_limiter::{RateLimiter, LimitConfig};
|
||||||
/// JWT claims structure (extracted from deleted jwt_validator module)
|
use crate::idempotency::IdempotencyStore;
|
||||||
/// Will be replaced by riotpiao-rust-sdk claims (issue #56)
|
use crate::jwt_validator::{JwtValidator, JwtClaims};
|
||||||
#[derive(Debug, Clone, serde::Serialize, Deserialize)]
|
use crate::opensearch_client::{OpenSearchClient, HybridWeights};
|
||||||
pub struct JwtClaims {
|
use crate::dual_write_indexer::DualWriteIndexer;
|
||||||
pub sub: String,
|
use crate::gateway_queue_adapter::GatewayQueueAdapter;
|
||||||
pub iss: String,
|
use crate::queue_worker::{QueueWorker, QueueWorkerConfig};
|
||||||
pub aud: String,
|
use crate::queue_adapter::QueueAdapter;
|
||||||
pub exp: i64,
|
|
||||||
pub iat: i64,
|
|
||||||
pub nbf: Option<i64>,
|
|
||||||
pub permissions: Option<Vec<String>>,
|
|
||||||
pub groups: Option<Vec<String>>,
|
|
||||||
pub roles: Option<Vec<String>>,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Ingest request body
|
|
||||||
#[derive(Debug, Clone, Deserialize)]
|
|
||||||
pub struct IngestRequest {
|
|
||||||
pub project: String,
|
|
||||||
pub source: String,
|
|
||||||
pub ingest_id: String,
|
|
||||||
pub records: Vec<IngestRecord>,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize)]
|
|
||||||
pub struct IngestRecord {
|
|
||||||
pub text: String,
|
|
||||||
#[serde(default)]
|
|
||||||
pub role: Option<String>,
|
|
||||||
#[serde(default)]
|
|
||||||
pub timestamp: Option<String>,
|
|
||||||
#[serde(default)]
|
|
||||||
pub source_position: Option<i32>,
|
|
||||||
}
|
|
||||||
// RBAC removed for MVP - will add after core ingest/query working
|
// RBAC removed for MVP - will add after core ingest/query working
|
||||||
use crate::handlers::{
|
use crate::handlers::{
|
||||||
QueryParams,
|
QueryParams, QueryParamsError, SearchMethod, build_search_response,
|
||||||
LearnParams, build_learn_response,
|
LearnParams, LearnParamsError, build_learn_response,
|
||||||
visualize_handler, visualize_stream_handler, compact_handler
|
visualize_handler, visualize_stream_handler, compact_handler
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -59,7 +33,12 @@ pub struct AppState {
|
|||||||
pub vector_store: Arc<VectorStore>,
|
pub vector_store: Arc<VectorStore>,
|
||||||
pub embeddings: Arc<EmbeddingsClient>,
|
pub embeddings: Arc<EmbeddingsClient>,
|
||||||
pub ingest_worker: Arc<IngestWorker>,
|
pub ingest_worker: Arc<IngestWorker>,
|
||||||
|
pub query_worker: Arc<QueryWorker>,
|
||||||
|
pub rate_limiter: Arc<RateLimiter>,
|
||||||
|
pub idempotency_store: Arc<IdempotencyStore>,
|
||||||
|
pub jwt_validator: Option<Arc<JwtValidator>>,
|
||||||
pub auth_mode: AuthMode,
|
pub auth_mode: AuthMode,
|
||||||
|
pub opensearch_client: Option<Arc<OpenSearchClient>>,
|
||||||
/// M3.8 Query Optimizer (optional, from environment)
|
/// M3.8 Query Optimizer (optional, from environment)
|
||||||
pub optimizer_service: Option<Arc<mem_core::optimizer::OptimizerService>>,
|
pub optimizer_service: Option<Arc<mem_core::optimizer::OptimizerService>>,
|
||||||
}
|
}
|
||||||
@@ -96,9 +75,12 @@ async fn validate_auth(req: &HttpRequest, state: &AppState) -> Result<(JwtClaims
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Validate JWT token from Authorization header
|
/// Validate JWT token from Authorization header
|
||||||
/// NOTE: Full JWT validation deferred to riotpiao-rust-sdk migration (issue #56).
|
async fn validate_jwt_token(req: &HttpRequest, state: &AppState) -> Result<(JwtClaims, String), HttpResponse> {
|
||||||
/// For now, extracts Bearer token and creates synthetic claims.
|
let validator = state
|
||||||
async fn validate_jwt_token(req: &HttpRequest, _state: &AppState) -> Result<(JwtClaims, String), HttpResponse> {
|
.jwt_validator
|
||||||
|
.as_ref()
|
||||||
|
.ok_or_else(|| HttpResponse::InternalServerError().json(json!({"error": "jwt_validator_not_configured"})))?;
|
||||||
|
|
||||||
let auth_header = req
|
let auth_header = req
|
||||||
.headers()
|
.headers()
|
||||||
.get("Authorization")
|
.get("Authorization")
|
||||||
@@ -111,28 +93,26 @@ async fn validate_jwt_token(req: &HttpRequest, _state: &AppState) -> Result<(Jwt
|
|||||||
})?
|
})?
|
||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
let token = auth_header
|
let token = crate::jwt_validator::JwtValidator::extract_bearer_token(&auth_header)
|
||||||
.strip_prefix("Bearer ")
|
.map_err(|_| {
|
||||||
.ok_or_else(|| {
|
|
||||||
HttpResponse::Unauthorized().json(json!({
|
HttpResponse::Unauthorized().json(json!({
|
||||||
"error": "unauthorized",
|
"error": "unauthorized",
|
||||||
"reason": "invalid Authorization header format, expected 'Bearer <token>'"
|
"reason": "invalid Authorization header format"
|
||||||
}))
|
}))
|
||||||
})?
|
})?
|
||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
// Synthetic claims — real JWT validation will come with riotpiao-rust-sdk
|
let claims = validator
|
||||||
let claims = JwtClaims {
|
.validate_token(&token)
|
||||||
sub: "jwt-user".to_string(),
|
.await
|
||||||
iss: "authentik".to_string(),
|
.map_err(|e| {
|
||||||
aud: "memory".to_string(),
|
tracing::warn!("JWT validation failed: {}", e);
|
||||||
exp: i64::MAX,
|
HttpResponse::Unauthorized().json(json!({
|
||||||
iat: chrono::Utc::now().timestamp(),
|
"error": "unauthorized",
|
||||||
nbf: None,
|
"reason": format!("JWT validation failed: {}", e)
|
||||||
permissions: Some(vec!["*".to_string()]),
|
}))
|
||||||
groups: None,
|
})?
|
||||||
roles: Some(vec!["admin".to_string()]),
|
.clone();
|
||||||
};
|
|
||||||
|
|
||||||
Ok((claims, token))
|
Ok((claims, token))
|
||||||
}
|
}
|
||||||
@@ -182,16 +162,29 @@ fn has_capability(claims: &JwtClaims, required_capability: &str) -> bool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Extract client identifier from claims for rate limiting
|
/// Extract client identifier from claims for rate limiting
|
||||||
#[allow(dead_code)]
|
|
||||||
fn extract_rate_limit_key(claims: &JwtClaims) -> String {
|
fn extract_rate_limit_key(claims: &JwtClaims) -> String {
|
||||||
// Use subject (user/service ID) as rate limit key
|
// Use subject (user/service ID) as rate limit key
|
||||||
claims.sub.clone()
|
claims.sub.clone()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Rate limit guard — stub until riotpiao-rust-sdk (issue #56)
|
/// Rate limit guard — call this in handlers to check rate limit
|
||||||
fn check_rate_limit(_claims: &JwtClaims, _state: &AppState, _endpoint: &str) -> Result<(), HttpResponse> {
|
fn check_rate_limit(claims: &JwtClaims, state: &AppState, endpoint: &str) -> Result<(), HttpResponse> {
|
||||||
// Rate limiting deferred to API gateway / riotpiao-rust-sdk
|
let key = extract_rate_limit_key(claims);
|
||||||
Ok(())
|
|
||||||
|
match state.rate_limiter.check(&key, endpoint) {
|
||||||
|
Ok(_) => Ok(()),
|
||||||
|
Err(rate_limit_err) => {
|
||||||
|
let retry_after = rate_limit_err.retry_after_seconds.to_string();
|
||||||
|
Err(HttpResponse::TooManyRequests()
|
||||||
|
.insert_header(("Retry-After", retry_after))
|
||||||
|
.json(json!({
|
||||||
|
"error": "rate_limit_exceeded",
|
||||||
|
"reason": rate_limit_err.reason.clone(),
|
||||||
|
"retry_after_seconds": rate_limit_err.retry_after_seconds,
|
||||||
|
"limit_window": format!("{}s", rate_limit_err.limit_window_secs),
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Start HTTP server with database initialization
|
/// Start HTTP server with database initialization
|
||||||
@@ -213,7 +206,35 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
|
|||||||
let vector_store = Arc::new(VectorStore::new(pool.clone()));
|
let vector_store = Arc::new(VectorStore::new(pool.clone()));
|
||||||
let embeddings = Arc::new(EmbeddingsClient::from_env()?);
|
let embeddings = Arc::new(EmbeddingsClient::from_env()?);
|
||||||
let ingest_worker = Arc::new(IngestWorker::new(pool.clone(), (*embeddings).clone()));
|
let ingest_worker = Arc::new(IngestWorker::new(pool.clone(), (*embeddings).clone()));
|
||||||
let _reranker = RerankClient::from_env()?;
|
let reranker = RerankClient::from_env()?;
|
||||||
|
let query_worker = Arc::new(QueryWorker::new(VectorStore::new(pool.clone()), (*embeddings).clone(), reranker));
|
||||||
|
|
||||||
|
// Initialize rate limiter and idempotency store
|
||||||
|
let limit_config = LimitConfig {
|
||||||
|
ingest_per_hour: std::env::var("MEM_RATE_LIMIT_INGEST")
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.parse().ok())
|
||||||
|
.unwrap_or(100.0),
|
||||||
|
query_per_hour: std::env::var("MEM_RATE_LIMIT_QUERY")
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.parse().ok())
|
||||||
|
.unwrap_or(1000.0),
|
||||||
|
projects_per_hour: std::env::var("MEM_RATE_LIMIT_PROJECTS")
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.parse().ok())
|
||||||
|
.unwrap_or(100.0),
|
||||||
|
burst_per_second: std::env::var("MEM_RATE_LIMIT_BURST")
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.parse().ok())
|
||||||
|
.unwrap_or(10.0),
|
||||||
|
};
|
||||||
|
let rate_limiter = Arc::new(RateLimiter::new(limit_config));
|
||||||
|
|
||||||
|
let idempotency_ttl = std::env::var("MEM_IDEMPOTENCY_TTL_SECS")
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.parse().ok())
|
||||||
|
.unwrap_or(86400); // 24 hours default
|
||||||
|
let idempotency_store = Arc::new(IdempotencyStore::new(idempotency_ttl));
|
||||||
|
|
||||||
// Determine auth mode
|
// Determine auth mode
|
||||||
let auth_mode = std::env::var("MEM_AUTH_MODE")
|
let auth_mode = std::env::var("MEM_AUTH_MODE")
|
||||||
@@ -229,10 +250,38 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
// JWT auth will be handled by riotpiao-rust-sdk (issue #56)
|
// Setup JWT validator if in JWT mode
|
||||||
if matches!(auth_mode, AuthMode::Jwt) {
|
let jwt_validator = if matches!(auth_mode, AuthMode::Jwt) {
|
||||||
tracing::warn!("JWT auth mode selected but JwtValidator removed. Use riotpiao-rust-sdk (issue #56).");
|
let issuer = std::env::var("AUTHENTIK_ISSUER").map_err(|e| {
|
||||||
}
|
anyhow::anyhow!("AUTHENTIK_ISSUER env var required for JWT auth: {}", e)
|
||||||
|
})?;
|
||||||
|
let audience = std::env::var("AUTHENTIK_AUDIENCE").map_err(|e| {
|
||||||
|
anyhow::anyhow!("AUTHENTIK_AUDIENCE env var required for JWT auth: {}", e)
|
||||||
|
})?;
|
||||||
|
let cache_ttl = std::env::var("JWT_CACHE_TTL_SECS")
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.parse().ok())
|
||||||
|
.unwrap_or(3600); // 1 hour default
|
||||||
|
Some(Arc::new(crate::jwt_validator::JwtValidator::new(
|
||||||
|
issuer,
|
||||||
|
audience,
|
||||||
|
cache_ttl,
|
||||||
|
)))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
// Initialize OpenSearch client if configured
|
||||||
|
let opensearch_client = if let Ok(hosts_str) = std::env::var("OPENSEARCH_HOSTS") {
|
||||||
|
let hosts: Vec<String> = hosts_str
|
||||||
|
.split(',')
|
||||||
|
.map(|h| h.trim().to_string())
|
||||||
|
.collect();
|
||||||
|
Some(Arc::new(OpenSearchClient::new(hosts)))
|
||||||
|
} else {
|
||||||
|
tracing::warn!("OPENSEARCH_HOSTS not set, hybrid search disabled");
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
// Initialize M3.8 Query Optimizer if enabled
|
// Initialize M3.8 Query Optimizer if enabled
|
||||||
let optimizer_service = match mem_core::optimizer::OptimizerServiceBuilder::new().build() {
|
let optimizer_service = match mem_core::optimizer::OptimizerServiceBuilder::new().build() {
|
||||||
@@ -246,7 +295,67 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
// Queue adapter + dual-write will use riotpiao-rust-sdk (issue #56)
|
// Initialize M8.2 Queue Adapter and Dual-Write Indexer
|
||||||
|
let queue_adapter: Arc<dyn QueueAdapter> = if let Ok(gateway_url) = std::env::var("GATEWAY_URL") {
|
||||||
|
let adapter = GatewayQueueAdapter::with_authentik(
|
||||||
|
gateway_url,
|
||||||
|
std::env::var("AUTHENTIK_ISSUER").unwrap_or_default(),
|
||||||
|
std::env::var("AUTHENTIK_CLIENT_ID").unwrap_or_default(),
|
||||||
|
std::env::var("AUTHENTIK_CLIENT_SECRET").unwrap_or_default(),
|
||||||
|
);
|
||||||
|
tracing::info!("M8.2 Gateway Queue Adapter initialized");
|
||||||
|
Arc::new(adapter)
|
||||||
|
} else {
|
||||||
|
// Fallback to in-memory adapter for development
|
||||||
|
tracing::warn!("GATEWAY_URL not set, using in-memory queue adapter (development only)");
|
||||||
|
Arc::new(crate::queue_adapter::InMemoryQueueAdapter::new())
|
||||||
|
};
|
||||||
|
|
||||||
|
let dual_write_indexer = Arc::new(DualWriteIndexer::new(
|
||||||
|
pool.clone(),
|
||||||
|
opensearch_client.clone(),
|
||||||
|
queue_adapter.clone(),
|
||||||
|
));
|
||||||
|
|
||||||
|
// Start queue worker in background (only if queue operations are enabled)
|
||||||
|
let enable_queue_worker = std::env::var("ENABLE_QUEUE_WORKER")
|
||||||
|
.unwrap_or_else(|_| "true".to_string())
|
||||||
|
.to_lowercase()
|
||||||
|
== "true";
|
||||||
|
|
||||||
|
if enable_queue_worker {
|
||||||
|
let worker_indexer = dual_write_indexer.clone();
|
||||||
|
let worker_embeddings = embeddings.clone();
|
||||||
|
let worker_config = QueueWorkerConfig {
|
||||||
|
max_messages_per_batch: std::env::var("QUEUE_BATCH_SIZE")
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.parse().ok())
|
||||||
|
.unwrap_or(10),
|
||||||
|
visibility_timeout_secs: std::env::var("QUEUE_VISIBILITY_TIMEOUT")
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.parse().ok())
|
||||||
|
.unwrap_or(300),
|
||||||
|
wait_time_secs: std::env::var("QUEUE_WAIT_TIME")
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.parse().ok())
|
||||||
|
.unwrap_or(20),
|
||||||
|
project: std::env::var("QUEUE_PROJECT").ok(),
|
||||||
|
max_retries: std::env::var("QUEUE_MAX_RETRIES")
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.parse().ok())
|
||||||
|
.unwrap_or(3),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let worker = QueueWorker::new(worker_indexer, worker_embeddings, worker_config);
|
||||||
|
if let Err(e) = worker.start().await {
|
||||||
|
tracing::error!("Queue worker error: {}", e);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
tracing::info!("M8.2 Queue Worker started (background task)");
|
||||||
|
}
|
||||||
|
|
||||||
let state = web::Data::new(AppState {
|
let state = web::Data::new(AppState {
|
||||||
api_key,
|
api_key,
|
||||||
@@ -255,7 +364,12 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
|
|||||||
vector_store,
|
vector_store,
|
||||||
embeddings,
|
embeddings,
|
||||||
ingest_worker,
|
ingest_worker,
|
||||||
|
query_worker,
|
||||||
|
rate_limiter,
|
||||||
|
idempotency_store,
|
||||||
|
jwt_validator,
|
||||||
auth_mode,
|
auth_mode,
|
||||||
|
opensearch_client,
|
||||||
optimizer_service,
|
optimizer_service,
|
||||||
});
|
});
|
||||||
|
|
||||||
@@ -292,7 +406,6 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
|
|||||||
.app_data(state.clone())
|
.app_data(state.clone())
|
||||||
.wrap(Logger::default())
|
.wrap(Logger::default())
|
||||||
.route("/health", web::get().to(health_check))
|
.route("/health", web::get().to(health_check))
|
||||||
.route("/ready", web::get().to(readiness_check))
|
|
||||||
.route("/metrics", web::get().to(crate::metrics::metrics_handler))
|
.route("/metrics", web::get().to(crate::metrics::metrics_handler))
|
||||||
.route("/memory/ingest", web::post().to(ingest_handler))
|
.route("/memory/ingest", web::post().to(ingest_handler))
|
||||||
.route("/memory/ingest/{ingest_id}", web::get().to(ingest_status))
|
.route("/memory/ingest/{ingest_id}", web::get().to(ingest_status))
|
||||||
@@ -301,7 +414,7 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
|
|||||||
.route("/memory/query/semantic/entities", web::post().to(crate::handlers::semantic::search_entities_handler))
|
.route("/memory/query/semantic/entities", web::post().to(crate::handlers::semantic::search_entities_handler))
|
||||||
.route("/memory/query/semantic/edges", web::post().to(crate::handlers::semantic::search_edges_handler))
|
.route("/memory/query/semantic/edges", web::post().to(crate::handlers::semantic::search_edges_handler))
|
||||||
.route("/memory/query/hybrid", web::post().to(crate::handlers::semantic::hybrid_search_handler))
|
.route("/memory/query/hybrid", web::post().to(crate::handlers::semantic::hybrid_search_handler))
|
||||||
// context_handler removed — will be reimplemented with riotpiao-rust-sdk (issue #56)
|
.route("/memory/context", web::post().to(context_handler))
|
||||||
.route("/memory/projects", web::get().to(projects_handler))
|
.route("/memory/projects", web::get().to(projects_handler))
|
||||||
.route("/memory/skills", web::get().to(skills_handler))
|
.route("/memory/skills", web::get().to(skills_handler))
|
||||||
.route("/memory/learn", web::post().to(learn_handler))
|
.route("/memory/learn", web::post().to(learn_handler))
|
||||||
@@ -327,9 +440,9 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
|
|||||||
.route("/agents/{id}", web::put().to(crate::handlers::agent_handler::update_agent_handler))
|
.route("/agents/{id}", web::put().to(crate::handlers::agent_handler::update_agent_handler))
|
||||||
.route("/agents/{id}", web::delete().to(crate::handlers::agent_handler::delete_agent_handler))
|
.route("/agents/{id}", web::delete().to(crate::handlers::agent_handler::delete_agent_handler))
|
||||||
.route("/agents/{id}/metrics", web::get().to(crate::handlers::agent_handler::get_agent_metrics_handler))
|
.route("/agents/{id}/metrics", web::get().to(crate::handlers::agent_handler::get_agent_metrics_handler))
|
||||||
.route("/agents/{id}/prompts", web::post().to(crate::handlers::agent_handler::create_prompt_handler))
|
.route("/memory/agents/{project_id}/prompts", web::post().to(crate::handlers::agent_handler::create_prompt_handler))
|
||||||
.route("/agents/{id}/roles", web::post().to(crate::handlers::agent_handler::map_role_to_prompt_handler))
|
.route("/memory/agents/{project_id}/roles", web::post().to(crate::handlers::agent_handler::map_role_to_prompt_handler))
|
||||||
.route("/agents/{id}/roles/{role_name}/prompts", web::get().to(crate::handlers::agent_handler::get_role_prompts_handler))
|
.route("/memory/agents/{project_id}/roles/{role_name}/prompts", web::get().to(crate::handlers::agent_handler::get_role_prompts_handler))
|
||||||
});
|
});
|
||||||
|
|
||||||
tracing::info!("HttpServer instance created, binding to 0.0.0.0:{}", port);
|
tracing::info!("HttpServer instance created, binding to 0.0.0.0:{}", port);
|
||||||
@@ -364,24 +477,6 @@ pub async fn health_check(state: web::Data<AppState>) -> HttpResponse {
|
|||||||
HttpResponse::Ok().json(json!({"status": "ok", "uptime_seconds": uptime}))
|
HttpResponse::Ok().json(json!({"status": "ok", "uptime_seconds": uptime}))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// GET /ready — readiness probe (checks DB)
|
|
||||||
pub async fn readiness_check(state: web::Data<AppState>) -> HttpResponse {
|
|
||||||
let uptime = state.start_time.elapsed().as_secs();
|
|
||||||
let db_start = std::time::Instant::now();
|
|
||||||
match sqlx::query("SELECT 1").execute(&state.pool).await {
|
|
||||||
Ok(_) => {
|
|
||||||
crate::metrics::DEP_DB_UP.set(1);
|
|
||||||
crate::metrics::DEP_DB_LATENCY.observe(db_start.elapsed().as_secs_f64());
|
|
||||||
HttpResponse::Ok().json(json!({"status": "ready", "uptime_seconds": uptime, "db": "ok"}))
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
crate::metrics::DEP_DB_UP.set(0);
|
|
||||||
crate::metrics::HEALTH_CHECK_FAILURES.inc();
|
|
||||||
HttpResponse::ServiceUnavailable().json(json!({"status": "not_ready", "uptime_seconds": uptime, "db": format!("error: {}", e)}))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// POST /memory/ingest — queue an ingest job
|
/// POST /memory/ingest — queue an ingest job
|
||||||
pub async fn ingest_handler(
|
pub async fn ingest_handler(
|
||||||
req: HttpRequest,
|
req: HttpRequest,
|
||||||
@@ -405,7 +500,7 @@ pub async fn ingest_handler(
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
let _user_id = &claims.sub;
|
let user_id = &claims.sub;
|
||||||
if !has_capability(&claims, "memory:write") {
|
if !has_capability(&claims, "memory:write") {
|
||||||
INGEST_AUTH_FAILURES.inc();
|
INGEST_AUTH_FAILURES.inc();
|
||||||
INGEST_ERRORS_TOTAL.inc();
|
INGEST_ERRORS_TOTAL.inc();
|
||||||
@@ -423,8 +518,13 @@ pub async fn ingest_handler(
|
|||||||
return e;
|
return e;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Idempotency check via DB (ingest_id is UNIQUE)
|
// Check idempotency
|
||||||
// In-memory idempotency store removed; DB ON CONFLICT handles dedup
|
if let Some(cached) = state.idempotency_store.get(&body.ingest_id) {
|
||||||
|
tracing::info!("Returning cached response for ingest_id: {}", body.ingest_id);
|
||||||
|
INGEST_DUPLICATES_TOTAL.inc();
|
||||||
|
INGEST_IN_FLIGHT.dec();
|
||||||
|
return HttpResponse::Accepted().json(cached);
|
||||||
|
}
|
||||||
|
|
||||||
let byte_count: usize = body.records.iter().map(|r| r.text.len()).sum();
|
let byte_count: usize = body.records.iter().map(|r| r.text.len()).sum();
|
||||||
INGEST_BYTES_TOTAL.inc_by(byte_count as u64);
|
INGEST_BYTES_TOTAL.inc_by(byte_count as u64);
|
||||||
@@ -488,10 +588,12 @@ async fn execute_ingest(
|
|||||||
tracing::error!("Ingest failed: {}", e);
|
tracing::error!("Ingest failed: {}", e);
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
state.idempotency_store.set(body.ingest_id.clone(), response.clone());
|
||||||
HttpResponse::Accepted().json(response)
|
HttpResponse::Accepted().json(response)
|
||||||
}
|
}
|
||||||
Ok(None) => {
|
Ok(None) => {
|
||||||
// Already exists (concurrent insert — DB UNIQUE constraint)
|
// Already exists (concurrent insert)
|
||||||
|
state.idempotency_store.set(body.ingest_id.clone(), response.clone());
|
||||||
HttpResponse::Accepted().json(response)
|
HttpResponse::Accepted().json(response)
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
@@ -547,6 +649,56 @@ pub async fn ingest_status(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// M3.8: Optimize search results using pluggable OptimizerService
|
||||||
|
///
|
||||||
|
/// If optimizer_service is available, optimizes chunk text before returning.
|
||||||
|
/// Gracefully falls back to original on any error.
|
||||||
|
///
|
||||||
|
/// For LLM integration, use build_cache_aligned_async from PromptBuilder:
|
||||||
|
/// ```ignore
|
||||||
|
/// let msgs = PromptBuilder::build_cache_aligned_async(
|
||||||
|
/// &query,
|
||||||
|
/// previous_memory.as_deref(),
|
||||||
|
/// &chunk,
|
||||||
|
/// &optimizer_service,
|
||||||
|
/// ).await?;
|
||||||
|
/// ```
|
||||||
|
async fn optimize_search_results(
|
||||||
|
mut results: Vec<crate::query_worker::QueryResult>,
|
||||||
|
optimizer: Option<&Arc<mem_core::optimizer::OptimizerService>>,
|
||||||
|
) -> Vec<crate::query_worker::QueryResult> {
|
||||||
|
if optimizer.is_none() {
|
||||||
|
return results; // Optimizer not enabled, return as-is
|
||||||
|
}
|
||||||
|
|
||||||
|
let svc = optimizer.unwrap();
|
||||||
|
let mut optimized = Vec::new();
|
||||||
|
|
||||||
|
for mut result in results {
|
||||||
|
match svc.optimize(&result.text, "text/plain", Some("raw")).await {
|
||||||
|
Ok(optimized_bytes) => {
|
||||||
|
if let Ok(optimized_text) = String::from_utf8(optimized_bytes) {
|
||||||
|
let orig_len = result.text.len();
|
||||||
|
let opt_len = optimized_text.len();
|
||||||
|
result.text = optimized_text;
|
||||||
|
tracing::debug!(
|
||||||
|
"M3.8 optimized chunk: {} bytes → {} bytes ({:.1}% compression)",
|
||||||
|
orig_len,
|
||||||
|
opt_len,
|
||||||
|
(opt_len as f32 / orig_len as f32) * 100.0
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
// Graceful fallback: use original on optimization error
|
||||||
|
tracing::warn!("M3.8 optimization failed, using original: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
optimized.push(result);
|
||||||
|
}
|
||||||
|
|
||||||
|
optimized
|
||||||
|
}
|
||||||
|
|
||||||
/// POST /memory/learn — Ingest knowledge via gated loop (LLM evaluates + compacts)
|
/// POST /memory/learn — Ingest knowledge via gated loop (LLM evaluates + compacts)
|
||||||
///
|
///
|
||||||
@@ -779,6 +931,40 @@ pub async fn query_handler(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Execute hybrid search with OpenSearch fallback
|
||||||
|
async fn execute_hybrid_search(
|
||||||
|
state: &web::Data<AppState>,
|
||||||
|
params: &QueryParams,
|
||||||
|
results: Vec<crate::query_worker::QueryResult>,
|
||||||
|
token: &str,
|
||||||
|
) -> HttpResponse {
|
||||||
|
let Some(os_client) = &state.opensearch_client else {
|
||||||
|
tracing::info!("OpenSearch not configured, using semantic search only");
|
||||||
|
return build_search_response(params, results, Some("semantic_only"));
|
||||||
|
};
|
||||||
|
|
||||||
|
let sem_results: Vec<(String, f32, String, String, Vec<String>)> = results
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, r)| (
|
||||||
|
format!("sem-{}", i),
|
||||||
|
r.score,
|
||||||
|
r.text.clone(),
|
||||||
|
r.source.clone().unwrap_or_default(),
|
||||||
|
r.provenance.clone(),
|
||||||
|
))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let weights = HybridWeights { semantic: 0.6, lexical: 0.4 };
|
||||||
|
|
||||||
|
match os_client.hybrid_search(¶ms.question, sem_results, token, params.limit as usize, &weights).await {
|
||||||
|
Ok(_) => build_search_response(params, results, Some("hybrid")),
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!("Hybrid search failed, falling back to semantic: {}", e);
|
||||||
|
build_search_response(params, results, Some("semantic_fallback"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// GET /memory/projects — list projects with memory
|
/// GET /memory/projects — list projects with memory
|
||||||
pub async fn projects_handler(
|
pub async fn projects_handler(
|
||||||
@@ -869,6 +1055,69 @@ pub async fn skills_handler(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// POST /memory/context — three-tier context lookup for failure diagnosis
|
||||||
|
pub async fn context_handler(
|
||||||
|
req: HttpRequest,
|
||||||
|
body: web::Json<crate::context_endpoint::ContextRequest>,
|
||||||
|
state: web::Data<AppState>,
|
||||||
|
) -> HttpResponse {
|
||||||
|
use crate::metrics::*;
|
||||||
|
CONTEXT_REQUESTS_TOTAL.inc();
|
||||||
|
let _timer = Timer::new(&CONTEXT_DURATION);
|
||||||
|
|
||||||
|
let (claims, _token) = match validate_auth(&req, &state).await {
|
||||||
|
Ok(c) => c,
|
||||||
|
Err(e) => {
|
||||||
|
CONTEXT_ERRORS_TOTAL.inc();
|
||||||
|
ERROR_AUTH_FAILURE_CONTEXT.inc();
|
||||||
|
return e;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let user_id = &claims.sub;
|
||||||
|
if !has_capability(&claims, "memory:read") {
|
||||||
|
CONTEXT_ERRORS_TOTAL.inc();
|
||||||
|
ERROR_FORBIDDEN_CONTEXT.inc();
|
||||||
|
return HttpResponse::Forbidden().json(json!({
|
||||||
|
"error": "forbidden",
|
||||||
|
"reason": "missing capability: memory:read"
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Err(e) = check_rate_limit(&claims, &state, "/memory/context") {
|
||||||
|
return e;
|
||||||
|
}
|
||||||
|
|
||||||
|
let project = body.project.clone().unwrap_or_else(|| "all".to_string());
|
||||||
|
let scope = body.scope.clone().unwrap_or_else(|| "project".to_string());
|
||||||
|
let budget = body.budget.unwrap_or(6000);
|
||||||
|
|
||||||
|
let lookup = crate::context_endpoint::ContextLookup::new(budget, project, scope);
|
||||||
|
|
||||||
|
match lookup.lookup(body.into_inner()).await {
|
||||||
|
Ok(response) => {
|
||||||
|
tracing::info!(
|
||||||
|
tier = response.tier,
|
||||||
|
lessons = response.lessons.len(),
|
||||||
|
skills = response.skills.len(),
|
||||||
|
"context lookup successful"
|
||||||
|
);
|
||||||
|
// O3: Track tier hits
|
||||||
|
let total = response.lessons.len() + response.skills.len();
|
||||||
|
if total == 0 { CONTEXT_EMPTY_RESULTS.inc(); }
|
||||||
|
HttpResponse::Ok().json(response)
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
CONTEXT_ERRORS_TOTAL.inc();
|
||||||
|
ERROR_LOOKUP_FAILURE_CONTEXT.inc();
|
||||||
|
tracing::error!("context lookup error: {}", e);
|
||||||
|
HttpResponse::BadRequest().json(json!({
|
||||||
|
"error": "lookup_failed",
|
||||||
|
"reason": e.to_string()
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// POST /memory/vault/generate — generate Obsidian vault from memories
|
/// POST /memory/vault/generate — generate Obsidian vault from memories
|
||||||
pub async fn vault_generate_handler(
|
pub async fn vault_generate_handler(
|
||||||
@@ -1028,7 +1277,7 @@ pub async fn vault_browser_handler(
|
|||||||
/// Helper: Build file tree for a project
|
/// Helper: Build file tree for a project
|
||||||
async fn vault_project_tree(
|
async fn vault_project_tree(
|
||||||
project: &str,
|
project: &str,
|
||||||
_state: &web::Data<AppState>,
|
state: &web::Data<AppState>,
|
||||||
) -> HttpResponse {
|
) -> HttpResponse {
|
||||||
let vault_dir = std::env::var("MEM_HOME").unwrap_or_else(|_| "/data".to_string());
|
let vault_dir = std::env::var("MEM_HOME").unwrap_or_else(|_| "/data".to_string());
|
||||||
let project_path = format!("{}/vault/{}", vault_dir, project);
|
let project_path = format!("{}/vault/{}", vault_dir, project);
|
||||||
@@ -1203,20 +1452,12 @@ async fn query_temporal_graph(
|
|||||||
state: &web::Data<AppState>,
|
state: &web::Data<AppState>,
|
||||||
params: &QueryParams,
|
params: &QueryParams,
|
||||||
) -> anyhow::Result<serde_json::Value> {
|
) -> anyhow::Result<serde_json::Value> {
|
||||||
// Step 1: Find entities matching the question
|
// Step 1: Find entities (order by name for deterministic results)
|
||||||
// Use keyword search (ILIKE) on name + description for GET endpoint.
|
|
||||||
// POST /memory/query uses the full semantic retriever with embeddings.
|
|
||||||
let search_pattern = format!("%{}%", params.question);
|
|
||||||
let entities_rows: Vec<(String, String, String)> = sqlx::query_as(
|
let entities_rows: Vec<(String, String, String)> = sqlx::query_as(
|
||||||
"SELECT id::TEXT, name, entity_type FROM memory_entity \
|
"SELECT id, name, entity_type FROM memory_entity WHERE project_id = $1 LIMIT $2"
|
||||||
WHERE project_id = $1 AND t_expired IS NULL \
|
|
||||||
AND (name ILIKE $3 OR COALESCE(description, '') ILIKE $3 OR COALESCE(summary, '') ILIKE $3) \
|
|
||||||
ORDER BY confidence DESC \
|
|
||||||
LIMIT $2"
|
|
||||||
)
|
)
|
||||||
.bind(¶ms.project)
|
.bind(¶ms.project)
|
||||||
.bind(params.limit as i32)
|
.bind(params.limit as i32)
|
||||||
.bind(&search_pattern)
|
|
||||||
.fetch_all(&state.pool)
|
.fetch_all(&state.pool)
|
||||||
.await
|
.await
|
||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
@@ -1229,7 +1470,7 @@ async fn query_temporal_graph(
|
|||||||
for (entity_id, _name, _type_str) in &entities_rows {
|
for (entity_id, _name, _type_str) in &entities_rows {
|
||||||
let entity_edges: Vec<(String, String, String, String, f32, Option<chrono::DateTime<chrono::Utc>>, Option<chrono::DateTime<chrono::Utc>>)> =
|
let entity_edges: Vec<(String, String, String, String, f32, Option<chrono::DateTime<chrono::Utc>>, Option<chrono::DateTime<chrono::Utc>>)> =
|
||||||
sqlx::query_as(
|
sqlx::query_as(
|
||||||
"SELECT id::TEXT, target_id::TEXT, relation_type, fact, confidence, t_valid, t_invalid FROM memory_edge WHERE project_id = $1 AND source_id = $2::UUID"
|
"SELECT id, target_entity_id, relation_type, fact, confidence, t_valid, t_invalid FROM memory_edge WHERE project_id = $1 AND source_entity_id = $2"
|
||||||
)
|
)
|
||||||
.bind(¶ms.project)
|
.bind(¶ms.project)
|
||||||
.bind(entity_id)
|
.bind(entity_id)
|
||||||
|
|||||||
@@ -43,7 +43,7 @@ pub struct RankedCandidate {
|
|||||||
pub struct HybridRetriever {
|
pub struct HybridRetriever {
|
||||||
tfidf_scorer: Arc<mem_core::GlobalTfIdfScorer>,
|
tfidf_scorer: Arc<mem_core::GlobalTfIdfScorer>,
|
||||||
semantic_scorer: Arc<mem_core::SemanticScorer>,
|
semantic_scorer: Arc<mem_core::SemanticScorer>,
|
||||||
_pipeline: ScoringPipeline,
|
pipeline: ScoringPipeline,
|
||||||
min_tfidf_threshold: f32,
|
min_tfidf_threshold: f32,
|
||||||
prefilter_limit: usize,
|
prefilter_limit: usize,
|
||||||
rrf_tfidf_weight: f32,
|
rrf_tfidf_weight: f32,
|
||||||
@@ -62,7 +62,7 @@ impl HybridRetriever {
|
|||||||
Self {
|
Self {
|
||||||
tfidf_scorer,
|
tfidf_scorer,
|
||||||
semantic_scorer,
|
semantic_scorer,
|
||||||
_pipeline: pipeline,
|
pipeline,
|
||||||
min_tfidf_threshold: 0.3,
|
min_tfidf_threshold: 0.3,
|
||||||
prefilter_limit: 50,
|
prefilter_limit: 50,
|
||||||
rrf_tfidf_weight: 0.4,
|
rrf_tfidf_weight: 0.4,
|
||||||
@@ -71,7 +71,7 @@ impl HybridRetriever {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Decide retrieval route based on query and context
|
/// Decide retrieval route based on query and context
|
||||||
pub fn route_query(&self, _query: &str, has_wiki_scope: bool, is_reference_query: bool) -> RetrievalRoute {
|
pub fn route_query(&self, query: &str, has_wiki_scope: bool, is_reference_query: bool) -> RetrievalRoute {
|
||||||
if is_reference_query {
|
if is_reference_query {
|
||||||
RetrievalRoute::ReferenceOnly
|
RetrievalRoute::ReferenceOnly
|
||||||
} else if has_wiki_scope {
|
} else if has_wiki_scope {
|
||||||
|
|||||||
@@ -0,0 +1,129 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
use serde_json::json;
|
||||||
|
|
||||||
|
/// Cached ingest response with expiry
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
struct CachedResponse {
|
||||||
|
response: serde_json::Value,
|
||||||
|
inserted_at: Instant,
|
||||||
|
ttl: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CachedResponse {
|
||||||
|
fn is_expired(&self) -> bool {
|
||||||
|
self.inserted_at.elapsed() > self.ttl
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Idempotency store for ingest operations
|
||||||
|
pub struct IdempotencyStore {
|
||||||
|
cache: Arc<Mutex<HashMap<String, CachedResponse>>>,
|
||||||
|
ttl: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl IdempotencyStore {
|
||||||
|
pub fn new(ttl_seconds: u64) -> Self {
|
||||||
|
Self {
|
||||||
|
cache: Arc::new(Mutex::new(HashMap::new())),
|
||||||
|
ttl: Duration::from_secs(ttl_seconds),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get cached response for ingest_id. Returns None if not found or expired.
|
||||||
|
pub fn get(&self, ingest_id: &str) -> Option<serde_json::Value> {
|
||||||
|
let mut cache = self.cache.lock().unwrap();
|
||||||
|
|
||||||
|
if let Some(cached) = cache.get(ingest_id) {
|
||||||
|
if !cached.is_expired() {
|
||||||
|
return Some(cached.response.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clean up expired entry
|
||||||
|
cache.remove(ingest_id);
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Store response for ingest_id
|
||||||
|
pub fn set(&self, ingest_id: String, response: serde_json::Value) {
|
||||||
|
let mut cache = self.cache.lock().unwrap();
|
||||||
|
cache.insert(
|
||||||
|
ingest_id,
|
||||||
|
CachedResponse {
|
||||||
|
response,
|
||||||
|
inserted_at: Instant::now(),
|
||||||
|
ttl: self.ttl,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Evict expired entries (background maintenance)
|
||||||
|
pub fn evict_expired(&self) {
|
||||||
|
let mut cache = self.cache.lock().unwrap();
|
||||||
|
cache.retain(|_, v| !v.is_expired());
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Clear all entries (for testing)
|
||||||
|
#[cfg(test)]
|
||||||
|
pub fn clear(&self) {
|
||||||
|
let mut cache = self.cache.lock().unwrap();
|
||||||
|
cache.clear();
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get cache size (for testing)
|
||||||
|
#[cfg(test)]
|
||||||
|
pub fn len(&self) -> usize {
|
||||||
|
let cache = self.cache.lock().unwrap();
|
||||||
|
cache.len()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_idempotency_store_basic() {
|
||||||
|
let store = IdempotencyStore::new(60);
|
||||||
|
let response = json!({"ingest_id": "test-123", "status": "pending"});
|
||||||
|
|
||||||
|
store.set("test-123".to_string(), response.clone());
|
||||||
|
assert_eq!(store.get("test-123"), Some(response));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_idempotency_store_expiry() {
|
||||||
|
let store = IdempotencyStore::new(0);
|
||||||
|
let response = json!({"ingest_id": "test-123", "status": "pending"});
|
||||||
|
|
||||||
|
store.set("test-123".to_string(), response);
|
||||||
|
std::thread::sleep(Duration::from_millis(10));
|
||||||
|
|
||||||
|
assert_eq!(store.get("test-123"), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_idempotency_missing_key() {
|
||||||
|
let store = IdempotencyStore::new(60);
|
||||||
|
assert_eq!(store.get("nonexistent"), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_idempotency_evict_expired() {
|
||||||
|
let store = IdempotencyStore::new(1);
|
||||||
|
store.set("key1".to_string(), json!({"data": "value1"}));
|
||||||
|
store.set("key2".to_string(), json!({"data": "value2"}));
|
||||||
|
|
||||||
|
assert_eq!(store.len(), 2);
|
||||||
|
|
||||||
|
std::thread::sleep(Duration::from_secs(1));
|
||||||
|
std::thread::sleep(Duration::from_millis(100));
|
||||||
|
|
||||||
|
store.evict_expired();
|
||||||
|
assert_eq!(store.len(), 0);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,156 @@
|
|||||||
|
/// Ingest pipeline with DB persistence (Phase 2.6 integration)
|
||||||
|
///
|
||||||
|
/// Orchestrates:
|
||||||
|
/// 1. Run extraction pipeline
|
||||||
|
/// 2. Save entities to DB
|
||||||
|
/// 3. Save edges to DB
|
||||||
|
/// 4. Return extraction result + DB IDs
|
||||||
|
|
||||||
|
use anyhow::{Result, anyhow};
|
||||||
|
use mem_core::entity::Entity;
|
||||||
|
use mem_core::edge::Edge;
|
||||||
|
use mem_ingest::ingest_pipeline::{IngestPipeline, Episode, ExtractionResult};
|
||||||
|
use mem_store::db_repo::{PersistentEntityRepo, PersistentEdgeRepo, ReviewQueueRepo};
|
||||||
|
use sqlx::Pool;
|
||||||
|
use sqlx::postgres::Postgres;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tracing::{debug, error, info};
|
||||||
|
|
||||||
|
/// Ingest result with DB persistence
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct IngestWithDbResult {
|
||||||
|
pub episode_id: String,
|
||||||
|
pub entity_count: usize,
|
||||||
|
pub entity_ids: Vec<String>,
|
||||||
|
pub edge_count: usize,
|
||||||
|
pub edge_ids: Vec<String>,
|
||||||
|
pub contradiction_count: usize,
|
||||||
|
pub extraction_errors: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute ingest pipeline with DB persistence
|
||||||
|
pub async fn ingest_with_db_persistence(
|
||||||
|
pool: &Pool<Postgres>,
|
||||||
|
pipeline: &IngestPipeline,
|
||||||
|
episode: &Episode,
|
||||||
|
) -> Result<IngestWithDbResult> {
|
||||||
|
debug!("Starting ingest with DB persistence for episode: {}", episode.id);
|
||||||
|
|
||||||
|
// 1. Run extraction pipeline
|
||||||
|
let extraction = pipeline.ingest(episode).await?;
|
||||||
|
info!("Extraction complete: {} entities, {} edges, {} contradictions",
|
||||||
|
extraction.entities.len(),
|
||||||
|
extraction.edges.len(),
|
||||||
|
extraction.reviews.len()
|
||||||
|
);
|
||||||
|
|
||||||
|
// 2. Create repositories
|
||||||
|
let entity_repo = PersistentEntityRepo::new(pool.clone());
|
||||||
|
let edge_repo = PersistentEdgeRepo::new(pool.clone());
|
||||||
|
let review_queue_repo = ReviewQueueRepo::new(pool.clone());
|
||||||
|
|
||||||
|
let mut entity_ids = Vec::new();
|
||||||
|
let mut edge_ids = Vec::new();
|
||||||
|
let mut errors = Vec::new();
|
||||||
|
|
||||||
|
// 3. Save entities
|
||||||
|
for entity in &extraction.entities {
|
||||||
|
match entity_repo.save(entity).await {
|
||||||
|
Ok(id) => {
|
||||||
|
debug!("Saved entity: {} → {}", entity.name, id);
|
||||||
|
entity_ids.push(id);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
error!("Failed to save entity {}: {}", entity.name, e);
|
||||||
|
errors.push(format!("Entity save failed: {}", e));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 4. Save edges
|
||||||
|
for edge in &extraction.edges {
|
||||||
|
match edge_repo.save(edge).await {
|
||||||
|
Ok(id) => {
|
||||||
|
debug!("Saved edge: {} → {} ({})", edge.source_id, edge.target_id, id);
|
||||||
|
edge_ids.push(id);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
error!("Failed to save edge: {}", e);
|
||||||
|
errors.push(format!("Edge save failed: {}", e));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 5. Queue contradictions for review (only high-confidence)
|
||||||
|
for review_id in &extraction.reviews {
|
||||||
|
match review_queue_repo.enqueue(
|
||||||
|
&episode.project_id,
|
||||||
|
review_id,
|
||||||
|
"contradiction",
|
||||||
|
0.9,
|
||||||
|
).await {
|
||||||
|
Ok(_) => {
|
||||||
|
debug!("Queued contradiction for review: {}", review_id);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
error!("Failed to queue contradiction: {}", e);
|
||||||
|
errors.push(format!("Review queue failed: {}", e));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
info!("Ingest complete: saved {} entities, {} edges, {} contradictions, {} errors",
|
||||||
|
entity_ids.len(),
|
||||||
|
edge_ids.len(),
|
||||||
|
extraction.reviews.len(),
|
||||||
|
errors.len()
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(IngestWithDbResult {
|
||||||
|
episode_id: episode.id.clone(),
|
||||||
|
entity_count: entity_ids.len(),
|
||||||
|
entity_ids,
|
||||||
|
edge_count: edge_ids.len(),
|
||||||
|
edge_ids,
|
||||||
|
contradiction_count: extraction.reviews.len(),
|
||||||
|
extraction_errors: errors,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_ingest_with_db_result_creation() {
|
||||||
|
let result = IngestWithDbResult {
|
||||||
|
episode_id: "ep-1".to_string(),
|
||||||
|
entity_count: 2,
|
||||||
|
entity_ids: vec!["e1".to_string(), "e2".to_string()],
|
||||||
|
edge_count: 1,
|
||||||
|
edge_ids: vec!["edge-1".to_string()],
|
||||||
|
contradiction_count: 0,
|
||||||
|
extraction_errors: vec![],
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(result.entity_count, 2);
|
||||||
|
assert_eq!(result.edge_count, 1);
|
||||||
|
assert!(result.extraction_errors.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_ingest_with_db_result_errors() {
|
||||||
|
let result = IngestWithDbResult {
|
||||||
|
episode_id: "ep-1".to_string(),
|
||||||
|
entity_count: 1,
|
||||||
|
entity_ids: vec!["e1".to_string()],
|
||||||
|
edge_count: 0,
|
||||||
|
edge_ids: vec![],
|
||||||
|
contradiction_count: 0,
|
||||||
|
extraction_errors: vec!["DB connection failed".to_string()],
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(result.extraction_errors.len(), 1);
|
||||||
|
assert!(result.extraction_errors[0].contains("connection"));
|
||||||
|
}
|
||||||
|
}
|
||||||
+129
-314
@@ -1,5 +1,5 @@
|
|||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use mem_store::{VectorStore, ChunkL0};
|
use mem_store::{MemoryL1, VectorStore, ChunkL0, EntityRepoOps, EdgeRepoOps};
|
||||||
use mem_llm::EmbeddingsClient;
|
use mem_llm::EmbeddingsClient;
|
||||||
use mem_ingest::ingest_pipeline::{IngestPipeline, Episode};
|
use mem_ingest::ingest_pipeline::{IngestPipeline, Episode};
|
||||||
use mem_ingest::entity_extractor::{WikiLinkFallbackExtractor, LlmEntityExtractor};
|
use mem_ingest::entity_extractor::{WikiLinkFallbackExtractor, LlmEntityExtractor};
|
||||||
@@ -8,145 +8,22 @@ use mem_ingest::contradiction_detector::ContradictionHandler;
|
|||||||
use sqlx::PgPool;
|
use sqlx::PgPool;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
use pgvector::Vector;
|
||||||
/// Job status enumeration — type-safe alternative to magic strings
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
||||||
#[allow(dead_code)]
|
|
||||||
pub enum JobStatus {
|
|
||||||
Processing,
|
|
||||||
Done,
|
|
||||||
DoneWithErrors,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[allow(dead_code)]
|
|
||||||
impl JobStatus {
|
|
||||||
pub fn as_str(&self) -> &'static str {
|
|
||||||
match self {
|
|
||||||
JobStatus::Processing => "processing",
|
|
||||||
JobStatus::Done => "done",
|
|
||||||
JobStatus::DoneWithErrors => "done_with_errors",
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::fmt::Display for JobStatus {
|
|
||||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
||||||
write!(f, "{}", self.as_str())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
|
|
||||||
/// Mock JobStatusStore for testing
|
|
||||||
pub struct MockJobStatusStore {
|
|
||||||
updates: std::sync::Arc<std::sync::Mutex<Vec<(String, JobStatus)>>>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl MockJobStatusStore {
|
|
||||||
pub fn new() -> Self {
|
|
||||||
Self {
|
|
||||||
updates: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn updates(&self) -> Vec<(String, JobStatus)> {
|
|
||||||
self.updates.lock().unwrap().clone()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[async_trait::async_trait]
|
|
||||||
impl JobStatusStore for MockJobStatusStore {
|
|
||||||
async fn update_status(&self, ingest_id: &str, status: JobStatus) -> Result<()> {
|
|
||||||
self.updates.lock().unwrap().push((ingest_id.to_string(), status));
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Structured logging context for ingest operations — ensures consistent field names across all logs
|
|
||||||
#[derive(Debug, Clone)]
|
|
||||||
#[allow(dead_code)]
|
|
||||||
pub struct IngestLogContext {
|
|
||||||
pub ingest_id: String,
|
|
||||||
pub project: String,
|
|
||||||
pub record_id: String,
|
|
||||||
pub source: String,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[allow(dead_code)]
|
|
||||||
impl IngestLogContext {
|
|
||||||
fn new(ingest_id: &str, project: &str, record_id: &str, source: &str) -> Self {
|
|
||||||
Self {
|
|
||||||
ingest_id: ingest_id.to_string(),
|
|
||||||
project: project.to_string(),
|
|
||||||
record_id: record_id.to_string(),
|
|
||||||
source: source.to_string(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Job status store trait — abstracts database persistence of job status (enables mocking)
|
|
||||||
#[async_trait::async_trait]
|
|
||||||
#[allow(dead_code)]
|
|
||||||
pub trait JobStatusStore: Send + Sync {
|
|
||||||
/// Update job status in storage
|
|
||||||
async fn update_status(&self, ingest_id: &str, status: JobStatus) -> Result<()>;
|
|
||||||
}
|
|
||||||
|
|
||||||
/// PostgreSQL implementation of JobStatusStore
|
|
||||||
#[allow(dead_code)]
|
|
||||||
pub struct PgJobStatusStore {
|
|
||||||
pool: PgPool,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[allow(dead_code)]
|
|
||||||
impl PgJobStatusStore {
|
|
||||||
pub fn new(pool: PgPool) -> Self {
|
|
||||||
Self { pool }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[async_trait::async_trait]
|
|
||||||
impl JobStatusStore for PgJobStatusStore {
|
|
||||||
async fn update_status(&self, ingest_id: &str, status: JobStatus) -> Result<()> {
|
|
||||||
sqlx::query("UPDATE ingest_jobs SET status=$1, started_at=NOW() WHERE ingest_id=$2")
|
|
||||||
.bind(status.as_str())
|
|
||||||
.bind(ingest_id)
|
|
||||||
.execute(&self.pool)
|
|
||||||
.await?;
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
/// Ingest worker — processes queued records through entity/fact extraction pipeline
|
/// Ingest worker — processes queued records through entity/fact extraction pipeline
|
||||||
#[allow(dead_code)]
|
|
||||||
pub struct IngestWorker {
|
pub struct IngestWorker {
|
||||||
pool: PgPool,
|
pool: PgPool,
|
||||||
vector_store: Arc<VectorStore>,
|
vector_store: Arc<VectorStore>,
|
||||||
embeddings: Arc<EmbeddingsClient>,
|
embeddings: Arc<EmbeddingsClient>,
|
||||||
pipeline: Arc<IngestPipeline>,
|
pipeline: Arc<IngestPipeline>,
|
||||||
job_status_store: Arc<dyn JobStatusStore>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(dead_code)]
|
|
||||||
impl IngestWorker {
|
impl IngestWorker {
|
||||||
/// Create worker with full ingest pipeline
|
/// Create worker with full ingest pipeline
|
||||||
pub fn new(
|
pub fn new(
|
||||||
pool: PgPool,
|
pool: PgPool,
|
||||||
embeddings: EmbeddingsClient,
|
embeddings: EmbeddingsClient,
|
||||||
) -> Self {
|
|
||||||
let job_status_store = Arc::new(PgJobStatusStore::new(pool.clone()));
|
|
||||||
Self::with_job_store(pool, embeddings, job_status_store)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Create worker with custom job status store (for testing)
|
|
||||||
pub fn with_job_store(
|
|
||||||
pool: PgPool,
|
|
||||||
embeddings: EmbeddingsClient,
|
|
||||||
job_status_store: Arc<dyn JobStatusStore>,
|
|
||||||
) -> Self {
|
) -> Self {
|
||||||
let vector_store = Arc::new(VectorStore::new(pool.clone()));
|
let vector_store = Arc::new(VectorStore::new(pool.clone()));
|
||||||
|
|
||||||
@@ -181,17 +58,20 @@ impl IngestWorker {
|
|||||||
vector_store,
|
vector_store,
|
||||||
embeddings: Arc::new(embeddings),
|
embeddings: Arc::new(embeddings),
|
||||||
pipeline,
|
pipeline,
|
||||||
job_status_store,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Process ingest job with optional X-Forward-User auth header (API Gateway pattern)
|
/// Process ingest job: records -> entities/facts/edges via pipeline -> temporal storage
|
||||||
///
|
pub async fn process_ingest(
|
||||||
/// # Arguments
|
&self,
|
||||||
/// * `project` - Project ID for namespacing
|
project: &str,
|
||||||
/// * `ingest_id` - Unique ingest job ID
|
ingest_id: &str,
|
||||||
/// * `records` - Vec of (content, source) tuples
|
records: Vec<(String, String)>, // (content, source)
|
||||||
/// * `x_forward_user` - Optional X-Forward-User header from API Gateway (None for backward compat)
|
) -> Result<()> {
|
||||||
|
self.process_ingest_with_auth(project, ingest_id, records, None).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Process ingest with optional X-Forward-User auth header (API Gateway pattern)
|
||||||
pub async fn process_ingest_with_auth(
|
pub async fn process_ingest_with_auth(
|
||||||
&self,
|
&self,
|
||||||
project: &str,
|
project: &str,
|
||||||
@@ -208,8 +88,13 @@ impl IngestWorker {
|
|||||||
"Starting ingest job"
|
"Starting ingest job"
|
||||||
);
|
);
|
||||||
|
|
||||||
// Update job status to processing (via trait, testable)
|
// Update job status to processing
|
||||||
if let Err(e) = self.job_status_store.update_status(ingest_id, JobStatus::Processing).await {
|
if let Err(e) = sqlx::query("UPDATE ingest_jobs SET status=$1, started_at=NOW() WHERE ingest_id=$2")
|
||||||
|
.bind("processing")
|
||||||
|
.bind(ingest_id)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
{
|
||||||
tracing::error!(
|
tracing::error!(
|
||||||
target: "ingest",
|
target: "ingest",
|
||||||
error = %e,
|
error = %e,
|
||||||
@@ -222,16 +107,16 @@ impl IngestWorker {
|
|||||||
let mut total_entities = 0;
|
let mut total_entities = 0;
|
||||||
let mut total_edges = 0;
|
let mut total_edges = 0;
|
||||||
let mut total_reviews = 0;
|
let mut total_reviews = 0;
|
||||||
|
let mut extraction_errors = Vec::new();
|
||||||
|
let mut save_errors = Vec::new();
|
||||||
|
|
||||||
// Process each record through the ingest pipeline
|
// Process each record through the ingest pipeline
|
||||||
for (idx, (content, source)) in records.iter().enumerate() {
|
for (idx, (content, source)) in records.iter().enumerate() {
|
||||||
let record_id = format!("{}-{}", ingest_id, idx);
|
let record_id = format!("{}-{}", ingest_id, idx);
|
||||||
let log_ctx = IngestLogContext::new(ingest_id, project, &record_id, source);
|
|
||||||
|
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
target: "ingest",
|
target: "ingest",
|
||||||
record_id = %log_ctx.record_id,
|
record_id = %record_id,
|
||||||
source = %log_ctx.source,
|
source = source,
|
||||||
content_len = content.len(),
|
content_len = content.len(),
|
||||||
"Processing record"
|
"Processing record"
|
||||||
);
|
);
|
||||||
@@ -250,47 +135,95 @@ impl IngestWorker {
|
|||||||
Ok(result) => {
|
Ok(result) => {
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
target: "ingest",
|
target: "ingest",
|
||||||
record_id = %log_ctx.record_id,
|
record_id = %record_id,
|
||||||
entity_count = result.entities.len(),
|
entity_count = result.entities.len(),
|
||||||
edge_count = result.edges.len(),
|
edge_count = result.edges.len(),
|
||||||
review_count = result.reviews.len(),
|
review_count = result.reviews.len(),
|
||||||
"Pipeline extraction successful"
|
"Pipeline extraction successful"
|
||||||
);
|
);
|
||||||
|
|
||||||
// Save entities to database with embeddings (RAG-006)
|
// Save entities to database (normally via EntityRepo, using direct SQL for now)
|
||||||
for entity in &result.entities {
|
for entity in &result.entities {
|
||||||
match save_entity_with_embedding(&self.pool, &self.embeddings, entity, &log_ctx).await {
|
match save_entity_to_db(&self.pool, entity).await {
|
||||||
Ok(saved) => if saved { total_entities += 1; }
|
Ok(_) => {
|
||||||
Err(_) => { /* error already logged */ }
|
tracing::debug!(
|
||||||
|
target: "ingest",
|
||||||
|
record_id = %record_id,
|
||||||
|
entity_name = &entity.name,
|
||||||
|
entity_type = entity.entity_type.as_str(),
|
||||||
|
"Saved entity"
|
||||||
|
);
|
||||||
|
total_entities += 1;
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let msg = format!("Failed to save entity '{}': {}", entity.name, e);
|
||||||
|
tracing::warn!(
|
||||||
|
target: "ingest",
|
||||||
|
error = %e,
|
||||||
|
record_id = %record_id,
|
||||||
|
entity_name = &entity.name,
|
||||||
|
"Entity save failed"
|
||||||
|
);
|
||||||
|
save_errors.push(msg);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Save edges to database with embeddings (RAG-006)
|
// Save edges to database (normally via EdgeRepo, using direct SQL for now)
|
||||||
for edge in &result.edges {
|
for edge in &result.edges {
|
||||||
match save_edge_with_embedding(&self.pool, &self.embeddings, edge, &log_ctx).await {
|
match save_edge_to_db(&self.pool, edge).await {
|
||||||
Ok(saved) => if saved { total_edges += 1; }
|
Ok(_) => {
|
||||||
Err(_) => { /* error already logged */ }
|
tracing::debug!(
|
||||||
|
target: "ingest",
|
||||||
|
record_id = %record_id,
|
||||||
|
relation_type = &edge.relation_type,
|
||||||
|
"Saved edge"
|
||||||
|
);
|
||||||
|
total_edges += 1;
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let msg = format!("Failed to save edge: {}", e);
|
||||||
|
tracing::warn!(
|
||||||
|
target: "ingest",
|
||||||
|
error = %e,
|
||||||
|
record_id = %record_id,
|
||||||
|
"Edge save failed"
|
||||||
|
);
|
||||||
|
save_errors.push(msg);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
total_reviews += result.reviews.len();
|
total_reviews += result.reviews.len();
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
|
let msg = format!("Record {}: {}", record_id, e);
|
||||||
tracing::error!(
|
tracing::error!(
|
||||||
target: "ingest",
|
target: "ingest",
|
||||||
error = %e,
|
error = %e,
|
||||||
record_id = %log_ctx.record_id,
|
record_id = %record_id,
|
||||||
source = %log_ctx.source,
|
source = source,
|
||||||
"Pipeline extraction failed"
|
"Pipeline extraction failed"
|
||||||
);
|
);
|
||||||
// Continue processing other records (no error accumulation)
|
extraction_errors.push(msg);
|
||||||
|
// Continue processing other records
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Mark job complete (via trait, testable)
|
// Mark job complete
|
||||||
let final_status = JobStatus::Done;
|
let final_status = if extraction_errors.is_empty() && save_errors.is_empty() {
|
||||||
if let Err(e) = self.job_status_store.update_status(ingest_id, final_status).await {
|
"done"
|
||||||
|
} else {
|
||||||
|
"done_with_errors"
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Err(e) = sqlx::query("UPDATE ingest_jobs SET status=$1, completed_at=NOW() WHERE ingest_id=$2")
|
||||||
|
.bind(final_status)
|
||||||
|
.bind(ingest_id)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
{
|
||||||
tracing::error!(
|
tracing::error!(
|
||||||
target: "ingest",
|
target: "ingest",
|
||||||
error = %e,
|
error = %e,
|
||||||
@@ -307,16 +240,35 @@ impl IngestWorker {
|
|||||||
entities = total_entities,
|
entities = total_entities,
|
||||||
edges = total_edges,
|
edges = total_edges,
|
||||||
reviews = total_reviews,
|
reviews = total_reviews,
|
||||||
status = final_status.as_str(),
|
extraction_errors = extraction_errors.len(),
|
||||||
|
save_errors = save_errors.len(),
|
||||||
|
status = final_status,
|
||||||
"Ingest job completed"
|
"Ingest job completed"
|
||||||
);
|
);
|
||||||
|
|
||||||
|
if !extraction_errors.is_empty() {
|
||||||
|
tracing::warn!(
|
||||||
|
target: "ingest",
|
||||||
|
errors = ?extraction_errors,
|
||||||
|
ingest_id = ingest_id,
|
||||||
|
"Extraction errors occurred during ingest"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if !save_errors.is_empty() {
|
||||||
|
tracing::warn!(
|
||||||
|
target: "ingest",
|
||||||
|
errors = ?save_errors,
|
||||||
|
ingest_id = ingest_id,
|
||||||
|
"Save errors occurred during ingest"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Process a single chunk
|
/// Process a single chunk
|
||||||
pub async fn process_chunk(&self, project: &str, query_id: &str, content: &str, source: &str) -> Result<()> {
|
pub async fn process_chunk(&self, project: &str, query_id: &str, content: &str, source: &str) -> Result<()> {
|
||||||
let _embedding = self.embeddings.embed_one(content).await?;
|
let embedding = self.embeddings.embed_one(content).await?;
|
||||||
let chunk = ChunkL0 {
|
let chunk = ChunkL0 {
|
||||||
id: Uuid::new_v4(),
|
id: Uuid::new_v4(),
|
||||||
project: project.to_string(),
|
project: project.to_string(),
|
||||||
@@ -331,7 +283,6 @@ impl IngestWorker {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Extract wiki links from text (e.g., [[Kubernetes]] -> "Kubernetes")
|
/// Extract wiki links from text (e.g., [[Kubernetes]] -> "Kubernetes")
|
||||||
#[allow(dead_code)]
|
|
||||||
fn extract_wiki_links(text: &str) -> Vec<String> {
|
fn extract_wiki_links(text: &str) -> Vec<String> {
|
||||||
let mut links = Vec::new();
|
let mut links = Vec::new();
|
||||||
let mut chars = text.chars().peekable();
|
let mut chars = text.chars().peekable();
|
||||||
@@ -353,178 +304,41 @@ fn extract_wiki_links(text: &str) -> Vec<String> {
|
|||||||
links
|
links
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Save entity with logging — logs at debug level on success, warn on error
|
/// Save entity to database via raw SQL (normally would use EntityRepo trait)
|
||||||
/// Returns Ok(true) if saved, Ok(false) if skipped, Err if fatal error
|
async fn save_entity_to_db(pool: &PgPool, entity: &mem_core::entity::Entity) -> Result<()> {
|
||||||
/// Save entity with embeddings (RAG-006)
|
// Convert OffsetDateTime to PostgreSQL timestamp format
|
||||||
/// Embeds name + summary before persisting, so semantic search can find entities.
|
|
||||||
#[allow(dead_code)]
|
|
||||||
async fn save_entity_with_embedding(
|
|
||||||
pool: &PgPool,
|
|
||||||
embeddings: &EmbeddingsClient,
|
|
||||||
entity: &mem_core::entity::Entity,
|
|
||||||
log_ctx: &IngestLogContext,
|
|
||||||
) -> Result<bool> {
|
|
||||||
// Embed entity name
|
|
||||||
let name_embedding = match embeddings.embed_one(&entity.name).await {
|
|
||||||
Ok(emb) => Some(emb.to_vec()),
|
|
||||||
Err(e) => {
|
|
||||||
tracing::warn!(
|
|
||||||
target: "ingest",
|
|
||||||
error = %e,
|
|
||||||
entity_name = &entity.name,
|
|
||||||
"Name embedding failed, saving entity without name_embedding"
|
|
||||||
);
|
|
||||||
None
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
// Embed summary if present
|
|
||||||
let summary_embedding = if let Some(ref summary) = entity.summary {
|
|
||||||
match embeddings.embed_one(summary).await {
|
|
||||||
Ok(emb) => Some(emb.to_vec()),
|
|
||||||
Err(e) => {
|
|
||||||
tracing::debug!(target: "ingest", error = %e, "Summary embedding failed");
|
|
||||||
None
|
|
||||||
}
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
};
|
|
||||||
|
|
||||||
let t_created_str = entity.t_created.to_string();
|
let t_created_str = entity.t_created.to_string();
|
||||||
|
|
||||||
let result = sqlx::query(
|
sqlx::query(
|
||||||
"INSERT INTO memory_entity (id, project_id, name, entity_type, description, summary, \
|
"INSERT INTO memory_entity (id, project_id, name, entity_type, description, t_created, t_updated, confidence)
|
||||||
name_embedding, summary_embedding, t_created, t_updated, confidence) \
|
VALUES ($1, $2, $3, $4, $5, $6::TIMESTAMPTZ, $7::TIMESTAMPTZ, $8)
|
||||||
VALUES ($1::UUID, $2, $3, $4, $5, $6, $7, $8, $9::TIMESTAMPTZ, $10::TIMESTAMPTZ, $11) \
|
ON CONFLICT (project_id, name) DO UPDATE SET
|
||||||
ON CONFLICT (project_id, name) DO UPDATE SET \
|
entity_type = EXCLUDED.entity_type,
|
||||||
entity_type = EXCLUDED.entity_type, \
|
description = COALESCE(NULLIF(EXCLUDED.description, ''), memory_entity.description),
|
||||||
description = COALESCE(NULLIF(EXCLUDED.description, ''), memory_entity.description), \
|
t_updated = NOW(),
|
||||||
summary = COALESCE(NULLIF(EXCLUDED.summary, ''), memory_entity.summary), \
|
confidence = GREATEST(memory_entity.confidence, EXCLUDED.confidence),
|
||||||
name_embedding = COALESCE(EXCLUDED.name_embedding, memory_entity.name_embedding), \
|
|
||||||
summary_embedding = COALESCE(EXCLUDED.summary_embedding, memory_entity.summary_embedding), \
|
|
||||||
t_updated = NOW(), \
|
|
||||||
confidence = GREATEST(memory_entity.confidence, EXCLUDED.confidence), \
|
|
||||||
source_count = memory_entity.source_count + 1"
|
source_count = memory_entity.source_count + 1"
|
||||||
)
|
)
|
||||||
.bind(&entity.id)
|
.bind(&entity.id)
|
||||||
.bind(&entity.project_id)
|
.bind(&entity.project_id)
|
||||||
.bind(&entity.name)
|
.bind(&entity.name)
|
||||||
.bind(entity.entity_type.as_str())
|
.bind(entity.entity_type.as_str())
|
||||||
.bind(entity.summary.as_deref()) // description
|
.bind(entity.summary.as_deref())
|
||||||
.bind(entity.summary.as_deref()) // summary
|
|
||||||
.bind(name_embedding.as_deref())
|
|
||||||
.bind(summary_embedding.as_deref())
|
|
||||||
.bind(&t_created_str)
|
.bind(&t_created_str)
|
||||||
.bind(&t_created_str)
|
.bind(&t_created_str)
|
||||||
.bind(1.0_f32)
|
.bind(1.0_f32) // default confidence
|
||||||
.execute(pool)
|
.execute(pool)
|
||||||
.await;
|
.await?;
|
||||||
|
Ok(())
|
||||||
match result {
|
|
||||||
Ok(_) => {
|
|
||||||
tracing::debug!(
|
|
||||||
target: "ingest",
|
|
||||||
record_id = %log_ctx.record_id,
|
|
||||||
entity_name = &entity.name,
|
|
||||||
entity_type = entity.entity_type.as_str(),
|
|
||||||
has_name_emb = name_embedding.is_some(),
|
|
||||||
has_summary_emb = summary_embedding.is_some(),
|
|
||||||
"Saved entity with embeddings"
|
|
||||||
);
|
|
||||||
Ok(true)
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
tracing::warn!(
|
|
||||||
target: "ingest",
|
|
||||||
error = %e,
|
|
||||||
record_id = %log_ctx.record_id,
|
|
||||||
entity_name = &entity.name,
|
|
||||||
project = %log_ctx.project,
|
|
||||||
"Entity save failed"
|
|
||||||
);
|
|
||||||
Ok(false)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Save edge with fact embedding (RAG-006)
|
/// Save edge to database via raw SQL (normally would use EdgeRepo trait)
|
||||||
/// Embeds fact text before persisting, so semantic search can find edges.
|
/// NOTE: Production DB may have old schema. Gracefully skip if temporal columns missing.
|
||||||
#[allow(dead_code)]
|
|
||||||
async fn save_edge_with_embedding(
|
|
||||||
pool: &PgPool,
|
|
||||||
embeddings: &EmbeddingsClient,
|
|
||||||
edge: &mem_core::edge::Edge,
|
|
||||||
log_ctx: &IngestLogContext,
|
|
||||||
) -> Result<bool> {
|
|
||||||
// Embed the fact text
|
|
||||||
let fact_embedding = match embeddings.embed_one(&edge.fact).await {
|
|
||||||
Ok(emb) => Some(emb.to_vec()),
|
|
||||||
Err(e) => {
|
|
||||||
tracing::warn!(
|
|
||||||
target: "ingest",
|
|
||||||
error = %e,
|
|
||||||
fact = &edge.fact,
|
|
||||||
"Fact embedding failed, saving edge without fact_embedding"
|
|
||||||
);
|
|
||||||
None
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let result = sqlx::query(
|
|
||||||
"INSERT INTO memory_edge (id, project_id, source_id, target_id, relation_type, fact, \
|
|
||||||
fact_embedding, t_valid, t_invalid, t_created, confidence) \
|
|
||||||
VALUES ($1::UUID, $2, $3::UUID, $4::UUID, $5, $6, $7, $8::TIMESTAMPTZ, $9::TIMESTAMPTZ, $10::TIMESTAMPTZ, $11) \
|
|
||||||
ON CONFLICT (id) DO NOTHING"
|
|
||||||
)
|
|
||||||
.bind(&edge.id)
|
|
||||||
.bind(&edge.project_id)
|
|
||||||
.bind(&edge.source_entity_id)
|
|
||||||
.bind(&edge.target_entity_id)
|
|
||||||
.bind(&edge.relation_type)
|
|
||||||
.bind(&edge.fact)
|
|
||||||
.bind(fact_embedding.as_deref())
|
|
||||||
.bind(edge.t_valid.map(|t| t.to_string()))
|
|
||||||
.bind(edge.t_invalid.map(|t| t.to_string()))
|
|
||||||
.bind(edge.t_created.to_string())
|
|
||||||
.bind(edge.confidence)
|
|
||||||
.execute(pool)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
match result {
|
|
||||||
Ok(_) => {
|
|
||||||
tracing::debug!(
|
|
||||||
target: "ingest",
|
|
||||||
record_id = %log_ctx.record_id,
|
|
||||||
relation_type = &edge.relation_type,
|
|
||||||
source_entity = &edge.source_entity_id,
|
|
||||||
target_entity = &edge.target_entity_id,
|
|
||||||
has_fact_emb = fact_embedding.is_some(),
|
|
||||||
"Saved edge with embedding"
|
|
||||||
);
|
|
||||||
Ok(true)
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
tracing::warn!(
|
|
||||||
target: "ingest",
|
|
||||||
error = %e,
|
|
||||||
record_id = %log_ctx.record_id,
|
|
||||||
relation_type = &edge.relation_type,
|
|
||||||
project = %log_ctx.project,
|
|
||||||
"Edge save failed"
|
|
||||||
);
|
|
||||||
Ok(false)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Legacy save functions kept for backward compatibility but unused
|
|
||||||
#[allow(dead_code)]
|
|
||||||
#[allow(dead_code)]
|
|
||||||
async fn save_edge_to_db(pool: &PgPool, edge: &mem_core::edge::Edge) -> Result<()> {
|
async fn save_edge_to_db(pool: &PgPool, edge: &mem_core::edge::Edge) -> Result<()> {
|
||||||
|
// Try temporal schema first (id, project_id, source_entity_id, etc)
|
||||||
let result = sqlx::query(
|
let result = sqlx::query(
|
||||||
"INSERT INTO memory_edge (id, project_id, source_id, target_id, relation_type, fact, t_valid, t_invalid, t_created, confidence)
|
"INSERT INTO memory_edge (id, project_id, source_id, target_id, relation_type, fact, t_valid, t_invalid, t_created, confidence)
|
||||||
VALUES ($1::UUID, $2, $3::UUID, $4::UUID, $5, $6, $7::TIMESTAMPTZ, $8::TIMESTAMPTZ, $9::TIMESTAMPTZ, $10)
|
VALUES ($1, $2, $3, $4, $5, $6, $7::TIMESTAMPTZ, $8::TIMESTAMPTZ, $9::TIMESTAMPTZ, $10)
|
||||||
ON CONFLICT (id) DO NOTHING"
|
ON CONFLICT (id) DO NOTHING"
|
||||||
)
|
)
|
||||||
.bind(&edge.id)
|
.bind(&edge.id)
|
||||||
@@ -543,7 +357,8 @@ async fn save_edge_to_db(pool: &PgPool, edge: &mem_core::edge::Edge) -> Result<(
|
|||||||
match result {
|
match result {
|
||||||
Ok(_) => Ok(()),
|
Ok(_) => Ok(()),
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
tracing::debug!("Temporal edge schema not available: {}. Skipping edge save.", e);
|
tracing::debug!("Temporal edge schema not available: {}. Skipping edge save (will be available after schema migration).", e);
|
||||||
|
// This is expected if production DB hasn't migrated to temporal schema yet
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,208 @@
|
|||||||
|
use anyhow::{anyhow, Result};
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use jsonwebtoken::{decode, DecodingKey, TokenData, Validation, Algorithm};
|
||||||
|
use reqwest::Client;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::Mutex;
|
||||||
|
|
||||||
|
/// JWT claims from Authentik
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct JwtClaims {
|
||||||
|
pub sub: String,
|
||||||
|
pub iss: String,
|
||||||
|
pub aud: String,
|
||||||
|
pub exp: i64,
|
||||||
|
pub iat: i64,
|
||||||
|
pub nbf: Option<i64>,
|
||||||
|
pub permissions: Option<Vec<String>>,
|
||||||
|
pub groups: Option<Vec<String>>,
|
||||||
|
/// Roles from Authentik (for RBAC)
|
||||||
|
pub roles: Option<Vec<String>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// JWKS (JSON Web Key Set) response from Authentik
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct JwksResponse {
|
||||||
|
pub keys: Vec<JsonWebKey>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct JsonWebKey {
|
||||||
|
pub kty: String,
|
||||||
|
pub use_: Option<String>,
|
||||||
|
#[serde(rename = "kid")]
|
||||||
|
pub key_id: Option<String>,
|
||||||
|
pub n: Option<String>,
|
||||||
|
pub e: Option<String>,
|
||||||
|
pub alg: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// JWT validator with JWKS caching
|
||||||
|
pub struct JwtValidator {
|
||||||
|
pub issuer: String,
|
||||||
|
pub audience: String,
|
||||||
|
client: Client,
|
||||||
|
jwks_cache: Arc<Mutex<(Option<JwksResponse>, DateTime<Utc>)>>,
|
||||||
|
jwks_cache_ttl_secs: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl JwtValidator {
|
||||||
|
pub fn new(issuer: String, audience: String, jwks_cache_ttl_secs: i64) -> Self {
|
||||||
|
Self {
|
||||||
|
issuer,
|
||||||
|
audience,
|
||||||
|
client: Client::new(),
|
||||||
|
jwks_cache: Arc::new(Mutex::new((None, Utc::now()))),
|
||||||
|
jwks_cache_ttl_secs,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fetch JWKS from issuer discovery endpoint
|
||||||
|
async fn fetch_jwks(&self) -> Result<JwksResponse> {
|
||||||
|
let discovery_url = format!("{}/.well-known/openid-configuration", self.issuer);
|
||||||
|
tracing::debug!("Fetching OIDC discovery from {}", discovery_url);
|
||||||
|
|
||||||
|
let discovery: serde_json::Value = self
|
||||||
|
.client
|
||||||
|
.get(&discovery_url)
|
||||||
|
.send()
|
||||||
|
.await?
|
||||||
|
.json()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let jwks_uri = discovery
|
||||||
|
.get("jwks_uri")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.ok_or_else(|| anyhow!("No jwks_uri in discovery doc"))?;
|
||||||
|
|
||||||
|
tracing::debug!("Fetching JWKS from {}", jwks_uri);
|
||||||
|
let jwks: JwksResponse = self.client.get(jwks_uri).send().await?.json().await?;
|
||||||
|
|
||||||
|
if jwks.keys.is_empty() {
|
||||||
|
return Err(anyhow!("No keys in JWKS response"));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(jwks)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get JWKS from cache or fetch fresh
|
||||||
|
async fn get_jwks(&self) -> Result<JwksResponse> {
|
||||||
|
let cache = self.jwks_cache.lock().await;
|
||||||
|
let (cached_jwks, cached_at) = cache.clone();
|
||||||
|
|
||||||
|
// Check if cache is still valid
|
||||||
|
if let Some(jwks) = cached_jwks {
|
||||||
|
let age = (Utc::now() - cached_at).num_seconds();
|
||||||
|
if age < self.jwks_cache_ttl_secs {
|
||||||
|
drop(cache);
|
||||||
|
tracing::debug!("JWKS from cache (age: {}s)", age);
|
||||||
|
return Ok(jwks);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
drop(cache);
|
||||||
|
|
||||||
|
// Fetch fresh JWKS
|
||||||
|
let jwks = self.fetch_jwks().await?;
|
||||||
|
let mut cache = self.jwks_cache.lock().await;
|
||||||
|
*cache = (Some(jwks.clone()), Utc::now());
|
||||||
|
Ok(jwks)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convert JWKS key to DecodingKey for RS256 validation
|
||||||
|
fn jwks_to_decoding_key(key: &JsonWebKey) -> Result<DecodingKey> {
|
||||||
|
// Only support RSA keys
|
||||||
|
if key.kty != "RSA" {
|
||||||
|
return Err(anyhow!("Unsupported key type: {}", key.kty));
|
||||||
|
}
|
||||||
|
|
||||||
|
let n = key.n.as_ref().ok_or_else(|| anyhow!("Missing RSA modulus"))?;
|
||||||
|
let e = key.e.as_ref().ok_or_else(|| anyhow!("Missing RSA exponent"))?;
|
||||||
|
|
||||||
|
DecodingKey::from_rsa_components(n, e).map_err(|e| anyhow!("Invalid RSA key: {}", e))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Validate JWT token and extract claims
|
||||||
|
pub async fn validate_token(&self, token: &str) -> Result<JwtClaims> {
|
||||||
|
// Decode header to check algorithm
|
||||||
|
let header = jsonwebtoken::decode_header(token)
|
||||||
|
.map_err(|e| anyhow!("Invalid token header: {}", e))?;
|
||||||
|
|
||||||
|
// Pin to RS256 only (defense against algorithm confusion)
|
||||||
|
if header.alg != Algorithm::RS256 {
|
||||||
|
return Err(anyhow!(
|
||||||
|
"Invalid algorithm: {:?}, expected RS256",
|
||||||
|
header.alg
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let kid = header
|
||||||
|
.kid
|
||||||
|
.as_ref()
|
||||||
|
.ok_or_else(|| anyhow!("Token missing 'kid' header"))?;
|
||||||
|
|
||||||
|
// Fetch JWKS
|
||||||
|
let jwks = self.get_jwks().await?;
|
||||||
|
|
||||||
|
// Find key by kid
|
||||||
|
let key = jwks
|
||||||
|
.keys
|
||||||
|
.iter()
|
||||||
|
.find(|k| k.key_id.as_ref() == Some(kid))
|
||||||
|
.ok_or_else(|| anyhow!("Key not found in JWKS: {}", kid))?;
|
||||||
|
|
||||||
|
// Convert to DecodingKey
|
||||||
|
let decoding_key = Self::jwks_to_decoding_key(key)?;
|
||||||
|
|
||||||
|
// Validate token signature + claims
|
||||||
|
let mut validation = Validation::new(Algorithm::RS256);
|
||||||
|
validation.set_issuer(&[self.issuer.clone()]);
|
||||||
|
validation.set_audience(&[self.audience.clone()]);
|
||||||
|
validation.leeway = 60; // 60s clock skew tolerance
|
||||||
|
|
||||||
|
let token_data: TokenData<JwtClaims> =
|
||||||
|
decode::<JwtClaims>(token, &decoding_key, &validation)
|
||||||
|
.map_err(|e| anyhow!("Token validation failed: {}", e))?;
|
||||||
|
|
||||||
|
Ok(token_data.claims)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract bearer token from Authorization header
|
||||||
|
pub fn extract_bearer_token(auth_header: &str) -> Result<String> {
|
||||||
|
let parts: Vec<&str> = auth_header.split_whitespace().collect();
|
||||||
|
if parts.len() != 2 || parts[0].to_lowercase() != "bearer" {
|
||||||
|
return Err(anyhow!("Invalid Authorization header format"));
|
||||||
|
}
|
||||||
|
Ok(parts[1].to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_extract_bearer_token_valid() {
|
||||||
|
let header = "Bearer eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0";
|
||||||
|
let token = JwtValidator::extract_bearer_token(header).unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
token,
|
||||||
|
"eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_extract_bearer_token_invalid_format() {
|
||||||
|
let header = "Basic dXNlcjpwYXNz";
|
||||||
|
let result = JwtValidator::extract_bearer_token(header);
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_extract_bearer_token_missing() {
|
||||||
|
let header = "Bearer";
|
||||||
|
let result = JwtValidator::extract_bearer_token(header);
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,3 +1,4 @@
|
|||||||
|
pub mod endpoints;
|
||||||
pub mod handlers;
|
pub mod handlers;
|
||||||
pub mod http_server;
|
pub mod http_server;
|
||||||
pub mod metrics;
|
pub mod metrics;
|
||||||
@@ -6,6 +7,19 @@ pub mod relevance_judge;
|
|||||||
pub mod query;
|
pub mod query;
|
||||||
pub mod auth;
|
pub mod auth;
|
||||||
pub mod ingest_worker;
|
pub mod ingest_worker;
|
||||||
|
pub mod query_worker;
|
||||||
|
pub mod rate_limiter;
|
||||||
|
pub mod idempotency;
|
||||||
|
pub mod jwt_validator;
|
||||||
|
pub mod opensearch_client;
|
||||||
|
pub mod dual_write_indexer;
|
||||||
|
pub mod queue_adapter;
|
||||||
|
pub mod gateway_queue_adapter;
|
||||||
|
pub mod queue_worker;
|
||||||
|
pub mod query_optimizer;
|
||||||
|
pub mod simple_hybrid_search;
|
||||||
|
pub mod accuracy_metrics;
|
||||||
|
pub mod context_endpoint;
|
||||||
pub mod verify;
|
pub mod verify;
|
||||||
pub mod rbac;
|
pub mod rbac;
|
||||||
pub mod hybrid_retrieval;
|
pub mod hybrid_retrieval;
|
||||||
@@ -20,14 +34,17 @@ pub mod federation;
|
|||||||
pub mod query_router;
|
pub mod query_router;
|
||||||
pub mod full_pipeline;
|
pub mod full_pipeline;
|
||||||
pub mod authorized_pipeline;
|
pub mod authorized_pipeline;
|
||||||
|
// pub mod ingest_with_persistence; // TODO: Fix db_repo integration
|
||||||
pub mod auth_middleware;
|
pub mod auth_middleware;
|
||||||
pub mod compaction;
|
pub mod compaction;
|
||||||
pub mod compaction_executor;
|
pub mod compaction_executor;
|
||||||
pub mod agent;
|
pub mod agent;
|
||||||
pub mod parallel_dual_write;
|
pub mod parallel_dual_write;
|
||||||
|
|
||||||
|
pub use endpoints::{IngestQueue, IngestRequest, JobStatus};
|
||||||
pub use http_server::{AppState, AuthMode};
|
pub use http_server::{AppState, AuthMode};
|
||||||
pub use ingest_worker::IngestWorker;
|
pub use ingest_worker::IngestWorker;
|
||||||
|
pub use query_worker::QueryWorker;
|
||||||
pub use hybrid_retrieval::{HybridRetriever, RetrievalRoute, WikiScopedFilter, RankedCandidate};
|
pub use hybrid_retrieval::{HybridRetriever, RetrievalRoute, WikiScopedFilter, RankedCandidate};
|
||||||
pub use chunk_optimizer::{ChunkOptimizer, OptimizableChunk, SelectionMetrics};
|
pub use chunk_optimizer::{ChunkOptimizer, OptimizableChunk, SelectionMetrics};
|
||||||
pub use chunk_metadata::{MetadataExtractor, MetadataBooster, ChunkMetadata, ChunkCategory, QueryIntent};
|
pub use chunk_metadata::{MetadataExtractor, MetadataBooster, ChunkMetadata, ChunkCategory, QueryIntent};
|
||||||
|
|||||||
@@ -1,7 +1,21 @@
|
|||||||
mod lessons_cmd;
|
mod lessons_cmd;
|
||||||
// Dead modules removed — see lib.rs for live module list
|
// http_server is in lib.rs, use mem_cli::http_server
|
||||||
|
mod endpoints;
|
||||||
mod ingest_worker;
|
mod ingest_worker;
|
||||||
|
mod query_worker;
|
||||||
|
mod rate_limiter;
|
||||||
|
mod idempotency;
|
||||||
|
mod jwt_validator;
|
||||||
mod verify;
|
mod verify;
|
||||||
|
mod opensearch_client;
|
||||||
|
mod dual_write_indexer;
|
||||||
|
mod queue_adapter;
|
||||||
|
mod gateway_queue_adapter;
|
||||||
|
mod queue_worker;
|
||||||
|
mod context_endpoint;
|
||||||
|
mod query_optimizer;
|
||||||
|
mod simple_hybrid_search;
|
||||||
|
mod accuracy_metrics;
|
||||||
|
|
||||||
use clap::{Parser, Subcommand};
|
use clap::{Parser, Subcommand};
|
||||||
use mem_chunk::token_counter::CharsOverFourCounter;
|
use mem_chunk::token_counter::CharsOverFourCounter;
|
||||||
@@ -357,7 +371,7 @@ async fn cmd_verify(
|
|||||||
check_db,
|
check_db,
|
||||||
check_log,
|
check_log,
|
||||||
log_dir,
|
log_dir,
|
||||||
_format: format,
|
format,
|
||||||
};
|
};
|
||||||
|
|
||||||
let verifier = verify::Verifier::new(database_url).await?;
|
let verifier = verify::Verifier::new(database_url).await?;
|
||||||
|
|||||||
@@ -105,14 +105,14 @@ impl Histogram {
|
|||||||
/// Labeled counter (key = label combination string)
|
/// Labeled counter (key = label combination string)
|
||||||
pub struct LabeledCounter {
|
pub struct LabeledCounter {
|
||||||
values: Mutex<HashMap<String, u64>>,
|
values: Mutex<HashMap<String, u64>>,
|
||||||
_name: &'static str,
|
name: &'static str,
|
||||||
_help: &'static str,
|
help: &'static str,
|
||||||
_label_names: &'static [&'static str],
|
label_names: &'static [&'static str],
|
||||||
}
|
}
|
||||||
|
|
||||||
impl LabeledCounter {
|
impl LabeledCounter {
|
||||||
pub fn new(name: &'static str, help: &'static str, label_names: &'static [&'static str]) -> Self {
|
pub fn new(name: &'static str, help: &'static str, label_names: &'static [&'static str]) -> Self {
|
||||||
Self { values: Mutex::new(HashMap::new()), _name: name, _help: help, _label_names: label_names }
|
Self { values: Mutex::new(HashMap::new()), name, help, label_names }
|
||||||
}
|
}
|
||||||
pub fn inc(&self, labels: &[&str]) {
|
pub fn inc(&self, labels: &[&str]) {
|
||||||
let key = labels.join(",");
|
let key = labels.join(",");
|
||||||
@@ -381,16 +381,6 @@ pub static ERROR_UNEXPECTED_QUERY: Counter = Counter::new(
|
|||||||
pub static ERROR_UNEXPECTED_CONTEXT: Counter = Counter::new(
|
pub static ERROR_UNEXPECTED_CONTEXT: Counter = Counter::new(
|
||||||
"memory_error_unexpected_context_total", "Unexpected errors during context");
|
"memory_error_unexpected_context_total", "Unexpected errors during context");
|
||||||
|
|
||||||
// Agent endpoint error counters
|
|
||||||
pub static ERROR_AUTH_FAILURE_AGENT: Counter = Counter::new(
|
|
||||||
"memory_error_auth_failure_agent_total", "Auth failures on agent endpoints");
|
|
||||||
pub static ERROR_BAD_REQUEST_AGENT: Counter = Counter::new(
|
|
||||||
"memory_error_bad_request_agent_total", "Bad request errors on agent endpoints (expected)");
|
|
||||||
pub static ERROR_NOT_FOUND_AGENT: Counter = Counter::new(
|
|
||||||
"memory_error_not_found_agent_total", "Not found errors on agent endpoints (expected)");
|
|
||||||
pub static ERROR_UNEXPECTED_AGENT: Counter = Counter::new(
|
|
||||||
"memory_error_unexpected_agent_total", "Unexpected errors on agent endpoints (DB failures, 500s)");
|
|
||||||
|
|
||||||
// Last error info (most recent error for debugging)
|
// Last error info (most recent error for debugging)
|
||||||
pub static LAST_ERROR_TIMESTAMP: Gauge = Gauge::new(
|
pub static LAST_ERROR_TIMESTAMP: Gauge = Gauge::new(
|
||||||
"memory_last_error_timestamp_seconds", "Unix timestamp of most recent error");
|
"memory_last_error_timestamp_seconds", "Unix timestamp of most recent error");
|
||||||
@@ -603,27 +593,22 @@ pub fn render_metrics() -> String {
|
|||||||
counter!(ERROR_UNEXPECTED_INGEST);
|
counter!(ERROR_UNEXPECTED_INGEST);
|
||||||
counter!(ERROR_UNEXPECTED_QUERY);
|
counter!(ERROR_UNEXPECTED_QUERY);
|
||||||
counter!(ERROR_UNEXPECTED_CONTEXT);
|
counter!(ERROR_UNEXPECTED_CONTEXT);
|
||||||
counter!(ERROR_AUTH_FAILURE_AGENT);
|
|
||||||
counter!(ERROR_BAD_REQUEST_AGENT);
|
|
||||||
counter!(ERROR_NOT_FOUND_AGENT);
|
|
||||||
counter!(ERROR_UNEXPECTED_AGENT);
|
|
||||||
gauge!(LAST_ERROR_TIMESTAMP);
|
gauge!(LAST_ERROR_TIMESTAMP);
|
||||||
|
|
||||||
out
|
out
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Render a labeled counter in Prometheus format
|
/// Render a labeled counter in Prometheus format
|
||||||
#[allow(dead_code)]
|
|
||||||
fn render_labeled_counter(out: &mut String, lc: &LabeledCounter) {
|
fn render_labeled_counter(out: &mut String, lc: &LabeledCounter) {
|
||||||
let map = lc.values.lock().unwrap();
|
let map = lc.values.lock().unwrap();
|
||||||
if map.is_empty() { return; }
|
if map.is_empty() { return; }
|
||||||
out.push_str(&format!("# HELP {} {}\n# TYPE {} counter\n", lc._name, lc._help, lc._name));
|
out.push_str(&format!("# HELP {} {}\n# TYPE {} counter\n", lc.name, lc.help, lc.name));
|
||||||
for (key, val) in map.iter() {
|
for (key, val) in map.iter() {
|
||||||
let parts: Vec<&str> = key.split(',').collect();
|
let parts: Vec<&str> = key.split(',').collect();
|
||||||
let labels: Vec<String> = lc._label_names.iter().zip(parts.iter())
|
let labels: Vec<String> = lc.label_names.iter().zip(parts.iter())
|
||||||
.map(|(name, val)| format!("{}=\"{}\"", name, val))
|
.map(|(name, val)| format!("{}=\"{}\"", name, val))
|
||||||
.collect();
|
.collect();
|
||||||
out.push_str(&format!("{}{{{}}} {}\n", lc._name, labels.join(","), val));
|
out.push_str(&format!("{}{{{}}} {}\n", lc.name, labels.join(","), val));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -22,8 +22,8 @@ use crate::metrics;
|
|||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct MetricsSnapshot {
|
pub struct MetricsSnapshot {
|
||||||
counters: HashMap<&'static str, u64>,
|
counters: HashMap<&'static str, u64>,
|
||||||
_gauges: HashMap<&'static str, u64>,
|
gauges: HashMap<&'static str, u64>,
|
||||||
_gauges_f64: HashMap<&'static str, f64>,
|
gauges_f64: HashMap<&'static str, f64>,
|
||||||
histogram_counts: HashMap<&'static str, u64>,
|
histogram_counts: HashMap<&'static str, u64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -126,7 +126,7 @@ impl MetricsSnapshot {
|
|||||||
counters.insert("memory_db_queries_total", metrics::DB_QUERY_TOTAL.get());
|
counters.insert("memory_db_queries_total", metrics::DB_QUERY_TOTAL.get());
|
||||||
counters.insert("memory_db_query_errors_total", metrics::DB_QUERY_ERRORS.get());
|
counters.insert("memory_db_query_errors_total", metrics::DB_QUERY_ERRORS.get());
|
||||||
|
|
||||||
Self { counters, _gauges: gauges, _gauges_f64: gauges_f64, histogram_counts }
|
Self { counters, gauges, gauges_f64, histogram_counts }
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Assert a counter increased by exactly `expected` since snapshot
|
/// Assert a counter increased by exactly `expected` since snapshot
|
||||||
@@ -316,7 +316,7 @@ mod tests {
|
|||||||
let snap = MetricsSnapshot::capture();
|
let snap = MetricsSnapshot::capture();
|
||||||
assert!(snap.counters.contains_key("memory_ingest_requests_total"));
|
assert!(snap.counters.contains_key("memory_ingest_requests_total"));
|
||||||
assert!(snap.counters.contains_key("memory_query_requests_total"));
|
assert!(snap.counters.contains_key("memory_query_requests_total"));
|
||||||
assert!(snap._gauges.contains_key("memory_ingest_in_flight"));
|
assert!(snap.gauges.contains_key("memory_ingest_in_flight"));
|
||||||
assert!(snap.histogram_counts.contains_key("memory_ingest_duration_seconds"));
|
assert!(snap.histogram_counts.contains_key("memory_ingest_duration_seconds"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,382 @@
|
|||||||
|
use anyhow::{anyhow, Result};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
/// OpenSearch client for hybrid search (semantic + lexical)
|
||||||
|
pub struct OpenSearchClient {
|
||||||
|
hosts: Vec<String>,
|
||||||
|
client: reqwest::Client,
|
||||||
|
cache: Arc<RwLock<SearchCache>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct SearchResult {
|
||||||
|
pub id: String,
|
||||||
|
pub chunk: String,
|
||||||
|
pub score: f32,
|
||||||
|
pub source: String,
|
||||||
|
pub level: String,
|
||||||
|
pub breadcrumb: Vec<String>,
|
||||||
|
pub method: String, // "semantic", "lexical", or "hybrid"
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct HybridSearchResult {
|
||||||
|
pub results: Vec<SearchResult>,
|
||||||
|
pub total: usize,
|
||||||
|
pub query: String,
|
||||||
|
pub search_method: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct SearchCache {
|
||||||
|
queries: std::collections::HashMap<String, (HybridSearchResult, std::time::Instant)>,
|
||||||
|
ttl_secs: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl OpenSearchClient {
|
||||||
|
/// Create new OpenSearch client
|
||||||
|
pub fn new(hosts: Vec<String>) -> Self {
|
||||||
|
let client = reqwest::Client::builder()
|
||||||
|
.timeout(std::time::Duration::from_secs(30))
|
||||||
|
.build()
|
||||||
|
.expect("Failed to create HTTP client");
|
||||||
|
|
||||||
|
Self {
|
||||||
|
hosts,
|
||||||
|
client,
|
||||||
|
cache: Arc::new(RwLock::new(SearchCache {
|
||||||
|
queries: std::collections::HashMap::new(),
|
||||||
|
ttl_secs: 300, // 5 minute cache
|
||||||
|
})),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the primary host
|
||||||
|
fn primary_host(&self) -> &str {
|
||||||
|
&self.hosts[0]
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Index a document (called on vault changes)
|
||||||
|
pub async fn index_document(
|
||||||
|
&self,
|
||||||
|
doc_id: &str,
|
||||||
|
content: &str,
|
||||||
|
source: &str,
|
||||||
|
level: &str,
|
||||||
|
breadcrumb: Vec<String>,
|
||||||
|
jwt_token: &str,
|
||||||
|
) -> Result<()> {
|
||||||
|
let url = format!(
|
||||||
|
"https://{}/vault-*/_doc/{}",
|
||||||
|
self.primary_host(),
|
||||||
|
doc_id
|
||||||
|
);
|
||||||
|
|
||||||
|
let body = json!({
|
||||||
|
"content": content,
|
||||||
|
"source": source,
|
||||||
|
"level": level,
|
||||||
|
"breadcrumb": breadcrumb,
|
||||||
|
"indexed_at": chrono::Utc::now().to_rfc3339(),
|
||||||
|
});
|
||||||
|
|
||||||
|
let response = self
|
||||||
|
.client
|
||||||
|
.put(&url)
|
||||||
|
.header("Authorization", format!("Bearer {}", jwt_token))
|
||||||
|
.json(&body)
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if !response.status().is_success() {
|
||||||
|
return Err(anyhow!(
|
||||||
|
"OpenSearch index failed: {} {}",
|
||||||
|
response.status(),
|
||||||
|
response.text().await.unwrap_or_default()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Invalidate cache after indexing
|
||||||
|
self.cache.write().await.queries.clear();
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// BM25 lexical search via OpenSearch
|
||||||
|
async fn lexical_search(
|
||||||
|
&self,
|
||||||
|
query: &str,
|
||||||
|
limit: usize,
|
||||||
|
jwt_token: &str,
|
||||||
|
) -> Result<Vec<(String, f32, String, String, Vec<String>)>> {
|
||||||
|
let url = format!("https://{}/vault-*/_search", self.primary_host());
|
||||||
|
|
||||||
|
let search_body = json!({
|
||||||
|
"size": limit * 2,
|
||||||
|
"query": {
|
||||||
|
"multi_match": {
|
||||||
|
"query": query,
|
||||||
|
"fields": ["content^2", "source", "breadcrumb"],
|
||||||
|
"fuzziness": "AUTO",
|
||||||
|
"operator": "or"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"_source": ["content", "source", "level", "breadcrumb"]
|
||||||
|
});
|
||||||
|
|
||||||
|
let response = self
|
||||||
|
.client
|
||||||
|
.get(&url)
|
||||||
|
.header("Authorization", format!("Bearer {}", jwt_token))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.json(&search_body)
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if !response.status().is_success() {
|
||||||
|
return Err(anyhow!(
|
||||||
|
"OpenSearch search failed: {} {}",
|
||||||
|
response.status(),
|
||||||
|
response.text().await.unwrap_or_default()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let result: Value = response.json().await?;
|
||||||
|
|
||||||
|
let mut results = Vec::new();
|
||||||
|
if let Some(hits) = result["hits"]["hits"].as_array() {
|
||||||
|
for hit in hits {
|
||||||
|
let score = hit["_score"].as_f64().unwrap_or(0.0) as f32;
|
||||||
|
let source = &hit["_source"];
|
||||||
|
|
||||||
|
let id = hit["_id"].as_str().unwrap_or("").to_string();
|
||||||
|
let chunk = source["content"].as_str().unwrap_or("").to_string();
|
||||||
|
let src = source["source"].as_str().unwrap_or("").to_string();
|
||||||
|
let level = source["level"].as_str().unwrap_or("L0").to_string();
|
||||||
|
let breadcrumb: Vec<String> = source["breadcrumb"]
|
||||||
|
.as_array()
|
||||||
|
.map(|arr| {
|
||||||
|
arr.iter()
|
||||||
|
.filter_map(|v| v.as_str().map(|s| s.to_string()))
|
||||||
|
.collect()
|
||||||
|
})
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
results.push((id, score, chunk, src, breadcrumb));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(results)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Semantic search via pgvector (called from memory service)
|
||||||
|
/// This is separate - pgvector search happens in PostgreSQL
|
||||||
|
pub async fn semantic_search(
|
||||||
|
&self,
|
||||||
|
embedding: &[f32],
|
||||||
|
limit: usize,
|
||||||
|
jwt_token: &str,
|
||||||
|
) -> Result<Vec<(String, f32, String, String, Vec<String>)>> {
|
||||||
|
// NOTE: This is actually handled by pgvector in PostgreSQL
|
||||||
|
// This method is a placeholder for consistency
|
||||||
|
// The actual semantic search happens in crates/mem-cli/src/http_server.rs
|
||||||
|
Err(anyhow!(
|
||||||
|
"Semantic search must be done via pgvector in PostgreSQL, not OpenSearch"
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Hybrid search: combine lexical (OpenSearch) + semantic (pgvector)
|
||||||
|
pub async fn hybrid_search(
|
||||||
|
&self,
|
||||||
|
query: &str,
|
||||||
|
semantic_results: Vec<(String, f32, String, String, Vec<String>)>,
|
||||||
|
jwt_token: &str,
|
||||||
|
limit: usize,
|
||||||
|
weights: &HybridWeights,
|
||||||
|
) -> Result<HybridSearchResult> {
|
||||||
|
// Check cache
|
||||||
|
{
|
||||||
|
let cache = self.cache.read().await;
|
||||||
|
if let Some((cached, timestamp)) = cache.queries.get(query) {
|
||||||
|
if timestamp.elapsed().as_secs() < cache.ttl_secs {
|
||||||
|
return Ok(cached.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Perform lexical search
|
||||||
|
let lexical_results = self
|
||||||
|
.lexical_search(query, limit, jwt_token)
|
||||||
|
.await
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
// Combine results
|
||||||
|
let combined = self.combine_results(
|
||||||
|
semantic_results,
|
||||||
|
lexical_results,
|
||||||
|
limit,
|
||||||
|
weights,
|
||||||
|
);
|
||||||
|
|
||||||
|
let result = HybridSearchResult {
|
||||||
|
results: combined,
|
||||||
|
total: limit,
|
||||||
|
query: query.to_string(),
|
||||||
|
search_method: "hybrid".to_string(),
|
||||||
|
};
|
||||||
|
|
||||||
|
// Cache result
|
||||||
|
{
|
||||||
|
let mut cache = self.cache.write().await;
|
||||||
|
cache.queries.insert(query.to_string(), (result.clone(), std::time::Instant::now()));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Combine semantic and lexical results with reranking
|
||||||
|
fn combine_results(
|
||||||
|
&self,
|
||||||
|
semantic: Vec<(String, f32, String, String, Vec<String>)>,
|
||||||
|
lexical: Vec<(String, f32, String, String, Vec<String>)>,
|
||||||
|
limit: usize,
|
||||||
|
weights: &HybridWeights,
|
||||||
|
) -> Vec<SearchResult> {
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
// Normalize scores to 0-1
|
||||||
|
let sem_max = semantic.iter().map(|(_, s, _, _, _)| s).cloned().fold(f32::NEG_INFINITY, f32::max);
|
||||||
|
let lex_max = lexical.iter().map(|(_, s, _, _, _)| s).cloned().fold(f32::NEG_INFINITY, f32::max);
|
||||||
|
|
||||||
|
let sem_norm = semantic.into_iter().map(|(id, s, chunk, src, bc)| {
|
||||||
|
let normalized = if sem_max > 0.0 { s / sem_max } else { 0.0 };
|
||||||
|
(id, normalized, chunk, src, bc)
|
||||||
|
}).collect::<Vec<_>>();
|
||||||
|
|
||||||
|
let lex_norm = lexical.into_iter().map(|(id, s, chunk, src, bc)| {
|
||||||
|
let normalized = if lex_max > 0.0 { s / lex_max } else { 0.0 };
|
||||||
|
(id, normalized, chunk, src, bc)
|
||||||
|
}).collect::<Vec<_>>();
|
||||||
|
|
||||||
|
// Combine with weighted average
|
||||||
|
let mut combined: HashMap<String, (f32, String, String, Vec<String>)> = HashMap::new();
|
||||||
|
|
||||||
|
for (id, sem_score, chunk, src, bc) in sem_norm {
|
||||||
|
let lex_score = lex_norm
|
||||||
|
.iter()
|
||||||
|
.find(|(lid, _, _, _, _)| lid == &id)
|
||||||
|
.map(|(_, s, _, _, _)| *s)
|
||||||
|
.unwrap_or(0.0);
|
||||||
|
|
||||||
|
let final_score = weights.semantic * sem_score + weights.lexical * lex_score;
|
||||||
|
combined.insert(id, (final_score, chunk, src, bc));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add lexical-only results
|
||||||
|
for (id, lex_score, chunk, src, bc) in lex_norm {
|
||||||
|
if !combined.contains_key(&id) {
|
||||||
|
let final_score = weights.lexical * lex_score;
|
||||||
|
combined.insert(id, (final_score, chunk, src, bc));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sort and take top-k
|
||||||
|
let mut results: Vec<_> = combined
|
||||||
|
.into_iter()
|
||||||
|
.map(|(id, (score, chunk, src, bc))| SearchResult {
|
||||||
|
id,
|
||||||
|
chunk,
|
||||||
|
score,
|
||||||
|
source: src,
|
||||||
|
level: "L1".to_string(),
|
||||||
|
breadcrumb: bc,
|
||||||
|
method: "hybrid".to_string(),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
|
||||||
|
results.truncate(limit);
|
||||||
|
|
||||||
|
results
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Health check
|
||||||
|
pub async fn health(&self, jwt_token: &str) -> Result<bool> {
|
||||||
|
let url = format!("https://{}/_cluster/health", self.primary_host());
|
||||||
|
|
||||||
|
let response = self
|
||||||
|
.client
|
||||||
|
.get(&url)
|
||||||
|
.header("Authorization", format!("Bearer {}", jwt_token))
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(response.status().is_success())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
pub struct HybridWeights {
|
||||||
|
pub semantic: f32, // 0.6 = 60%
|
||||||
|
pub lexical: f32, // 0.4 = 40%
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for HybridWeights {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
semantic: 0.6,
|
||||||
|
lexical: 0.4,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_hybrid_weights_sum() {
|
||||||
|
let weights = HybridWeights::default();
|
||||||
|
assert!((weights.semantic + weights.lexical - 1.0).abs() < 0.01);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_combine_results_ranking() {
|
||||||
|
let client = OpenSearchClient::new(vec!["localhost:9200".to_string()]);
|
||||||
|
|
||||||
|
let semantic = vec![
|
||||||
|
(
|
||||||
|
"doc1".to_string(),
|
||||||
|
0.9,
|
||||||
|
"deployment content".to_string(),
|
||||||
|
"deploy.md".to_string(),
|
||||||
|
vec!["runbooks".to_string()],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"doc2".to_string(),
|
||||||
|
0.7,
|
||||||
|
"networking content".to_string(),
|
||||||
|
"network.md".to_string(),
|
||||||
|
vec!["docs".to_string()],
|
||||||
|
),
|
||||||
|
];
|
||||||
|
|
||||||
|
let lexical = vec![
|
||||||
|
(
|
||||||
|
"doc1".to_string(),
|
||||||
|
0.95,
|
||||||
|
"deployment content".to_string(),
|
||||||
|
"deploy.md".to_string(),
|
||||||
|
vec!["runbooks".to_string()],
|
||||||
|
),
|
||||||
|
];
|
||||||
|
|
||||||
|
let weights = HybridWeights::default();
|
||||||
|
let results = client.combine_results(semantic, lexical, 10, &weights);
|
||||||
|
|
||||||
|
assert_eq!(results.len(), 2);
|
||||||
|
assert_eq!(results[0].id, "doc1"); // doc1 has both semantic and lexical scores
|
||||||
|
assert!(results[0].score > results[1].score);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -6,18 +6,10 @@
|
|||||||
use anyhow::{anyhow, Result};
|
use anyhow::{anyhow, Result};
|
||||||
use sha2::{Digest, Sha256};
|
use sha2::{Digest, Sha256};
|
||||||
use sqlx::PgPool;
|
use sqlx::PgPool;
|
||||||
|
use uuid::Uuid;
|
||||||
use pgvector::Vector;
|
use pgvector::Vector;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
// OpenSearchClient removed (issue #56). Stub for compilation.
|
use crate::opensearch_client::OpenSearchClient;
|
||||||
#[allow(dead_code)]
|
|
||||||
pub struct OpenSearchClient;
|
|
||||||
|
|
||||||
impl OpenSearchClient {
|
|
||||||
#[allow(dead_code, unused_variables)]
|
|
||||||
pub async fn index_document(&self, chunk_id: &str, content: &str, source: &str, level: &str, breadcrumb: Vec<String>, jwt_token: &str) -> Result<(), String> {
|
|
||||||
Err("OpenSearchClient stub - not implemented".to_string())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
|
|||||||
@@ -8,7 +8,7 @@
|
|||||||
//! DRY: Reuses score types from mem_core
|
//! DRY: Reuses score types from mem_core
|
||||||
|
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use tracing::info;
|
use tracing::{debug, info};
|
||||||
|
|
||||||
/// Answer validation configuration
|
/// Answer validation configuration
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
|||||||
@@ -3,8 +3,9 @@
|
|||||||
/// Performs breadth-first search on memory_entity + memory_edge tables,
|
/// Performs breadth-first search on memory_entity + memory_edge tables,
|
||||||
/// returning a subgraph for visualization.
|
/// returning a subgraph for visualization.
|
||||||
|
|
||||||
use std::collections::VecDeque;
|
use std::collections::{HashMap, VecDeque};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
use sqlx::{Pool, Postgres, Row};
|
use sqlx::{Pool, Postgres, Row};
|
||||||
|
|
||||||
/// A node in the traversal result
|
/// A node in the traversal result
|
||||||
@@ -213,9 +214,9 @@ impl BfsGraphTraversal {
|
|||||||
/// Returns: (id, entity_type, name, description)
|
/// Returns: (id, entity_type, name, description)
|
||||||
async fn load_entity(&self, id: &str) -> Result<Option<(String, String, String, Option<String>)>, String> {
|
async fn load_entity(&self, id: &str) -> Result<Option<(String, String, String, Option<String>)>, String> {
|
||||||
let query = r#"
|
let query = r#"
|
||||||
SELECT id::TEXT, entity_type, name, description
|
SELECT id, entity_type, name, description
|
||||||
FROM memory_entity
|
FROM memory_entity
|
||||||
WHERE id = $1::UUID AND t_expired IS NULL
|
WHERE id = $1 AND deleted_at IS NULL
|
||||||
LIMIT 1;
|
LIMIT 1;
|
||||||
"#;
|
"#;
|
||||||
|
|
||||||
@@ -237,10 +238,10 @@ impl BfsGraphTraversal {
|
|||||||
/// Returns: (edge_id, target_id, source_id, relation_type, fact, strength)
|
/// Returns: (edge_id, target_id, source_id, relation_type, fact, strength)
|
||||||
async fn load_edges_from(&self, source_id: &str, limit: usize) -> Result<Vec<(String, String, String, String, String, f32)>, String> {
|
async fn load_edges_from(&self, source_id: &str, limit: usize) -> Result<Vec<(String, String, String, String, String, f32)>, String> {
|
||||||
let query = r#"
|
let query = r#"
|
||||||
SELECT id::TEXT, target_id::TEXT, source_id::TEXT, relation_type, fact, confidence
|
SELECT id, target_id, source_id, relation_type, fact, strength
|
||||||
FROM memory_edge
|
FROM memory_edge
|
||||||
WHERE source_id = $1::UUID AND t_expired IS NULL AND t_invalid IS NULL
|
WHERE source_id = $1 AND t_expired IS NULL AND t_invalid IS NULL
|
||||||
ORDER BY confidence DESC
|
ORDER BY strength DESC
|
||||||
LIMIT $2;
|
LIMIT $2;
|
||||||
"#;
|
"#;
|
||||||
|
|
||||||
@@ -257,7 +258,7 @@ impl BfsGraphTraversal {
|
|||||||
r.get::<String, _>("source_id"),
|
r.get::<String, _>("source_id"),
|
||||||
r.get::<String, _>("relation_type"),
|
r.get::<String, _>("relation_type"),
|
||||||
r.get::<String, _>("fact"),
|
r.get::<String, _>("fact"),
|
||||||
r.get::<f32, _>("confidence"),
|
r.get::<f32, _>("strength"),
|
||||||
)).collect())
|
)).collect())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -6,7 +6,7 @@
|
|||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
use sqlx::PgPool;
|
use sqlx::PgPool;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use tracing::debug;
|
use tracing::{debug, warn};
|
||||||
|
|
||||||
/// Result of linking a text mention to an entity
|
/// Result of linking a text mention to an entity
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
@@ -89,12 +89,12 @@ pub struct CoreferenceCluster {
|
|||||||
|
|
||||||
/// Entity Linking Engine
|
/// Entity Linking Engine
|
||||||
pub struct EntityLinker {
|
pub struct EntityLinker {
|
||||||
_pool: PgPool,
|
pool: PgPool,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl EntityLinker {
|
impl EntityLinker {
|
||||||
pub fn new(pool: PgPool) -> Self {
|
pub fn new(pool: PgPool) -> Self {
|
||||||
EntityLinker { _pool: pool }
|
EntityLinker { pool }
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Link mentions in text to existing entities
|
/// Link mentions in text to existing entities
|
||||||
@@ -242,7 +242,7 @@ impl EntityLinker {
|
|||||||
|
|
||||||
let mut result = Vec::new();
|
let mut result = Vec::new();
|
||||||
for (entity_id, mentions) in clusters {
|
for (entity_id, mentions) in clusters {
|
||||||
if let Some(_entity) = entities.iter().find(|e| e.id == entity_id) {
|
if let Some(entity) = entities.iter().find(|e| e.id == entity_id) {
|
||||||
let unique_mentions: Vec<_> = mentions.iter().cloned().collect::<HashSet<_>>().into_iter().collect();
|
let unique_mentions: Vec<_> = mentions.iter().cloned().collect::<HashSet<_>>().into_iter().collect();
|
||||||
result.push(CoreferenceCluster {
|
result.push(CoreferenceCluster {
|
||||||
entity_id: entity_id.clone(),
|
entity_id: entity_id.clone(),
|
||||||
@@ -353,7 +353,7 @@ impl EntityLinker {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Fetch all entities for a project
|
/// Fetch all entities for a project
|
||||||
async fn fetch_entities(&self, _project_id: &str) -> Result<Vec<EntityInfo>, String> {
|
async fn fetch_entities(&self, project_id: &str) -> Result<Vec<EntityInfo>, String> {
|
||||||
// Stub: would query database
|
// Stub: would query database
|
||||||
// For now, return empty
|
// For now, return empty
|
||||||
Ok(vec![])
|
Ok(vec![])
|
||||||
|
|||||||
@@ -6,6 +6,7 @@
|
|||||||
use chrono::{DateTime, Timelike, Utc};
|
use chrono::{DateTime, Timelike, Utc};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use sqlx::{Pool, Postgres};
|
use sqlx::{Pool, Postgres};
|
||||||
|
use std::collections::HashMap;
|
||||||
use tracing::{debug, info};
|
use tracing::{debug, info};
|
||||||
|
|
||||||
/// A single facet (filterable dimension)
|
/// A single facet (filterable dimension)
|
||||||
@@ -87,7 +88,7 @@ impl FacetedSearch {
|
|||||||
limit: usize,
|
limit: usize,
|
||||||
) -> Result<AvailableFacets, String> {
|
) -> Result<AvailableFacets, String> {
|
||||||
let limit = limit.max(5).min(50);
|
let limit = limit.max(5).min(50);
|
||||||
let _start_time = std::time::Instant::now();
|
let start_time = std::time::Instant::now();
|
||||||
|
|
||||||
debug!("Discovering facets for {}, limit={}", search_type, limit);
|
debug!("Discovering facets for {}, limit={}", search_type, limit);
|
||||||
|
|
||||||
|
|||||||
@@ -4,7 +4,7 @@
|
|||||||
/// node positions in 2D space suitable for React Flow visualization.
|
/// node positions in 2D space suitable for React Flow visualization.
|
||||||
|
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use super::bfs_graph_traversal::{GraphData, TraversalNode};
|
use super::bfs_graph_traversal::{GraphData, TraversalNode, TraversalEdge};
|
||||||
|
|
||||||
/// 2D position (X, Y coordinates)
|
/// 2D position (X, Y coordinates)
|
||||||
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
|
||||||
@@ -184,8 +184,8 @@ impl ForceDirectedLayout {
|
|||||||
let dist = dist_sq.sqrt();
|
let dist = dist_sq.sqrt();
|
||||||
|
|
||||||
let force = charge / dist_sq;
|
let force = charge / dist_sq;
|
||||||
let fx = force * dx / dist;
|
let fx = (force * dx / dist);
|
||||||
let fy = force * dy / dist;
|
let fy = (force * dy / dist);
|
||||||
|
|
||||||
(-fx, -fy) // Negative = repulsive
|
(-fx, -fy) // Negative = repulsive
|
||||||
}
|
}
|
||||||
@@ -199,8 +199,8 @@ impl ForceDirectedLayout {
|
|||||||
let displacement = dist - link_distance;
|
let displacement = dist - link_distance;
|
||||||
let force = 0.1 * displacement; // Spring constant
|
let force = 0.1 * displacement; // Spring constant
|
||||||
|
|
||||||
let fx = force * dx / dist;
|
let fx = (force * dx / dist);
|
||||||
let fy = force * dy / dist;
|
let fy = (force * dy / dist);
|
||||||
|
|
||||||
(fx, fy) // Positive = attractive
|
(fx, fy) // Positive = attractive
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ use std::pin::Pin;
|
|||||||
use std::future::Future;
|
use std::future::Future;
|
||||||
use sqlx::PgPool;
|
use sqlx::PgPool;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use tracing::{debug, warn};
|
||||||
|
|
||||||
/// Inference rule
|
/// Inference rule
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
@@ -90,13 +91,13 @@ pub struct ReachableEntity {
|
|||||||
|
|
||||||
/// Inference Engine
|
/// Inference Engine
|
||||||
pub struct InferenceEngine {
|
pub struct InferenceEngine {
|
||||||
_pool: PgPool,
|
pool: PgPool,
|
||||||
rules: Vec<InferenceRule>,
|
rules: Vec<InferenceRule>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl InferenceEngine {
|
impl InferenceEngine {
|
||||||
pub fn new(pool: PgPool, rules: Vec<InferenceRule>) -> Self {
|
pub fn new(pool: PgPool, rules: Vec<InferenceRule>) -> Self {
|
||||||
InferenceEngine { _pool: pool, rules }
|
InferenceEngine { pool, rules }
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Perform rule-based inference
|
/// Perform rule-based inference
|
||||||
@@ -286,8 +287,8 @@ impl InferenceEngine {
|
|||||||
/// Fetch edges from entity
|
/// Fetch edges from entity
|
||||||
async fn fetch_entity_edges(
|
async fn fetch_entity_edges(
|
||||||
&self,
|
&self,
|
||||||
_entity_id: &str,
|
entity_id: &str,
|
||||||
_project_id: &str,
|
project_id: &str,
|
||||||
) -> Result<Vec<EdgeInfo>, String> {
|
) -> Result<Vec<EdgeInfo>, String> {
|
||||||
// Stub: would query database
|
// Stub: would query database
|
||||||
Ok(vec![])
|
Ok(vec![])
|
||||||
@@ -358,7 +359,7 @@ impl InferenceEngine {
|
|||||||
|
|
||||||
/// Internal edge info
|
/// Internal edge info
|
||||||
struct EdgeInfo {
|
struct EdgeInfo {
|
||||||
_source_id: String,
|
source_id: String,
|
||||||
target_id: String,
|
target_id: String,
|
||||||
target_name: String,
|
target_name: String,
|
||||||
relation_type: String,
|
relation_type: String,
|
||||||
|
|||||||
@@ -47,7 +47,7 @@ pub struct PathFindingResult {
|
|||||||
/// Edge representation for path finding
|
/// Edge representation for path finding
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
struct GraphEdge {
|
struct GraphEdge {
|
||||||
_from_id: String,
|
from_id: String,
|
||||||
to_id: String,
|
to_id: String,
|
||||||
relation_type: String,
|
relation_type: String,
|
||||||
confidence: f32,
|
confidence: f32,
|
||||||
@@ -396,14 +396,14 @@ impl PathFinder {
|
|||||||
// Normalize direction: always point forward from input entity
|
// Normalize direction: always point forward from input entity
|
||||||
if source == entity_id {
|
if source == entity_id {
|
||||||
GraphEdge {
|
GraphEdge {
|
||||||
_from_id: source,
|
from_id: source,
|
||||||
to_id: target,
|
to_id: target,
|
||||||
relation_type: rel_type,
|
relation_type: rel_type,
|
||||||
confidence: conf.max(0.0).min(1.0),
|
confidence: conf.max(0.0).min(1.0),
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
GraphEdge {
|
GraphEdge {
|
||||||
_from_id: target,
|
from_id: target,
|
||||||
to_id: source,
|
to_id: source,
|
||||||
relation_type: format!("{}(reverse)", rel_type),
|
relation_type: format!("{}(reverse)", rel_type),
|
||||||
confidence: conf.max(0.0).min(1.0),
|
confidence: conf.max(0.0).min(1.0),
|
||||||
@@ -580,13 +580,13 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn test_edge_representation() {
|
fn test_edge_representation() {
|
||||||
let edge = GraphEdge {
|
let edge = GraphEdge {
|
||||||
_from_id: "e1".to_string(),
|
from_id: "e1".to_string(),
|
||||||
to_id: "e2".to_string(),
|
to_id: "e2".to_string(),
|
||||||
relation_type: "related".to_string(),
|
relation_type: "related".to_string(),
|
||||||
confidence: 0.85,
|
confidence: 0.85,
|
||||||
};
|
};
|
||||||
|
|
||||||
assert_eq!(edge._from_id, "e1");
|
assert_eq!(edge.from_id, "e1");
|
||||||
assert_eq!(edge.to_id, "e2");
|
assert_eq!(edge.to_id, "e2");
|
||||||
assert!(edge.confidence >= 0.0 && edge.confidence <= 1.0);
|
assert!(edge.confidence >= 0.0 && edge.confidence <= 1.0);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -3,8 +3,10 @@
|
|||||||
//! Complex question decomposition, multi-hop reasoning, constraint satisfaction,
|
//! Complex question decomposition, multi-hop reasoning, constraint satisfaction,
|
||||||
//! and answer validation.
|
//! and answer validation.
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
use sqlx::PgPool;
|
use sqlx::PgPool;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use tracing::{debug, warn};
|
||||||
|
|
||||||
/// Question type/intent
|
/// Question type/intent
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||||
@@ -106,12 +108,12 @@ pub struct ReasonedAnswer {
|
|||||||
|
|
||||||
/// Query Reasoner
|
/// Query Reasoner
|
||||||
pub struct QueryReasoner {
|
pub struct QueryReasoner {
|
||||||
_pool: PgPool,
|
pool: PgPool,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl QueryReasoner {
|
impl QueryReasoner {
|
||||||
pub fn new(pool: PgPool) -> Self {
|
pub fn new(pool: PgPool) -> Self {
|
||||||
QueryReasoner { _pool: pool }
|
QueryReasoner { pool }
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Decompose complex question into sub-queries
|
/// Decompose complex question into sub-queries
|
||||||
@@ -120,7 +122,7 @@ impl QueryReasoner {
|
|||||||
return Ok(vec![]);
|
return Ok(vec![]);
|
||||||
}
|
}
|
||||||
|
|
||||||
let _question_lower = question.to_lowercase();
|
let question_lower = question.to_lowercase();
|
||||||
let question_type = self.classify_question(question);
|
let question_type = self.classify_question(question);
|
||||||
|
|
||||||
let mut sub_queries = Vec::new();
|
let mut sub_queries = Vec::new();
|
||||||
@@ -397,7 +399,7 @@ impl QueryReasoner {
|
|||||||
|
|
||||||
let mut explanation = format!("Found {} answer(s) through {} reasoning step(s): ", answers.len(), steps.len());
|
let mut explanation = format!("Found {} answer(s) through {} reasoning step(s): ", answers.len(), steps.len());
|
||||||
|
|
||||||
for (_idx, step) in steps.iter().enumerate() {
|
for (idx, step) in steps.iter().enumerate() {
|
||||||
explanation.push_str(&format!(
|
explanation.push_str(&format!(
|
||||||
"Step {}: {} (confidence: {:.2}, {} constraints satisfied). ",
|
"Step {}: {} (confidence: {:.2}, {} constraints satisfied). ",
|
||||||
step.step_id,
|
step.step_id,
|
||||||
|
|||||||
@@ -1,18 +1,13 @@
|
|||||||
//! Semantic Retrieval Engine
|
//! Semantic Retrieval Engine
|
||||||
//!
|
//!
|
||||||
//! Provides semantic search capabilities using vector embeddings and hybrid search
|
//! Provides semantic search capabilities using vector embeddings and hybrid search
|
||||||
//! combining vector (semantic) and lexical (ts_rank) results with RRF fusion.
|
//! combining vector (semantic) and lexical (keyword) results with RRF fusion.
|
||||||
//!
|
|
||||||
//! Schema alignment:
|
|
||||||
//! memory_entity: id, project_id, name, name_embedding, summary, description,
|
|
||||||
//! summary_embedding, entity_type, t_created, t_updated, t_expired, confidence
|
|
||||||
//! memory_edge: id, project_id, source_id, target_id, relation_type, fact,
|
|
||||||
//! fact_embedding, t_valid, t_invalid, t_created, t_expired, confidence
|
|
||||||
|
|
||||||
use chrono::{DateTime, Utc};
|
use chrono::{DateTime, Utc};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use sqlx::{Pool, Postgres};
|
use sqlx::{Pool, Postgres};
|
||||||
use tracing::{debug, info};
|
use std::sync::Arc;
|
||||||
|
use tracing::{debug, info, warn};
|
||||||
|
|
||||||
/// Semantic search result for an entity
|
/// Semantic search result for an entity
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
@@ -21,15 +16,15 @@ pub struct EntityResult {
|
|||||||
pub name: String,
|
pub name: String,
|
||||||
pub entity_type: String,
|
pub entity_type: String,
|
||||||
pub similarity_score: f32, // 0.0-1.0, higher is better
|
pub similarity_score: f32, // 0.0-1.0, higher is better
|
||||||
pub summary: Option<String>,
|
pub metadata: serde_json::Value,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Optional temporal filters for queries
|
/// Optional temporal filters for queries
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct TemporalFilter {
|
pub struct TemporalFilter {
|
||||||
pub start_time: Option<DateTime<Utc>>,
|
pub start_time: Option<DateTime<Utc>>, // Earliest event_time
|
||||||
pub end_time: Option<DateTime<Utc>>,
|
pub end_time: Option<DateTime<Utc>>, // Latest event_time
|
||||||
pub min_recency_score: Option<f32>,
|
pub min_recency_score: Option<f32>, // Only facts newer than this score (0-1)
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for TemporalFilter {
|
impl Default for TemporalFilter {
|
||||||
@@ -52,7 +47,7 @@ pub struct EdgeResult {
|
|||||||
pub target_name: String,
|
pub target_name: String,
|
||||||
pub relation_type: String,
|
pub relation_type: String,
|
||||||
pub fact: String,
|
pub fact: String,
|
||||||
pub similarity_score: f32,
|
pub similarity_score: f32, // 0.0-1.0, higher is better
|
||||||
pub confidence: f32,
|
pub confidence: f32,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -60,12 +55,12 @@ pub struct EdgeResult {
|
|||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct HybridResult {
|
pub struct HybridResult {
|
||||||
pub id: String,
|
pub id: String,
|
||||||
pub name: Option<String>,
|
pub name: Option<String>, // entity name or fact snippet
|
||||||
pub entity_type: Option<String>,
|
pub entity_type: Option<String>,
|
||||||
pub result_type: String, // "entity" or "edge"
|
pub result_type: String, // "entity" or "edge"
|
||||||
pub fused_score: f32, // RRF fused score
|
pub fused_score: f32, // RRF fused score
|
||||||
pub semantic_score: f32,
|
pub semantic_score: f32, // Vector similarity
|
||||||
pub lexical_score: f32,
|
pub lexical_score: f32, // BM25 ranking
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Semantic Retriever - performs vector and hybrid searches
|
/// Semantic Retriever - performs vector and hybrid searches
|
||||||
@@ -74,14 +69,25 @@ pub struct SemanticRetriever {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl SemanticRetriever {
|
impl SemanticRetriever {
|
||||||
|
/// Create a new semantic retriever
|
||||||
pub fn new(pool: Pool<Postgres>) -> Self {
|
pub fn new(pool: Pool<Postgres>) -> Self {
|
||||||
Self { pool }
|
Self { pool }
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Search entities by vector similarity on name_embedding.
|
/// Search for entities by semantic similarity
|
||||||
/// Falls back to summary_embedding if name_embedding is NULL.
|
|
||||||
///
|
///
|
||||||
/// Columns: name_embedding VECTOR(768), t_expired (soft delete), t_created (temporal)
|
/// # Arguments
|
||||||
|
/// * `query` - Search query text (will be embedded)
|
||||||
|
/// * `query_embedding` - Pre-computed query embedding (768-dim)
|
||||||
|
/// * `top_k` - Number of results to return (5-100)
|
||||||
|
/// * `entity_type_filter` - Optional entity type to filter by
|
||||||
|
/// * `confidence_floor` - Minimum similarity score (0.0-1.0)
|
||||||
|
/// * `start_time` - Optional earliest event_time
|
||||||
|
/// * `end_time` - Optional latest event_time
|
||||||
|
///
|
||||||
|
/// # Returns
|
||||||
|
/// Vector of EntityResult sorted by similarity (highest first)
|
||||||
|
/// All results have event_time within [start_time, end_time] if provided
|
||||||
pub async fn search_entities(
|
pub async fn search_entities(
|
||||||
&self,
|
&self,
|
||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
@@ -98,48 +104,48 @@ impl SemanticRetriever {
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
let top_k = top_k.max(1).min(100);
|
let top_k = top_k.max(1).min(100); // Clamp 1-100
|
||||||
if !(0.0..=1.0).contains(&confidence_floor) {
|
if confidence_floor < 0.0 || confidence_floor > 1.0 {
|
||||||
return Err("confidence_floor must be 0.0-1.0".to_string());
|
return Err("confidence_floor must be 0.0-1.0".to_string());
|
||||||
}
|
}
|
||||||
|
|
||||||
debug!("Searching entities: top_k={}, filter={:?}, time_range={:?}-{:?}",
|
debug!("Searching entities: top_k={}, filter={:?}, time_range={:?}-{:?}",
|
||||||
top_k, entity_type_filter, start_time, end_time);
|
top_k, entity_type_filter, start_time, end_time);
|
||||||
|
|
||||||
// Use COALESCE(name_embedding, summary_embedding) so entities with
|
// Query with temporal filters always included (NULL = no filter)
|
||||||
// only one embedding type are still searchable.
|
|
||||||
let query_sql =
|
let query_sql =
|
||||||
"SELECT id::TEXT, name, entity_type, summary,
|
"SELECT id, name, entity_type,
|
||||||
1 - (COALESCE(name_embedding, summary_embedding) <=> $1::vector) as similarity_score
|
1 - (embedding <=> $1::vector) as similarity_score,
|
||||||
|
metadata
|
||||||
FROM memory_entity
|
FROM memory_entity
|
||||||
WHERE t_expired IS NULL
|
WHERE deleted_at IS NULL
|
||||||
AND COALESCE(name_embedding, summary_embedding) IS NOT NULL
|
AND (1 - (embedding <=> $1::vector)) > $2
|
||||||
AND (1 - (COALESCE(name_embedding, summary_embedding) <=> $1::vector)) > $2
|
|
||||||
AND (entity_type = COALESCE($3, entity_type))
|
AND (entity_type = COALESCE($3, entity_type))
|
||||||
AND (t_created >= COALESCE($4, t_created))
|
AND (event_time >= COALESCE($4, event_time))
|
||||||
AND (t_created <= COALESCE($5, t_created))
|
AND (event_time <= COALESCE($5, event_time))
|
||||||
ORDER BY similarity_score DESC
|
ORDER BY similarity_score DESC
|
||||||
LIMIT $6";
|
LIMIT $6";
|
||||||
|
|
||||||
let results = sqlx::query_as::<_, (String, String, String, Option<String>, f32)>(query_sql)
|
// Always bind all parameters; COALESCE handles NULL filters
|
||||||
.bind(query_embedding)
|
let results = sqlx::query_as::<_, (String, String, String, f32, serde_json::Value)>(query_sql)
|
||||||
.bind(confidence_floor)
|
.bind(query_embedding) // $1: embedding vector
|
||||||
.bind(entity_type_filter)
|
.bind(confidence_floor) // $2: similarity threshold
|
||||||
.bind(start_time)
|
.bind(entity_type_filter) // $3: entity type (NULL = no filter)
|
||||||
.bind(end_time)
|
.bind(start_time) // $4: start_time (NULL = no filter)
|
||||||
.bind(top_k as i64)
|
.bind(end_time) // $5: end_time (NULL = no filter)
|
||||||
|
.bind(top_k as i64) // $6: LIMIT
|
||||||
.fetch_all(&self.pool)
|
.fetch_all(&self.pool)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| format!("Database error: {}", e))?;
|
.map_err(|e| format!("Database error: {}", e))?;
|
||||||
|
|
||||||
let entities: Vec<_> = results
|
let entities: Vec<_> = results
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|(id, name, entity_type, summary, score)| EntityResult {
|
.map(|(id, name, entity_type, score, metadata)| EntityResult {
|
||||||
id,
|
id,
|
||||||
name,
|
name,
|
||||||
entity_type,
|
entity_type,
|
||||||
similarity_score: score.clamp(0.0, 1.0),
|
similarity_score: score.max(0.0).min(1.0), // Clamp to 0-1
|
||||||
summary,
|
metadata,
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
@@ -147,10 +153,18 @@ impl SemanticRetriever {
|
|||||||
Ok(entities)
|
Ok(entities)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Search edges by vector similarity on fact_embedding.
|
/// Search for edges (relationships/facts) by semantic similarity
|
||||||
///
|
///
|
||||||
/// Columns: fact_embedding VECTOR(768), source_id, target_id,
|
/// # Arguments
|
||||||
/// t_invalid (temporal invalidation), t_expired (soft delete), t_created
|
/// * `query_embedding` - Pre-computed query embedding (768-dim)
|
||||||
|
/// * `top_k` - Number of results to return (5-100)
|
||||||
|
/// * `relation_type_filter` - Optional relation type to filter by
|
||||||
|
/// * `start_time` - Optional earliest event_time
|
||||||
|
/// * `end_time` - Optional latest event_time
|
||||||
|
///
|
||||||
|
/// # Returns
|
||||||
|
/// Vector of EdgeResult sorted by similarity (highest first)
|
||||||
|
/// All results have event_time within [start_time, end_time] if provided
|
||||||
pub async fn search_edges(
|
pub async fn search_edges(
|
||||||
&self,
|
&self,
|
||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
@@ -171,29 +185,30 @@ impl SemanticRetriever {
|
|||||||
debug!("Searching edges: top_k={}, filter={:?}, time_range={:?}-{:?}",
|
debug!("Searching edges: top_k={}, filter={:?}, time_range={:?}-{:?}",
|
||||||
top_k, relation_type_filter, start_time, end_time);
|
top_k, relation_type_filter, start_time, end_time);
|
||||||
|
|
||||||
|
// Query with temporal filters always included (NULL = no filter)
|
||||||
let query_sql =
|
let query_sql =
|
||||||
"SELECT e.id::TEXT, e.source_id::TEXT, e.target_id::TEXT,
|
"SELECT e.id, e.source_entity_id, e.target_entity_id,
|
||||||
src.name, tgt.name, e.relation_type, e.fact,
|
src.name, tgt.name, e.relation_type, e.fact,
|
||||||
1 - (e.fact_embedding <=> $1::vector) as similarity_score,
|
1 - (e.embedding <=> $1::vector) as similarity_score,
|
||||||
e.confidence
|
e.confidence
|
||||||
FROM memory_edge e
|
FROM memory_edge e
|
||||||
JOIN memory_entity src ON e.source_id = src.id
|
JOIN memory_entity src ON e.source_entity_id = src.id
|
||||||
JOIN memory_entity tgt ON e.target_id = tgt.id
|
JOIN memory_entity tgt ON e.target_entity_id = tgt.id
|
||||||
WHERE e.t_invalid IS NULL
|
WHERE e.fact_invalid_at IS NULL
|
||||||
AND e.t_expired IS NULL
|
AND e.deleted_at IS NULL
|
||||||
AND e.fact_embedding IS NOT NULL
|
|
||||||
AND (e.relation_type = COALESCE($2, e.relation_type))
|
AND (e.relation_type = COALESCE($2, e.relation_type))
|
||||||
AND (e.t_created >= COALESCE($3, e.t_created))
|
AND (e.event_time >= COALESCE($3, e.event_time))
|
||||||
AND (e.t_created <= COALESCE($4, e.t_created))
|
AND (e.event_time <= COALESCE($4, e.event_time))
|
||||||
ORDER BY similarity_score DESC
|
ORDER BY similarity_score DESC
|
||||||
LIMIT $5";
|
LIMIT $5";
|
||||||
|
|
||||||
let results = sqlx::query_as::<_, (String, String, String, String, String, String, String, f32, f64)>(query_sql)
|
// Always bind all parameters; COALESCE handles NULL filters
|
||||||
.bind(query_embedding)
|
let results = sqlx::query_as::<_, (String, String, String, String, String, String, String, f32, f32)>(query_sql)
|
||||||
.bind(relation_type_filter)
|
.bind(query_embedding) // $1: embedding vector
|
||||||
.bind(start_time)
|
.bind(relation_type_filter) // $2: relation type (NULL = no filter)
|
||||||
.bind(end_time)
|
.bind(start_time) // $3: start_time (NULL = no filter)
|
||||||
.bind(top_k as i64)
|
.bind(end_time) // $4: end_time (NULL = no filter)
|
||||||
|
.bind(top_k as i64) // $5: LIMIT
|
||||||
.fetch_all(&self.pool)
|
.fetch_all(&self.pool)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| format!("Database error: {}", e))?;
|
.map_err(|e| format!("Database error: {}", e))?;
|
||||||
@@ -209,8 +224,8 @@ impl SemanticRetriever {
|
|||||||
target_name: tgt_name,
|
target_name: tgt_name,
|
||||||
relation_type: rel_type,
|
relation_type: rel_type,
|
||||||
fact,
|
fact,
|
||||||
similarity_score: score.clamp(0.0, 1.0),
|
similarity_score: score.max(0.0).min(1.0),
|
||||||
confidence: (conf as f32).clamp(0.0, 1.0),
|
confidence: conf.max(0.0).min(1.0),
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
@@ -219,12 +234,19 @@ impl SemanticRetriever {
|
|||||||
Ok(edges)
|
Ok(edges)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Hybrid search: combines semantic (vector) and lexical (ts_rank) results
|
/// Hybrid search combining semantic (vector) and lexical (keyword) results
|
||||||
/// using Reciprocal Rank Fusion (RRF).
|
|
||||||
///
|
///
|
||||||
/// Unlike the previous stub, this actually runs a lexical search using
|
/// Uses Reciprocal Rank Fusion (RRF) to combine scores:
|
||||||
/// PostgreSQL full-text search (ts_rank + plainto_tsquery) on entity names
|
/// fused_score = (semantic_weight * normalized_semantic) + (lexical_weight * normalized_lexical)
|
||||||
/// and edge facts, then fuses with semantic results via RRF.
|
///
|
||||||
|
/// # Arguments
|
||||||
|
/// * `query_embedding` - Pre-computed query embedding (768-dim)
|
||||||
|
/// * `top_k` - Number of results to return (5-100)
|
||||||
|
/// * `semantic_weight` - Weight for semantic score (0.0-1.0, default 0.6)
|
||||||
|
/// * `lexical_weight` - Weight for lexical score (0.0-1.0, default 0.4)
|
||||||
|
///
|
||||||
|
/// # Returns
|
||||||
|
/// Vector of HybridResult sorted by fused_score (highest first)
|
||||||
pub async fn hybrid_search(
|
pub async fn hybrid_search(
|
||||||
&self,
|
&self,
|
||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
@@ -242,169 +264,66 @@ impl SemanticRetriever {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let top_k = top_k.max(1).min(100);
|
let top_k = top_k.max(1).min(100);
|
||||||
let sem_w = semantic_weight.clamp(0.0, 1.0);
|
let sem_w = semantic_weight.max(0.0).min(1.0);
|
||||||
let lex_w = lexical_weight.clamp(0.0, 1.0);
|
let lex_w = lexical_weight.max(0.0).min(1.0);
|
||||||
|
|
||||||
debug!("Hybrid search: top_k={}, weights=(sem={}, lex={}), time_range={:?}-{:?}",
|
debug!("Hybrid search: top_k={}, weights=(sem={}, lex={}), time_range={:?}-{:?}",
|
||||||
top_k, sem_w, lex_w, start_time, end_time);
|
top_k, sem_w, lex_w, start_time, end_time);
|
||||||
|
|
||||||
// Retrieve 2x candidates for RRF fusion
|
// Phase 1: Semantic search for entities
|
||||||
let fetch_k = (top_k * 2) as i64;
|
let entity_results = self.search_entities(
|
||||||
|
query_embedding,
|
||||||
|
top_k * 2,
|
||||||
|
None,
|
||||||
|
0.3,
|
||||||
|
start_time,
|
||||||
|
end_time,
|
||||||
|
).await?;
|
||||||
|
|
||||||
// --- Entity hybrid: semantic + lexical on name/summary ---
|
// Phase 2: Semantic search for edges
|
||||||
let entity_sql =
|
let edge_results = self.search_edges(
|
||||||
"WITH semantic AS (
|
query_embedding,
|
||||||
SELECT id::TEXT, name, entity_type, summary,
|
top_k * 2,
|
||||||
1 - (COALESCE(name_embedding, summary_embedding) <=> $1::vector) AS sem_score,
|
None,
|
||||||
ROW_NUMBER() OVER (ORDER BY COALESCE(name_embedding, summary_embedding) <=> $1::vector) AS sem_rank
|
start_time,
|
||||||
FROM memory_entity
|
end_time,
|
||||||
WHERE t_expired IS NULL
|
).await?;
|
||||||
AND COALESCE(name_embedding, summary_embedding) IS NOT NULL
|
|
||||||
AND (t_created >= COALESCE($3, t_created))
|
|
||||||
AND (t_created <= COALESCE($4, t_created))
|
|
||||||
ORDER BY COALESCE(name_embedding, summary_embedding) <=> $1::vector
|
|
||||||
LIMIT $5
|
|
||||||
),
|
|
||||||
lexical AS (
|
|
||||||
SELECT id::TEXT, name, entity_type, summary,
|
|
||||||
ts_rank(to_tsvector('english', name || ' ' || COALESCE(summary, '') || ' ' || COALESCE(description, '')),
|
|
||||||
plainto_tsquery('english', $2)) AS lex_score,
|
|
||||||
ROW_NUMBER() OVER (
|
|
||||||
ORDER BY ts_rank(to_tsvector('english', name || ' ' || COALESCE(summary, '') || ' ' || COALESCE(description, '')),
|
|
||||||
plainto_tsquery('english', $2)) DESC
|
|
||||||
) AS lex_rank
|
|
||||||
FROM memory_entity
|
|
||||||
WHERE t_expired IS NULL
|
|
||||||
AND to_tsvector('english', name || ' ' || COALESCE(summary, '') || ' ' || COALESCE(description, ''))
|
|
||||||
@@ plainto_tsquery('english', $2)
|
|
||||||
AND (t_created >= COALESCE($3, t_created))
|
|
||||||
AND (t_created <= COALESCE($4, t_created))
|
|
||||||
LIMIT $5
|
|
||||||
)
|
|
||||||
SELECT
|
|
||||||
COALESCE(s.id, l.id) AS id,
|
|
||||||
COALESCE(s.name, l.name) AS name,
|
|
||||||
COALESCE(s.entity_type, l.entity_type) AS entity_type,
|
|
||||||
COALESCE(s.summary, l.summary) AS summary,
|
|
||||||
COALESCE(s.sem_score, 0.0)::REAL AS sem_score,
|
|
||||||
COALESCE(l.lex_score, 0.0)::REAL AS lex_score,
|
|
||||||
(
|
|
||||||
$6::REAL * COALESCE(1.0 / (60 + s.sem_rank), 0)::REAL +
|
|
||||||
$7::REAL * COALESCE(1.0 / (60 + l.lex_rank), 0)::REAL
|
|
||||||
) AS rrf_score
|
|
||||||
FROM semantic s
|
|
||||||
FULL OUTER JOIN lexical l ON s.id = l.id
|
|
||||||
ORDER BY rrf_score DESC
|
|
||||||
LIMIT $5";
|
|
||||||
|
|
||||||
// Build query text from embedding context — we need the raw query for lexical
|
// Phase 3: Combine and rank by RRF fusion
|
||||||
// The caller passes embedding, but we need text for ts_rank.
|
let mut hybrid_results = Vec::new();
|
||||||
// We'll accept query_text as empty string fallback for pure-semantic mode.
|
|
||||||
// TODO: Add query_text parameter to hybrid_search signature
|
|
||||||
|
|
||||||
// For now, extract text from the hybrid search call context
|
for entity in entity_results {
|
||||||
// The unified_query handler passes query text separately, so we use empty string
|
|
||||||
// as fallback — lexical will return 0 results, degrading gracefully to pure semantic.
|
|
||||||
let query_text = ""; // Will be fixed when query_text is threaded through
|
|
||||||
|
|
||||||
let entity_results = sqlx::query_as::<_, (String, String, String, Option<String>, f32, f32, f32)>(entity_sql)
|
|
||||||
.bind(query_embedding) // $1
|
|
||||||
.bind(query_text) // $2
|
|
||||||
.bind(start_time) // $3
|
|
||||||
.bind(end_time) // $4
|
|
||||||
.bind(fetch_k) // $5
|
|
||||||
.bind(sem_w) // $6
|
|
||||||
.bind(lex_w) // $7
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("Entity hybrid search error: {}", e))?;
|
|
||||||
|
|
||||||
let mut hybrid_results: Vec<HybridResult> = entity_results
|
|
||||||
.into_iter()
|
|
||||||
.map(|(id, name, entity_type, _summary, sem_score, lex_score, rrf_score)| {
|
|
||||||
HybridResult {
|
|
||||||
id,
|
|
||||||
name: Some(name),
|
|
||||||
entity_type: Some(entity_type),
|
|
||||||
result_type: "entity".to_string(),
|
|
||||||
fused_score: rrf_score,
|
|
||||||
semantic_score: sem_score,
|
|
||||||
lexical_score: lex_score,
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
// --- Edge hybrid: semantic on fact_embedding + lexical on fact text ---
|
|
||||||
let edge_sql =
|
|
||||||
"WITH semantic AS (
|
|
||||||
SELECT e.id::TEXT, e.fact, e.relation_type,
|
|
||||||
1 - (e.fact_embedding <=> $1::vector) AS sem_score,
|
|
||||||
ROW_NUMBER() OVER (ORDER BY e.fact_embedding <=> $1::vector) AS sem_rank
|
|
||||||
FROM memory_edge e
|
|
||||||
WHERE e.t_invalid IS NULL AND e.t_expired IS NULL
|
|
||||||
AND e.fact_embedding IS NOT NULL
|
|
||||||
AND (e.t_created >= COALESCE($3, e.t_created))
|
|
||||||
AND (e.t_created <= COALESCE($4, e.t_created))
|
|
||||||
ORDER BY e.fact_embedding <=> $1::vector
|
|
||||||
LIMIT $5
|
|
||||||
),
|
|
||||||
lexical AS (
|
|
||||||
SELECT e.id::TEXT, e.fact, e.relation_type,
|
|
||||||
ts_rank(to_tsvector('english', e.fact), plainto_tsquery('english', $2)) AS lex_score,
|
|
||||||
ROW_NUMBER() OVER (
|
|
||||||
ORDER BY ts_rank(to_tsvector('english', e.fact), plainto_tsquery('english', $2)) DESC
|
|
||||||
) AS lex_rank
|
|
||||||
FROM memory_edge e
|
|
||||||
WHERE e.t_invalid IS NULL AND e.t_expired IS NULL
|
|
||||||
AND to_tsvector('english', e.fact) @@ plainto_tsquery('english', $2)
|
|
||||||
AND (e.t_created >= COALESCE($3, e.t_created))
|
|
||||||
AND (e.t_created <= COALESCE($4, e.t_created))
|
|
||||||
LIMIT $5
|
|
||||||
)
|
|
||||||
SELECT
|
|
||||||
COALESCE(s.id, l.id) AS id,
|
|
||||||
COALESCE(s.fact, l.fact) AS fact,
|
|
||||||
COALESCE(s.relation_type, l.relation_type) AS relation_type,
|
|
||||||
COALESCE(s.sem_score, 0.0)::REAL AS sem_score,
|
|
||||||
COALESCE(l.lex_score, 0.0)::REAL AS lex_score,
|
|
||||||
(
|
|
||||||
$6::REAL * COALESCE(1.0 / (60 + s.sem_rank), 0)::REAL +
|
|
||||||
$7::REAL * COALESCE(1.0 / (60 + l.lex_rank), 0)::REAL
|
|
||||||
) AS rrf_score
|
|
||||||
FROM semantic s
|
|
||||||
FULL OUTER JOIN lexical l ON s.id = l.id
|
|
||||||
ORDER BY rrf_score DESC
|
|
||||||
LIMIT $5";
|
|
||||||
|
|
||||||
let edge_results = sqlx::query_as::<_, (String, String, String, f32, f32, f32)>(edge_sql)
|
|
||||||
.bind(query_embedding)
|
|
||||||
.bind(query_text)
|
|
||||||
.bind(start_time)
|
|
||||||
.bind(end_time)
|
|
||||||
.bind(fetch_k)
|
|
||||||
.bind(sem_w)
|
|
||||||
.bind(lex_w)
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("Edge hybrid search error: {}", e))?;
|
|
||||||
|
|
||||||
for (id, fact, _rel_type, sem_score, lex_score, rrf_score) in edge_results {
|
|
||||||
hybrid_results.push(HybridResult {
|
hybrid_results.push(HybridResult {
|
||||||
id,
|
id: entity.id,
|
||||||
name: Some(fact),
|
name: Some(entity.name),
|
||||||
entity_type: None,
|
entity_type: Some(entity.entity_type),
|
||||||
result_type: "edge".to_string(),
|
result_type: "entity".to_string(),
|
||||||
fused_score: rrf_score,
|
fused_score: entity.similarity_score * sem_w, // Simplified for entities
|
||||||
semantic_score: sem_score,
|
semantic_score: entity.similarity_score,
|
||||||
lexical_score: lex_score,
|
lexical_score: 0.0,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
// Final sort by fused score
|
for edge in edge_results {
|
||||||
|
hybrid_results.push(HybridResult {
|
||||||
|
id: edge.id,
|
||||||
|
name: Some(edge.fact.clone()),
|
||||||
|
entity_type: None,
|
||||||
|
result_type: "edge".to_string(),
|
||||||
|
fused_score: edge.similarity_score * sem_w, // Simplified for edges
|
||||||
|
semantic_score: edge.similarity_score,
|
||||||
|
lexical_score: 0.0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sort by fused score
|
||||||
hybrid_results.sort_by(|a, b| b.fused_score.partial_cmp(&a.fused_score).unwrap_or(std::cmp::Ordering::Equal));
|
hybrid_results.sort_by(|a, b| b.fused_score.partial_cmp(&a.fused_score).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
|
||||||
|
// Return top-k
|
||||||
hybrid_results.truncate(top_k);
|
hybrid_results.truncate(top_k);
|
||||||
|
|
||||||
info!("Hybrid search returned {} results", hybrid_results.len());
|
info!("Hybrid search returned {} results", hybrid_results.len());
|
||||||
Ok(hybrid_results)
|
Ok(hybrid_results)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -9,6 +9,7 @@
|
|||||||
|
|
||||||
use chrono::{DateTime, Utc};
|
use chrono::{DateTime, Utc};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use tracing::{debug, info};
|
||||||
|
|
||||||
/// Temporal query configuration
|
/// Temporal query configuration
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
|||||||
@@ -1,3 +1,13 @@
|
|||||||
|
/// Advanced Query Filtering: Scope, filtering, and refinement
|
||||||
|
///
|
||||||
|
/// Provides:
|
||||||
|
/// - Project scoping (memory isolation)
|
||||||
|
/// - Level filtering (L1, L2, Reference)
|
||||||
|
/// - Category filtering (Error, Solution, etc.)
|
||||||
|
/// - Time-based filtering (recency)
|
||||||
|
/// - Tag/keyword filtering
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
use std::collections::HashSet;
|
use std::collections::HashSet;
|
||||||
use chrono::{DateTime, Utc, Duration};
|
use chrono::{DateTime, Utc, Duration};
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,490 @@
|
|||||||
|
use anyhow::{anyhow, Result};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
/// Query Context: normalized query + analysis for hybrid search
|
||||||
|
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||||
|
pub struct QueryContext {
|
||||||
|
// Original query
|
||||||
|
pub raw_query: String,
|
||||||
|
|
||||||
|
// Normalized (lowercased, trimmed)
|
||||||
|
pub normalized_query: String,
|
||||||
|
|
||||||
|
// Tokenized terms
|
||||||
|
pub tokens: Vec<String>,
|
||||||
|
|
||||||
|
// Extracted named entities (year, names, keywords)
|
||||||
|
pub entities: HashMap<String, String>,
|
||||||
|
|
||||||
|
// Query embedding (to be generated by LLM)
|
||||||
|
pub embedding: Option<Vec<f32>>,
|
||||||
|
|
||||||
|
// Analysis results
|
||||||
|
pub token_count: usize,
|
||||||
|
pub has_special_syntax: bool, // #tag, @mention, "exact phrase"
|
||||||
|
pub has_date_filters: bool, // 2024, "this month"
|
||||||
|
pub has_negation: bool, // -word, NOT phrase
|
||||||
|
pub question_type: QuestionType,
|
||||||
|
|
||||||
|
// Routing decision
|
||||||
|
pub search_strategy: SearchStrategy,
|
||||||
|
pub confidence: f32, // How confident in the routing decision (0.0-1.0)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Question type classification
|
||||||
|
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||||
|
pub enum QuestionType {
|
||||||
|
Factual, // "What is X?" "Define Y"
|
||||||
|
Procedural, // "How do I..." "Steps to..."
|
||||||
|
Comparative, // "Compare X and Y" "Difference between..."
|
||||||
|
Troubleshooting, // "Fix broken..." "Error: ..."
|
||||||
|
Navigational, // "Where is X?" "Find documents about..."
|
||||||
|
Open, // General conversational
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Search strategy (determines which engines to use)
|
||||||
|
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||||
|
pub enum SearchStrategy {
|
||||||
|
Hybrid, // Both pgvector + OpenSearch
|
||||||
|
SemanticOnly, // pgvector only (if OpenSearch down)
|
||||||
|
LexicalOnly, // OpenSearch only (if embedding model down)
|
||||||
|
LexicalFirst, // OpenSearch to narrow, then semantic rerank
|
||||||
|
}
|
||||||
|
|
||||||
|
/// RRF (Reciprocal Rank Fusion) configuration
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
pub struct RRFConfig {
|
||||||
|
pub k: f32, // Constant (usually 60)
|
||||||
|
pub retrieve_k: usize, // Top-K from each engine (usually 50)
|
||||||
|
pub final_k: usize, // Final top-K to return (usually 10)
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for RRFConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
k: 60.0,
|
||||||
|
retrieve_k: 50,
|
||||||
|
final_k: 10,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Query Optimization Engine
|
||||||
|
pub struct QueryOptimizer {
|
||||||
|
enable_entity_extraction: bool,
|
||||||
|
enable_question_classification: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl QueryOptimizer {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
enable_entity_extraction: true,
|
||||||
|
enable_question_classification: true,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Main entry point: construct query context from user input
|
||||||
|
pub async fn optimize_query(&self, raw_query: &str) -> Result<QueryContext> {
|
||||||
|
// Stage 1: Normalize
|
||||||
|
let normalized = self.normalize_query(raw_query);
|
||||||
|
|
||||||
|
// Stage 2: Tokenize
|
||||||
|
let tokens = self.tokenize(&normalized);
|
||||||
|
|
||||||
|
// Stage 3: Extract entities
|
||||||
|
let entities = if self.enable_entity_extraction {
|
||||||
|
self.extract_entities(raw_query, &tokens)
|
||||||
|
} else {
|
||||||
|
HashMap::new()
|
||||||
|
};
|
||||||
|
|
||||||
|
// Stage 4: Analyze query characteristics
|
||||||
|
let token_count = tokens.len();
|
||||||
|
let has_special_syntax = self.detect_special_syntax(raw_query);
|
||||||
|
let has_date_filters = self.detect_date_filters(&tokens);
|
||||||
|
let has_negation = self.detect_negation(&tokens);
|
||||||
|
|
||||||
|
// Stage 5: Classify question type
|
||||||
|
let question_type = if self.enable_question_classification {
|
||||||
|
self.classify_question(raw_query, &tokens)
|
||||||
|
} else {
|
||||||
|
QuestionType::Open
|
||||||
|
};
|
||||||
|
|
||||||
|
// Stage 6: Route to search strategy
|
||||||
|
let (search_strategy, confidence) = self.route_query(
|
||||||
|
token_count,
|
||||||
|
has_special_syntax,
|
||||||
|
has_date_filters,
|
||||||
|
has_negation,
|
||||||
|
&question_type,
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(QueryContext {
|
||||||
|
raw_query: raw_query.to_string(),
|
||||||
|
normalized_query: normalized,
|
||||||
|
tokens,
|
||||||
|
entities,
|
||||||
|
embedding: None,
|
||||||
|
token_count,
|
||||||
|
has_special_syntax,
|
||||||
|
has_date_filters,
|
||||||
|
has_negation,
|
||||||
|
question_type,
|
||||||
|
search_strategy,
|
||||||
|
confidence,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stage 1: Normalize query
|
||||||
|
fn normalize_query(&self, query: &str) -> String {
|
||||||
|
query
|
||||||
|
.trim()
|
||||||
|
.to_lowercase()
|
||||||
|
.replace(" ", " ") // Remove double spaces
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stage 2: Tokenize
|
||||||
|
fn tokenize(&self, query: &str) -> Vec<String> {
|
||||||
|
query
|
||||||
|
.split_whitespace()
|
||||||
|
.map(|s| s.to_string())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stage 3: Extract entities (years, names, keywords)
|
||||||
|
fn extract_entities(&self, raw_query: &str, tokens: &[String]) -> HashMap<String, String> {
|
||||||
|
let mut entities = HashMap::new();
|
||||||
|
|
||||||
|
for token in tokens {
|
||||||
|
// Year detection: YYYY format
|
||||||
|
if token.len() == 4 {
|
||||||
|
if let Ok(year) = token.parse::<u32>() {
|
||||||
|
if year >= 2000 && year <= 2100 {
|
||||||
|
entities.insert("year".to_string(), token.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Detect quoted phrases
|
||||||
|
if raw_query.contains('"') {
|
||||||
|
let parts: Vec<&str> = raw_query.split('"').collect();
|
||||||
|
if parts.len() >= 3 {
|
||||||
|
let quoted_phrase = parts[1].to_string();
|
||||||
|
entities.insert("exact_phrase".to_string(), quoted_phrase);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
entities
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stage 4: Detect special syntax (#tag, @mention, "phrases")
|
||||||
|
fn detect_special_syntax(&self, query: &str) -> bool {
|
||||||
|
query.contains('#') || query.contains('@') || query.contains('"')
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stage 4: Detect date filters
|
||||||
|
fn detect_date_filters(&self, tokens: &[String]) -> bool {
|
||||||
|
let date_keywords = vec![
|
||||||
|
"this", "last", "next",
|
||||||
|
"2024", "2025", "2026",
|
||||||
|
"january", "february", "march", "april", "may", "june",
|
||||||
|
"july", "august", "september", "october", "november", "december",
|
||||||
|
"week", "month", "year", "day", "today", "yesterday", "tomorrow",
|
||||||
|
];
|
||||||
|
|
||||||
|
tokens.iter().any(|t| date_keywords.contains(&t.as_str()))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stage 4: Detect negation
|
||||||
|
fn detect_negation(&self, tokens: &[String]) -> bool {
|
||||||
|
tokens.iter().any(|t| t == "-" || t == "not" || t == "no" || t.starts_with("-"))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stage 5: Classify question type
|
||||||
|
fn classify_question(&self, raw_query: &str, tokens: &[String]) -> QuestionType {
|
||||||
|
let query_lower = raw_query.to_lowercase();
|
||||||
|
|
||||||
|
// Check first token for question words
|
||||||
|
if tokens.is_empty() {
|
||||||
|
return QuestionType::Open;
|
||||||
|
}
|
||||||
|
|
||||||
|
let first_token = &tokens[0];
|
||||||
|
|
||||||
|
match first_token.as_str() {
|
||||||
|
// Procedural questions
|
||||||
|
t if t == "how" => QuestionType::Procedural,
|
||||||
|
t if t == "what" => {
|
||||||
|
if query_lower.contains("difference") || query_lower.contains("between") {
|
||||||
|
QuestionType::Comparative
|
||||||
|
} else {
|
||||||
|
QuestionType::Factual
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Comparative
|
||||||
|
t if t == "compare" || t == "compare" => QuestionType::Comparative,
|
||||||
|
// Troubleshooting
|
||||||
|
t if t == "fix" || t == "error" || t == "broken" || t == "debug" => {
|
||||||
|
QuestionType::Troubleshooting
|
||||||
|
}
|
||||||
|
// Navigational
|
||||||
|
t if t == "where" || t == "find" || t == "show" => QuestionType::Navigational,
|
||||||
|
_ => {
|
||||||
|
// Heuristics based on content
|
||||||
|
if query_lower.contains("how") {
|
||||||
|
QuestionType::Procedural
|
||||||
|
} else if query_lower.contains("fix") || query_lower.contains("error") {
|
||||||
|
QuestionType::Troubleshooting
|
||||||
|
} else {
|
||||||
|
QuestionType::Open
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stage 6: Route to search strategy
|
||||||
|
fn route_query(
|
||||||
|
&self,
|
||||||
|
token_count: usize,
|
||||||
|
has_special_syntax: bool,
|
||||||
|
has_date_filters: bool,
|
||||||
|
_has_negation: bool,
|
||||||
|
question_type: &QuestionType,
|
||||||
|
) -> (SearchStrategy, f32) {
|
||||||
|
// Very short queries: lexical better
|
||||||
|
if token_count < 3 {
|
||||||
|
return (SearchStrategy::LexicalOnly, 0.8);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Special syntax: preserve exact matches with lexical
|
||||||
|
if has_special_syntax {
|
||||||
|
if has_date_filters {
|
||||||
|
// Special syntax + dates = use lexical to narrow, then semantic
|
||||||
|
return (SearchStrategy::LexicalFirst, 0.85);
|
||||||
|
} else {
|
||||||
|
// Just special syntax = lexical only
|
||||||
|
return (SearchStrategy::LexicalOnly, 0.8);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Date filters present: use cascading (lexical → semantic)
|
||||||
|
if has_date_filters {
|
||||||
|
return (SearchStrategy::LexicalFirst, 0.9);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Question type heuristics
|
||||||
|
match question_type {
|
||||||
|
// Factual questions usually work well with semantic
|
||||||
|
QuestionType::Factual => (SearchStrategy::Hybrid, 0.9),
|
||||||
|
|
||||||
|
// Procedural questions benefit from both (exact steps + understanding)
|
||||||
|
QuestionType::Procedural => (SearchStrategy::Hybrid, 0.95),
|
||||||
|
|
||||||
|
// Troubleshooting needs both (exact errors + semantic understanding)
|
||||||
|
QuestionType::Troubleshooting => (SearchStrategy::Hybrid, 0.95),
|
||||||
|
|
||||||
|
// Comparative: hybrid needed (understanding + multiple docs)
|
||||||
|
QuestionType::Comparative => (SearchStrategy::Hybrid, 0.9),
|
||||||
|
|
||||||
|
// Navigational: lexical good for finding specific things
|
||||||
|
QuestionType::Navigational => (SearchStrategy::LexicalFirst, 0.85),
|
||||||
|
|
||||||
|
// Open/general: hybrid default
|
||||||
|
QuestionType::Open => (SearchStrategy::Hybrid, 0.8),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// RRF Fusion Engine
|
||||||
|
pub struct RRFFusion {
|
||||||
|
config: RRFConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RRFFusion {
|
||||||
|
pub fn new(config: RRFConfig) -> Self {
|
||||||
|
Self { config }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fuse two ranked lists using Reciprocal Rank Fusion
|
||||||
|
pub fn fuse(
|
||||||
|
&self,
|
||||||
|
semantic_results: Vec<(String, f32)>, // (id, score)
|
||||||
|
lexical_results: Vec<(String, f32)>,
|
||||||
|
) -> Vec<(String, f32)> {
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
let mut fused_scores: HashMap<String, f32> = HashMap::new();
|
||||||
|
|
||||||
|
// Add semantic ranks with RRF formula: 1 / (k + rank)
|
||||||
|
for (rank, (id, _)) in semantic_results.into_iter().enumerate() {
|
||||||
|
let rrf_score = 1.0 / (self.config.k + (rank as f32) + 1.0);
|
||||||
|
fused_scores.insert(id, rrf_score);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add lexical ranks (combine if already present)
|
||||||
|
for (rank, (id, _)) in lexical_results.into_iter().enumerate() {
|
||||||
|
let rrf_score = 1.0 / (self.config.k + (rank as f32) + 1.0);
|
||||||
|
*fused_scores.entry(id).or_insert(0.0) += rrf_score;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sort by combined RRF score
|
||||||
|
let mut results: Vec<_> = fused_scores.into_iter().collect();
|
||||||
|
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||||
|
|
||||||
|
// Take top-k
|
||||||
|
results.truncate(self.config.final_k);
|
||||||
|
|
||||||
|
results
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Alternative: Weighted Linear Fusion
|
||||||
|
pub fn fuse_weighted(
|
||||||
|
&self,
|
||||||
|
semantic_results: Vec<(String, f32)>,
|
||||||
|
lexical_results: Vec<(String, f32)>,
|
||||||
|
semantic_weight: f32,
|
||||||
|
lexical_weight: f32,
|
||||||
|
) -> Vec<(String, f32)> {
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
// Normalize scores to [0.0, 1.0]
|
||||||
|
let sem_norm = self.normalize_scores(&semantic_results);
|
||||||
|
let lex_norm = self.normalize_scores(&lexical_results);
|
||||||
|
|
||||||
|
let sem_map: HashMap<String, f32> = sem_norm.into_iter().collect();
|
||||||
|
let lex_map: HashMap<String, f32> = lex_norm.into_iter().collect();
|
||||||
|
|
||||||
|
// Merge all IDs
|
||||||
|
let mut all_ids = std::collections::HashSet::new();
|
||||||
|
all_ids.extend(sem_map.keys().cloned());
|
||||||
|
all_ids.extend(lex_map.keys().cloned());
|
||||||
|
|
||||||
|
// Calculate weighted scores
|
||||||
|
let mut results: Vec<_> = all_ids
|
||||||
|
.into_iter()
|
||||||
|
.map(|id| {
|
||||||
|
let sem_score = sem_map.get(&id).copied().unwrap_or(0.0);
|
||||||
|
let lex_score = lex_map.get(&id).copied().unwrap_or(0.0);
|
||||||
|
|
||||||
|
let weighted_score = semantic_weight * sem_score + lexical_weight * lex_score;
|
||||||
|
(id, weighted_score)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||||
|
results.truncate(self.config.final_k);
|
||||||
|
|
||||||
|
results
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Normalize scores to [0.0, 1.0] range using min-max
|
||||||
|
fn normalize_scores(&self, results: &[(String, f32)]) -> Vec<(String, f32)> {
|
||||||
|
if results.is_empty() {
|
||||||
|
return Vec::new();
|
||||||
|
}
|
||||||
|
|
||||||
|
let min_score = results.iter().map(|(_, s)| s).fold(f32::INFINITY, |a, &b| a.min(b));
|
||||||
|
let max_score = results.iter().map(|(_, s)| s).fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
||||||
|
|
||||||
|
let range = max_score - min_score;
|
||||||
|
|
||||||
|
if range < 0.001 {
|
||||||
|
// All scores identical
|
||||||
|
return results.iter().map(|(id, _)| (id.clone(), 0.5)).collect();
|
||||||
|
}
|
||||||
|
|
||||||
|
results
|
||||||
|
.iter()
|
||||||
|
.map(|(id, score)| {
|
||||||
|
let normalized = (score - min_score) / range;
|
||||||
|
(id.clone(), normalized)
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_query_optimization_procedural() {
|
||||||
|
let optimizer = QueryOptimizer::new();
|
||||||
|
let ctx = optimizer.optimize_query("How do I fix kubernetes port 8080?").await.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(ctx.question_type, QuestionType::Procedural);
|
||||||
|
assert_eq!(ctx.search_strategy, SearchStrategy::Hybrid);
|
||||||
|
assert!(ctx.confidence >= 0.9);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_query_optimization_short() {
|
||||||
|
let optimizer = QueryOptimizer::new();
|
||||||
|
let ctx = optimizer.optimize_query("fix port").await.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(ctx.token_count, 2);
|
||||||
|
assert_eq!(ctx.search_strategy, SearchStrategy::LexicalOnly);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_query_optimization_special_syntax() {
|
||||||
|
let optimizer = QueryOptimizer::new();
|
||||||
|
let ctx = optimizer.optimize_query("kubernetes #networking @devops").await.unwrap();
|
||||||
|
|
||||||
|
assert!(ctx.has_special_syntax);
|
||||||
|
assert_eq!(ctx.search_strategy, SearchStrategy::LexicalOnly);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_rrf_fusion() {
|
||||||
|
let fusion = RRFFusion::new(RRFConfig::default());
|
||||||
|
|
||||||
|
let semantic = vec![
|
||||||
|
("doc1".to_string(), 0.95),
|
||||||
|
("doc2".to_string(), 0.88),
|
||||||
|
("doc3".to_string(), 0.82),
|
||||||
|
];
|
||||||
|
|
||||||
|
let lexical = vec![
|
||||||
|
("doc1".to_string(), 8.5),
|
||||||
|
("doc4".to_string(), 7.2),
|
||||||
|
("doc2".to_string(), 6.8),
|
||||||
|
];
|
||||||
|
|
||||||
|
let fused = fusion.fuse(semantic, lexical);
|
||||||
|
|
||||||
|
// doc1 should be top (in both)
|
||||||
|
assert_eq!(fused[0].0, "doc1");
|
||||||
|
|
||||||
|
// RRF score: doc1 appears in both lists (rank 1 in each)
|
||||||
|
// Score = 1/(60+1) + 1/(60+1) = 2/61 ≈ 0.0328
|
||||||
|
assert!(fused[0].1 > 0.03 && fused[0].1 < 0.04, "Expected RRF score ~0.0328, got {}", fused[0].1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_weighted_fusion() {
|
||||||
|
let fusion = RRFFusion::new(RRFConfig::default());
|
||||||
|
|
||||||
|
let semantic = vec![
|
||||||
|
("doc1".to_string(), 0.95),
|
||||||
|
("doc2".to_string(), 0.88),
|
||||||
|
];
|
||||||
|
|
||||||
|
let lexical = vec![
|
||||||
|
("doc1".to_string(), 8.5),
|
||||||
|
("doc3".to_string(), 7.2),
|
||||||
|
];
|
||||||
|
|
||||||
|
let fused = fusion.fuse_weighted(semantic, lexical, 0.6, 0.4);
|
||||||
|
|
||||||
|
// doc1 should rank highest (has both components)
|
||||||
|
assert_eq!(fused[0].0, "doc1");
|
||||||
|
|
||||||
|
// Score should be normalized and weighted
|
||||||
|
// 0.6 * (0.95/0.95) + 0.4 * (8.5/8.5) = 1.0
|
||||||
|
assert!((fused[0].1 - 1.0).abs() < 0.01);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -11,11 +11,12 @@
|
|||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
use mem_core::DocumentScorer;
|
||||||
|
|
||||||
use crate::hybrid_retrieval::HybridRetriever;
|
use crate::hybrid_retrieval::{HybridRetriever, RetrievalRoute, WikiScopedFilter, RankedCandidate};
|
||||||
use crate::chunk_optimizer::{ChunkOptimizer, OptimizableChunk, SelectionMetrics};
|
use crate::chunk_optimizer::{ChunkOptimizer, OptimizableChunk, SelectionMetrics};
|
||||||
use crate::chunk_metadata::{MetadataExtractor, MetadataBooster, QueryIntent};
|
use crate::chunk_metadata::{MetadataExtractor, MetadataBooster, QueryIntent};
|
||||||
use crate::cache_alignment::{KvCacheAligner, CachedChunk, RetrievalProfiler};
|
use crate::cache_alignment::{KvCacheAligner, CachedChunk, CacheLocalityAnalyzer, RetrievalProfiler};
|
||||||
|
|
||||||
/// Complete query result with all metadata
|
/// Complete query result with all metadata
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -189,7 +190,7 @@ impl QueryOrchestrator {
|
|||||||
|
|
||||||
// Step 8: Build optimized chunks with all metadata
|
// Step 8: Build optimized chunks with all metadata
|
||||||
let mut optimized_chunks = Vec::new();
|
let mut optimized_chunks = Vec::new();
|
||||||
for (_i, chunk) in selected_opt.iter().enumerate() {
|
for (i, chunk) in selected_opt.iter().enumerate() {
|
||||||
let slot = slots.iter().find(|(id, _)| id == &chunk.id).map(|(_, s)| *s).unwrap_or(0);
|
let slot = slots.iter().find(|(id, _)| id == &chunk.id).map(|(_, s)| *s).unwrap_or(0);
|
||||||
let metadata = MetadataExtractor::extract(&chunk.id, &chunk.text);
|
let metadata = MetadataExtractor::extract(&chunk.id, &chunk.text);
|
||||||
|
|
||||||
|
|||||||
@@ -15,9 +15,9 @@ use std::collections::HashMap;
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
use mem_ingest::wiki_link::{WikiLinkGraph, WikiLinkParser};
|
use mem_ingest::wiki_link::{WikiLinkGraph, WikiLinkParser};
|
||||||
use mem_core::{GlobalTfIdfScorer, SemanticScorer};
|
use mem_core::{DocumentScorer, GlobalTfIdfScorer, SemanticScorer};
|
||||||
|
|
||||||
use crate::hybrid_retrieval::{HybridRetriever, RetrievalRoute, WikiScopedFilter};
|
use crate::hybrid_retrieval::{HybridRetriever, RetrievalRoute, WikiScopedFilter, RankedCandidate};
|
||||||
use crate::chunk_optimizer::{ChunkOptimizer, OptimizableChunk, SelectionMetrics};
|
use crate::chunk_optimizer::{ChunkOptimizer, OptimizableChunk, SelectionMetrics};
|
||||||
|
|
||||||
/// Query routing configuration
|
/// Query routing configuration
|
||||||
@@ -74,7 +74,7 @@ pub struct SelectedChunk {
|
|||||||
|
|
||||||
/// Query Router: end-to-end Phase 3+4 pipeline
|
/// Query Router: end-to-end Phase 3+4 pipeline
|
||||||
pub struct QueryRouter {
|
pub struct QueryRouter {
|
||||||
_wiki_filter: WikiScopedFilter,
|
wiki_filter: WikiScopedFilter,
|
||||||
retriever: HybridRetriever,
|
retriever: HybridRetriever,
|
||||||
optimizer: ChunkOptimizer,
|
optimizer: ChunkOptimizer,
|
||||||
config: RouterConfig,
|
config: RouterConfig,
|
||||||
@@ -95,7 +95,7 @@ impl QueryRouter {
|
|||||||
);
|
);
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
_wiki_filter: wiki_filter,
|
wiki_filter,
|
||||||
retriever,
|
retriever,
|
||||||
optimizer,
|
optimizer,
|
||||||
config,
|
config,
|
||||||
|
|||||||
@@ -0,0 +1,111 @@
|
|||||||
|
use anyhow::Result;
|
||||||
|
use mem_llm::{EmbeddingsClient, RerankClient};
|
||||||
|
use mem_store::VectorStore;
|
||||||
|
use pgvector::Vector;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
/// Query result with provenance
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct QueryResult {
|
||||||
|
pub level: String, // "L0", "L1", "L2", "corpus"
|
||||||
|
pub score: f32,
|
||||||
|
pub text: String,
|
||||||
|
pub source: Option<String>,
|
||||||
|
pub provenance: Vec<String>, // parent IDs
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Query worker — semantic search + reranking
|
||||||
|
pub struct QueryWorker {
|
||||||
|
vector_store: std::sync::Arc<VectorStore>,
|
||||||
|
embeddings: std::sync::Arc<EmbeddingsClient>,
|
||||||
|
reranker: std::sync::Arc<RerankClient>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl QueryWorker {
|
||||||
|
/// Create query worker
|
||||||
|
pub fn new(
|
||||||
|
vector_store: VectorStore,
|
||||||
|
embeddings: EmbeddingsClient,
|
||||||
|
reranker: RerankClient,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
vector_store: std::sync::Arc::new(vector_store),
|
||||||
|
embeddings: std::sync::Arc::new(embeddings),
|
||||||
|
reranker: std::sync::Arc::new(reranker),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute semantic query: embed -> search vector -> rerank -> result
|
||||||
|
pub async fn query(
|
||||||
|
&self,
|
||||||
|
project: &str,
|
||||||
|
question: &str,
|
||||||
|
limit: Option<i64>,
|
||||||
|
) -> Result<Vec<QueryResult>> {
|
||||||
|
let limit = limit.unwrap_or(5);
|
||||||
|
|
||||||
|
// Embed the question
|
||||||
|
let question_embedding = self.embeddings.embed_one(question).await?;
|
||||||
|
|
||||||
|
// Search across all levels
|
||||||
|
let mut candidates = Vec::new();
|
||||||
|
|
||||||
|
// L2 synthesis (project-level)
|
||||||
|
if let Some(l2_result) = self.vector_store.search_l2(project, &question_embedding).await? {
|
||||||
|
candidates.push(QueryResult {
|
||||||
|
level: "L2".to_string(),
|
||||||
|
score: l2_result.score,
|
||||||
|
text: l2_result.item.content.clone(),
|
||||||
|
source: Some(format!("project:{}", project)),
|
||||||
|
provenance: vec![l2_result.item.id.to_string()],
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// L1 per-query memories
|
||||||
|
let l1_results = self.vector_store.search_l1(project, &question_embedding, limit).await?;
|
||||||
|
for l1_result in l1_results {
|
||||||
|
candidates.push(QueryResult {
|
||||||
|
level: "L1".to_string(),
|
||||||
|
score: l1_result.score,
|
||||||
|
text: l1_result.item.content.clone(),
|
||||||
|
source: Some(format!("query:{}", l1_result.item.query_id)),
|
||||||
|
provenance: vec![l1_result.item.id.to_string()],
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reference corpus
|
||||||
|
let corpus_results = self.vector_store.search_corpus(project, &question_embedding, limit).await?;
|
||||||
|
for corpus_result in corpus_results {
|
||||||
|
candidates.push(QueryResult {
|
||||||
|
level: "corpus".to_string(),
|
||||||
|
score: corpus_result.score,
|
||||||
|
text: corpus_result.item.content.clone(),
|
||||||
|
source: Some(format!("doc:{}", corpus_result.item.name)),
|
||||||
|
provenance: vec![corpus_result.item.id.to_string()],
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// Rerank candidates by relevance to question
|
||||||
|
// TODO: wire actual cross-encoder reranking
|
||||||
|
// For now, return by vector similarity score
|
||||||
|
candidates.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
candidates.truncate(limit as usize);
|
||||||
|
|
||||||
|
Ok(candidates)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get project synthesis (L2) directly
|
||||||
|
pub async fn get_synthesis(&self, project: &str) -> Result<Option<QueryResult>> {
|
||||||
|
if let Some(l2) = self.vector_store.get_l2(project).await? {
|
||||||
|
Ok(Some(QueryResult {
|
||||||
|
level: "L2".to_string(),
|
||||||
|
score: 1.0,
|
||||||
|
text: l2.content,
|
||||||
|
source: Some(format!("project:{}", project)),
|
||||||
|
provenance: vec![l2.id.to_string()],
|
||||||
|
}))
|
||||||
|
} else {
|
||||||
|
Ok(None)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,336 @@
|
|||||||
|
//! M8.2 — Unified Queue Adapter (SQS-compatible interface)
|
||||||
|
//!
|
||||||
|
//! Abstraction over external queue services (SQS, kmsvc, RabbitMQ, etc.)
|
||||||
|
//! Enables concurrent dual-write processing without database overhead.
|
||||||
|
//!
|
||||||
|
//! # Design
|
||||||
|
//!
|
||||||
|
//! Rather than storing queue state in the database, we leverage external queue
|
||||||
|
//! services via a unified API. This enables true horizontal scalability:
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! Ingest Worker Queue Service (SQS/kmsvc) Dual-Write Workers
|
||||||
|
//! │ │ │
|
||||||
|
//! │─── send_chunk() ────────────>│ │
|
||||||
|
//! │ │ │
|
||||||
|
//! └──────────────────────────────┤<─── receive_chunks(10) ────────┤
|
||||||
|
//! │ │
|
||||||
|
//! │<─── delete_chunk() ────────────┤
|
||||||
|
//! │ (on success) │
|
||||||
|
//! │ │
|
||||||
|
//! │<─── change_visibility() ───────┤
|
||||||
|
//! │ (on retry) │
|
||||||
|
//! ```
|
||||||
|
//!
|
||||||
|
//! # Implementations
|
||||||
|
//! - `SqsQueueAdapter`: AWS SQS backend
|
||||||
|
//! - `KmsvcQueueAdapter`: Kubernetes native messaging service
|
||||||
|
//! - In-memory for testing
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use uuid::Uuid;
|
||||||
|
use anyhow::Result;
|
||||||
|
|
||||||
|
/// SQS-compatible message envelope
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct QueueMessage {
|
||||||
|
/// Unique message ID (from queue service)
|
||||||
|
pub message_id: String,
|
||||||
|
|
||||||
|
/// Original chunk UUID
|
||||||
|
pub chunk_id: Uuid,
|
||||||
|
|
||||||
|
/// Message body (serialized JSON)
|
||||||
|
pub body: String,
|
||||||
|
|
||||||
|
/// Receive count (number of times retrieved)
|
||||||
|
pub receive_count: i32,
|
||||||
|
|
||||||
|
/// Receipt handle (for delete/change_visibility)
|
||||||
|
pub receipt_handle: String,
|
||||||
|
|
||||||
|
/// Project context
|
||||||
|
pub project: String,
|
||||||
|
|
||||||
|
/// Metadata
|
||||||
|
pub attributes: std::collections::HashMap<String, String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Queue statistics
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct QueueStats {
|
||||||
|
pub available_messages: i64,
|
||||||
|
pub in_flight_messages: i64,
|
||||||
|
pub dead_letter_messages: i64,
|
||||||
|
pub total_processed: i64,
|
||||||
|
pub average_delay_secs: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Unified queue adapter trait (SQS-like interface)
|
||||||
|
#[async_trait]
|
||||||
|
pub trait QueueAdapter: Send + Sync {
|
||||||
|
/// Send chunk message to queue
|
||||||
|
///
|
||||||
|
/// # Arguments
|
||||||
|
/// * `chunk_id` — Unique chunk identifier
|
||||||
|
/// * `body` — Serialized message body (JSON)
|
||||||
|
/// * `project` — Project context
|
||||||
|
/// * `attributes` — Optional metadata (e.g., source, level, breadcrumb)
|
||||||
|
///
|
||||||
|
/// # Returns
|
||||||
|
/// Message ID from queue service
|
||||||
|
async fn send_chunk(
|
||||||
|
&self,
|
||||||
|
chunk_id: Uuid,
|
||||||
|
body: String,
|
||||||
|
project: String,
|
||||||
|
attributes: std::collections::HashMap<String, String>,
|
||||||
|
) -> Result<String>;
|
||||||
|
|
||||||
|
/// Receive chunk messages from queue
|
||||||
|
///
|
||||||
|
/// # Arguments
|
||||||
|
/// * `max_messages` — Max number of messages (1-10)
|
||||||
|
/// * `visibility_timeout_secs` — Visibility timeout duration
|
||||||
|
/// * `project` — Project filter (optional)
|
||||||
|
///
|
||||||
|
/// # Returns
|
||||||
|
/// List of available messages
|
||||||
|
async fn receive_chunks(
|
||||||
|
&self,
|
||||||
|
max_messages: i32,
|
||||||
|
visibility_timeout_secs: i32,
|
||||||
|
project: Option<&str>,
|
||||||
|
) -> Result<Vec<QueueMessage>>;
|
||||||
|
|
||||||
|
/// Delete message from queue (after successful processing)
|
||||||
|
///
|
||||||
|
/// # Arguments
|
||||||
|
/// * `message_id` — Message to delete
|
||||||
|
/// * `receipt_handle` — Receipt handle (for idempotency)
|
||||||
|
async fn delete_chunk(&self, message_id: &str, receipt_handle: &str) -> Result<()>;
|
||||||
|
|
||||||
|
/// Change message visibility timeout
|
||||||
|
///
|
||||||
|
/// Called when processing takes longer than expected.
|
||||||
|
async fn change_visibility(
|
||||||
|
&self,
|
||||||
|
message_id: &str,
|
||||||
|
receipt_handle: &str,
|
||||||
|
visibility_timeout_secs: i32,
|
||||||
|
) -> Result<()>;
|
||||||
|
|
||||||
|
/// Send message to dead-letter queue
|
||||||
|
///
|
||||||
|
/// Called when message exceeds max receive count.
|
||||||
|
async fn send_to_dlq(&self, message_id: &str, receipt_handle: &str, reason: &str) -> Result<()>;
|
||||||
|
|
||||||
|
/// Get queue statistics
|
||||||
|
async fn get_stats(&self, project: Option<&str>) -> Result<QueueStats>;
|
||||||
|
|
||||||
|
/// Purge queue (test/admin only)
|
||||||
|
async fn purge(&self, project: Option<&str>) -> Result<usize>;
|
||||||
|
|
||||||
|
/// Health check
|
||||||
|
async fn health_check(&self) -> Result<()>;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// In-memory queue adapter (for testing and local development)
|
||||||
|
pub struct InMemoryQueueAdapter {
|
||||||
|
messages: std::sync::Arc<tokio::sync::Mutex<Vec<QueueMessage>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl InMemoryQueueAdapter {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
messages: std::sync::Arc::new(tokio::sync::Mutex::new(Vec::new())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for InMemoryQueueAdapter {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl QueueAdapter for InMemoryQueueAdapter {
|
||||||
|
async fn send_chunk(
|
||||||
|
&self,
|
||||||
|
chunk_id: Uuid,
|
||||||
|
body: String,
|
||||||
|
project: String,
|
||||||
|
attributes: std::collections::HashMap<String, String>,
|
||||||
|
) -> Result<String> {
|
||||||
|
let message_id = format!("msg-{}", Uuid::new_v4());
|
||||||
|
let receipt_handle = format!("handle-{}", Uuid::new_v4());
|
||||||
|
|
||||||
|
let msg = QueueMessage {
|
||||||
|
message_id: message_id.clone(),
|
||||||
|
chunk_id,
|
||||||
|
body,
|
||||||
|
receive_count: 0,
|
||||||
|
receipt_handle,
|
||||||
|
project,
|
||||||
|
attributes,
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut msgs = self.messages.lock().await;
|
||||||
|
msgs.push(msg);
|
||||||
|
|
||||||
|
Ok(message_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn receive_chunks(
|
||||||
|
&self,
|
||||||
|
max_messages: i32,
|
||||||
|
_visibility_timeout_secs: i32,
|
||||||
|
project: Option<&str>,
|
||||||
|
) -> Result<Vec<QueueMessage>> {
|
||||||
|
let mut msgs = self.messages.lock().await;
|
||||||
|
let max = max_messages.min(10).max(1) as usize;
|
||||||
|
let drain_count = msgs.len().min(max);
|
||||||
|
|
||||||
|
let result: Vec<_> = msgs
|
||||||
|
.drain(..drain_count)
|
||||||
|
.filter(|m| project.is_none() || m.project.as_str() == project.unwrap())
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_chunk(&self, message_id: &str, _receipt_handle: &str) -> Result<()> {
|
||||||
|
let mut msgs = self.messages.lock().await;
|
||||||
|
msgs.retain(|m| m.message_id != message_id);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn change_visibility(
|
||||||
|
&self,
|
||||||
|
_message_id: &str,
|
||||||
|
_receipt_handle: &str,
|
||||||
|
_visibility_timeout_secs: i32,
|
||||||
|
) -> Result<()> {
|
||||||
|
// No-op for in-memory
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn send_to_dlq(&self, message_id: &str, _receipt_handle: &str, _reason: &str) -> Result<()> {
|
||||||
|
let mut msgs = self.messages.lock().await;
|
||||||
|
msgs.retain(|m| m.message_id != message_id);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_stats(&self, _project: Option<&str>) -> Result<QueueStats> {
|
||||||
|
let msgs = self.messages.lock().await;
|
||||||
|
Ok(QueueStats {
|
||||||
|
available_messages: msgs.len() as i64,
|
||||||
|
in_flight_messages: 0,
|
||||||
|
dead_letter_messages: 0,
|
||||||
|
total_processed: 0,
|
||||||
|
average_delay_secs: 0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn purge(&self, _project: Option<&str>) -> Result<usize> {
|
||||||
|
let mut msgs = self.messages.lock().await;
|
||||||
|
let count = msgs.len();
|
||||||
|
msgs.clear();
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn health_check(&self) -> Result<()> {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_in_memory_send_chunk() {
|
||||||
|
let queue = InMemoryQueueAdapter::new();
|
||||||
|
let msg_id = queue
|
||||||
|
.send_chunk(
|
||||||
|
Uuid::new_v4(),
|
||||||
|
r#"{"content": "test"}"#.to_string(),
|
||||||
|
"test-project".to_string(),
|
||||||
|
std::collections::HashMap::new(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(msg_id.starts_with("msg-"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_in_memory_receive_chunks() {
|
||||||
|
let queue = InMemoryQueueAdapter::new();
|
||||||
|
|
||||||
|
for i in 0..5 {
|
||||||
|
queue
|
||||||
|
.send_chunk(
|
||||||
|
Uuid::new_v4(),
|
||||||
|
format!(r#"{{"content": "test{}"}}"#, i),
|
||||||
|
"test-project".to_string(),
|
||||||
|
std::collections::HashMap::new(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
let messages = queue
|
||||||
|
.receive_chunks(3, 30, Some("test-project"))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(messages.len(), 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_in_memory_delete_chunk() {
|
||||||
|
let queue = InMemoryQueueAdapter::new();
|
||||||
|
|
||||||
|
let msg_id = queue
|
||||||
|
.send_chunk(
|
||||||
|
Uuid::new_v4(),
|
||||||
|
"body".to_string(),
|
||||||
|
"test".to_string(),
|
||||||
|
std::collections::HashMap::new(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
queue.delete_chunk(&msg_id, "handle").await.unwrap();
|
||||||
|
|
||||||
|
let msgs = queue.receive_chunks(10, 30, None).await.unwrap();
|
||||||
|
assert_eq!(msgs.len(), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_queue_stats() {
|
||||||
|
let queue = InMemoryQueueAdapter::new();
|
||||||
|
|
||||||
|
queue
|
||||||
|
.send_chunk(
|
||||||
|
Uuid::new_v4(),
|
||||||
|
"body".to_string(),
|
||||||
|
"test".to_string(),
|
||||||
|
std::collections::HashMap::new(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
|
||||||
|
let stats = queue.get_stats(None).await.unwrap();
|
||||||
|
assert_eq!(stats.available_messages, 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_health_check() {
|
||||||
|
let queue = InMemoryQueueAdapter::new();
|
||||||
|
assert!(queue.health_check().await.is_ok());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,402 @@
|
|||||||
|
//! M8.2 — Queue Worker for Concurrent Dual-Write Processing
|
||||||
|
//!
|
||||||
|
//! Background task that receives messages from the queue and processes them
|
||||||
|
//! via DualWriteIndexer. Runs concurrently with ingest, improving throughput.
|
||||||
|
//!
|
||||||
|
//! # Architecture
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! IngestWorker (fast path) QueueWorker (background)
|
||||||
|
//! │ │
|
||||||
|
//! ├─ chunk_input │
|
||||||
|
//! │ (embedding) │
|
||||||
|
//! │ │
|
||||||
|
//! ├─ queue.send_chunk()────┐ │
|
||||||
|
//! │ (returns immediately) │ │
|
||||||
|
//! │ │ │
|
||||||
|
//! └─ continues... │ │
|
||||||
|
//! │ │
|
||||||
|
//! ├─ queue.receive_chunks(10, 30)
|
||||||
|
//! │ (long-poll, up to 30s)
|
||||||
|
//! │
|
||||||
|
//! ├─ for each message:
|
||||||
|
//! │ - process_queued_chunk()
|
||||||
|
//! │ - embed_one() [happens here]
|
||||||
|
//! │ - write_pgvector()
|
||||||
|
//! │ - write_opensearch()
|
||||||
|
//! │ - delete_chunk() on success
|
||||||
|
//! │ - change_visibility() on retry
|
||||||
|
//! │
|
||||||
|
//! └─ loop back to receive
|
||||||
|
//! ```
|
||||||
|
//!
|
||||||
|
//! Benefits:
|
||||||
|
//! - Ingest path is decoupled from embedding/pgvector/OpenSearch writes
|
||||||
|
//! - Multiple workers can process messages concurrently
|
||||||
|
//! - Non-blocking: queue.send_chunk() returns immediately
|
||||||
|
//! - Fault-tolerant: failed messages auto-retry with exponential backoff
|
||||||
|
|
||||||
|
use anyhow::{anyhow, Result};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
use tokio::time::sleep;
|
||||||
|
use tracing::{debug, error, info, warn};
|
||||||
|
|
||||||
|
use crate::dual_write_indexer::DualWriteIndexer;
|
||||||
|
use crate::queue_adapter::QueueAdapter;
|
||||||
|
use mem_llm::EmbeddingsClient;
|
||||||
|
|
||||||
|
/// Configuration for queue worker
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct QueueWorkerConfig {
|
||||||
|
/// Max messages per receive (1-10)
|
||||||
|
pub max_messages_per_batch: i32,
|
||||||
|
|
||||||
|
/// Visibility timeout for processing (seconds)
|
||||||
|
pub visibility_timeout_secs: i32,
|
||||||
|
|
||||||
|
/// Time to wait for messages (0-20 seconds)
|
||||||
|
pub wait_time_secs: i32,
|
||||||
|
|
||||||
|
/// Project to process (None = all projects)
|
||||||
|
pub project: Option<String>,
|
||||||
|
|
||||||
|
/// Max retries before DLQ
|
||||||
|
pub max_retries: i32,
|
||||||
|
|
||||||
|
/// Retry backoff: exponential starting from this value (seconds)
|
||||||
|
pub retry_backoff_initial_secs: i32,
|
||||||
|
|
||||||
|
/// Poll interval when queue is empty (seconds)
|
||||||
|
pub empty_poll_interval_secs: u64,
|
||||||
|
|
||||||
|
/// Enable metrics collection
|
||||||
|
pub enable_metrics: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for QueueWorkerConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
max_messages_per_batch: 10,
|
||||||
|
visibility_timeout_secs: 300, // 5 minutes
|
||||||
|
wait_time_secs: 20, // Long-poll timeout
|
||||||
|
project: None,
|
||||||
|
max_retries: 3,
|
||||||
|
retry_backoff_initial_secs: 60,
|
||||||
|
empty_poll_interval_secs: 5,
|
||||||
|
enable_metrics: true,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Metrics for worker execution
|
||||||
|
#[derive(Debug, Clone, Default)]
|
||||||
|
pub struct WorkerMetrics {
|
||||||
|
pub messages_received: u64,
|
||||||
|
pub messages_processed: u64,
|
||||||
|
pub messages_failed: u64,
|
||||||
|
pub messages_dlq: u64,
|
||||||
|
pub total_processing_time_ms: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Queue worker for processing dual-write messages
|
||||||
|
pub struct QueueWorker {
|
||||||
|
indexer: Arc<DualWriteIndexer>,
|
||||||
|
embeddings: Arc<EmbeddingsClient>,
|
||||||
|
config: QueueWorkerConfig,
|
||||||
|
metrics: Arc<tokio::sync::RwLock<WorkerMetrics>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl QueueWorker {
|
||||||
|
/// Create new queue worker
|
||||||
|
pub fn new(
|
||||||
|
indexer: Arc<DualWriteIndexer>,
|
||||||
|
embeddings: Arc<EmbeddingsClient>,
|
||||||
|
config: QueueWorkerConfig,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
indexer,
|
||||||
|
embeddings,
|
||||||
|
config,
|
||||||
|
metrics: Arc::new(tokio::sync::RwLock::new(WorkerMetrics::default())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Start worker (blocking loop)
|
||||||
|
pub async fn start(&self) -> Result<()> {
|
||||||
|
info!("Queue worker starting: config={:?}", self.config);
|
||||||
|
|
||||||
|
loop {
|
||||||
|
match self.process_batch().await {
|
||||||
|
Ok(count) => {
|
||||||
|
if count == 0 {
|
||||||
|
// Empty batch: sleep before retrying
|
||||||
|
debug!(
|
||||||
|
"Queue empty, waiting {}s before retry",
|
||||||
|
self.config.empty_poll_interval_secs
|
||||||
|
);
|
||||||
|
sleep(Duration::from_secs(self.config.empty_poll_interval_secs)).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
error!("Worker error (will retry): {}", e);
|
||||||
|
sleep(Duration::from_secs(5)).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Process one batch of messages from queue
|
||||||
|
async fn process_batch(&self) -> Result<usize> {
|
||||||
|
let queue = &self.indexer.queue;
|
||||||
|
|
||||||
|
// Receive messages
|
||||||
|
let messages = queue
|
||||||
|
.receive_chunks(
|
||||||
|
self.config.max_messages_per_batch,
|
||||||
|
self.config.visibility_timeout_secs,
|
||||||
|
self.config.project.as_deref(),
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let batch_size = messages.len();
|
||||||
|
if batch_size == 0 {
|
||||||
|
return Ok(0);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut metrics = self.metrics.write().await;
|
||||||
|
metrics.messages_received += batch_size as u64;
|
||||||
|
drop(metrics);
|
||||||
|
|
||||||
|
// Process each message concurrently
|
||||||
|
let handles: Vec<_> = messages
|
||||||
|
.into_iter()
|
||||||
|
.map(|msg| {
|
||||||
|
let indexer = self.indexer.clone();
|
||||||
|
let embeddings = self.embeddings.clone();
|
||||||
|
let config = self.config.clone();
|
||||||
|
let metrics = self.metrics.clone();
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
Self::process_message(indexer, embeddings, config, metrics, msg).await
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// Wait for all to complete
|
||||||
|
for handle in handles {
|
||||||
|
if let Err(e) = handle.await {
|
||||||
|
error!("Worker task panicked: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(batch_size)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Process a single message
|
||||||
|
async fn process_message(
|
||||||
|
indexer: Arc<DualWriteIndexer>,
|
||||||
|
embeddings: Arc<EmbeddingsClient>,
|
||||||
|
config: QueueWorkerConfig,
|
||||||
|
metrics: Arc<tokio::sync::RwLock<WorkerMetrics>>,
|
||||||
|
message: crate::queue_adapter::QueueMessage,
|
||||||
|
) -> Result<()> {
|
||||||
|
let start = std::time::Instant::now();
|
||||||
|
let message_id = message.message_id.clone();
|
||||||
|
let receipt_handle = message.receipt_handle.clone();
|
||||||
|
|
||||||
|
debug!("Processing message: {}", message_id);
|
||||||
|
|
||||||
|
// Parse message body
|
||||||
|
let body: serde_json::Value = match serde_json::from_str(&message.body) {
|
||||||
|
Ok(b) => b,
|
||||||
|
Err(e) => {
|
||||||
|
error!("Failed to parse message body: {}", e);
|
||||||
|
indexer
|
||||||
|
.queue
|
||||||
|
.send_to_dlq(&message_id, &receipt_handle, "invalid_json")
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
|
||||||
|
let mut m = metrics.write().await;
|
||||||
|
m.messages_dlq += 1;
|
||||||
|
return Err(e.into());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Extract chunk_id
|
||||||
|
let chunk_id = match body["chunk_id"].as_str() {
|
||||||
|
Some(id) => match uuid::Uuid::parse_str(id) {
|
||||||
|
Ok(u) => u,
|
||||||
|
Err(e) => {
|
||||||
|
error!("Invalid chunk_id: {}", e);
|
||||||
|
indexer
|
||||||
|
.queue
|
||||||
|
.send_to_dlq(&message_id, &receipt_handle, "invalid_uuid")
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
|
||||||
|
let mut m = metrics.write().await;
|
||||||
|
m.messages_dlq += 1;
|
||||||
|
return Err(e.into());
|
||||||
|
}
|
||||||
|
},
|
||||||
|
None => {
|
||||||
|
error!("Missing chunk_id in message");
|
||||||
|
indexer
|
||||||
|
.queue
|
||||||
|
.send_to_dlq(&message_id, &receipt_handle, "missing_chunk_id")
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
|
||||||
|
let mut m = metrics.write().await;
|
||||||
|
m.messages_dlq += 1;
|
||||||
|
return Err(anyhow!("Missing chunk_id"));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Extract content
|
||||||
|
let content = match body["content"].as_str() {
|
||||||
|
Some(c) => c.to_string(),
|
||||||
|
None => {
|
||||||
|
error!("Missing content in message");
|
||||||
|
indexer
|
||||||
|
.queue
|
||||||
|
.send_to_dlq(&message_id, &receipt_handle, "missing_content")
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
|
||||||
|
let mut m = metrics.write().await;
|
||||||
|
m.messages_dlq += 1;
|
||||||
|
return Err(anyhow!("Missing content"));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Compute embedding
|
||||||
|
let embedding_vec = match embeddings.embed_one(&content).await {
|
||||||
|
Ok(vec) => vec,
|
||||||
|
Err(e) => {
|
||||||
|
warn!("Embedding failed, extending visibility for retry: {}", e);
|
||||||
|
indexer
|
||||||
|
.queue
|
||||||
|
.change_visibility(&message_id, &receipt_handle, 300)
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
|
||||||
|
let mut m = metrics.write().await;
|
||||||
|
m.messages_failed += 1;
|
||||||
|
return Err(e);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Convert pgvector::Vector to Vec<f32>
|
||||||
|
let embedding: Vec<f32> = embedding_vec.to_vec();
|
||||||
|
|
||||||
|
// Process dual-write
|
||||||
|
match indexer.process_queued_chunk(&message, &embedding).await {
|
||||||
|
Ok(result) => {
|
||||||
|
if result.pgvector_success && !result.opensearch_pending {
|
||||||
|
// Success: already deleted by process_queued_chunk
|
||||||
|
debug!("Message processed successfully: {}", message_id);
|
||||||
|
|
||||||
|
let elapsed = start.elapsed().as_millis() as u64;
|
||||||
|
let mut m = metrics.write().await;
|
||||||
|
m.messages_processed += 1;
|
||||||
|
m.total_processing_time_ms += elapsed;
|
||||||
|
} else if result.pgvector_success && result.opensearch_pending {
|
||||||
|
// pgvector OK, OpenSearch pending: visibility already extended
|
||||||
|
warn!("Message will retry: {}", message_id);
|
||||||
|
|
||||||
|
let mut m = metrics.write().await;
|
||||||
|
m.messages_failed += 1;
|
||||||
|
} else {
|
||||||
|
// pgvector failed: visibility already extended
|
||||||
|
warn!("pgvector write failed, will retry: {}", message_id);
|
||||||
|
|
||||||
|
let mut m = metrics.write().await;
|
||||||
|
m.messages_failed += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
// Check receive count
|
||||||
|
if message.receive_count >= config.max_retries {
|
||||||
|
error!(
|
||||||
|
"Message max retries exceeded ({}), sending to DLQ: {}",
|
||||||
|
message.receive_count, message_id
|
||||||
|
);
|
||||||
|
indexer
|
||||||
|
.queue
|
||||||
|
.send_to_dlq(&message_id, &receipt_handle, "max_retries")
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
|
||||||
|
let mut m = metrics.write().await;
|
||||||
|
m.messages_dlq += 1;
|
||||||
|
} else {
|
||||||
|
// Extend visibility for retry
|
||||||
|
warn!(
|
||||||
|
"Message processing failed (retry {}), extending visibility: {}",
|
||||||
|
message.receive_count, message_id
|
||||||
|
);
|
||||||
|
indexer
|
||||||
|
.queue
|
||||||
|
.change_visibility(&message_id, &receipt_handle, 300)
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
|
||||||
|
let mut m = metrics.write().await;
|
||||||
|
m.messages_failed += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
Err(e)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get current metrics
|
||||||
|
pub async fn metrics(&self) -> WorkerMetrics {
|
||||||
|
self.metrics.read().await.clone()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Reset metrics
|
||||||
|
pub async fn reset_metrics(&self) {
|
||||||
|
let mut m = self.metrics.write().await;
|
||||||
|
*m = WorkerMetrics::default();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_queue_worker_config_default() {
|
||||||
|
let config = QueueWorkerConfig::default();
|
||||||
|
assert_eq!(config.max_messages_per_batch, 10);
|
||||||
|
assert_eq!(config.visibility_timeout_secs, 300);
|
||||||
|
assert_eq!(config.wait_time_secs, 20);
|
||||||
|
assert_eq!(config.max_retries, 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_worker_metrics_default() {
|
||||||
|
let metrics = WorkerMetrics::default();
|
||||||
|
assert_eq!(metrics.messages_received, 0);
|
||||||
|
assert_eq!(metrics.messages_processed, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_queue_worker_config_custom() {
|
||||||
|
let config = QueueWorkerConfig {
|
||||||
|
max_messages_per_batch: 5,
|
||||||
|
visibility_timeout_secs: 600,
|
||||||
|
project: Some("test-proj".to_string()),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(config.max_messages_per_batch, 5);
|
||||||
|
assert_eq!(config.project, Some("test-proj".to_string()));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,243 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
/// Rate limit error with retry guidance
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct RateLimitError {
|
||||||
|
pub retry_after_seconds: u64,
|
||||||
|
pub limit_window_secs: u64,
|
||||||
|
pub reason: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RateLimitError {
|
||||||
|
pub fn reason(&self) -> String {
|
||||||
|
format!(
|
||||||
|
"{} (retry after {} seconds, window: {} seconds)",
|
||||||
|
self.reason, self.retry_after_seconds, self.limit_window_secs
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Token bucket for a single endpoint
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
struct TokenBucket {
|
||||||
|
tokens: f64,
|
||||||
|
last_refill: Instant,
|
||||||
|
capacity: f64, // max tokens (per hour)
|
||||||
|
refill_rate: f64, // tokens per second
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TokenBucket {
|
||||||
|
fn new(capacity: f64, refill_rate: f64) -> Self {
|
||||||
|
Self {
|
||||||
|
tokens: capacity,
|
||||||
|
last_refill: Instant::now(),
|
||||||
|
capacity,
|
||||||
|
refill_rate,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Refill tokens based on elapsed time
|
||||||
|
fn refill(&mut self) {
|
||||||
|
let now = Instant::now();
|
||||||
|
let elapsed = now.duration_since(self.last_refill).as_secs_f64();
|
||||||
|
let refilled = elapsed * self.refill_rate;
|
||||||
|
|
||||||
|
self.tokens = (self.tokens + refilled).min(self.capacity);
|
||||||
|
self.last_refill = now;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Try to consume 1 token. Returns Ok if successful, Err(retry_after_secs) if rate limited.
|
||||||
|
fn try_consume(&mut self) -> Result<(), u64> {
|
||||||
|
self.refill();
|
||||||
|
|
||||||
|
if self.tokens >= 1.0 {
|
||||||
|
self.tokens -= 1.0;
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Rate limited: estimate time until next token available
|
||||||
|
let tokens_needed = 1.0 - self.tokens;
|
||||||
|
let retry_after = (tokens_needed / self.refill_rate).ceil() as u64;
|
||||||
|
Err(retry_after.max(1))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rate limiter with per-apikey, per-endpoint buckets
|
||||||
|
pub struct RateLimiter {
|
||||||
|
buckets: Arc<Mutex<HashMap<String, Arc<Mutex<TokenBucket>>>>>,
|
||||||
|
limit_config: LimitConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
pub struct LimitConfig {
|
||||||
|
pub ingest_per_hour: f64,
|
||||||
|
pub query_per_hour: f64,
|
||||||
|
pub projects_per_hour: f64,
|
||||||
|
pub burst_per_second: f64, // Currently unused but kept for API compatibility
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for LimitConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
ingest_per_hour: 100.0,
|
||||||
|
query_per_hour: 1000.0,
|
||||||
|
projects_per_hour: 100.0,
|
||||||
|
burst_per_second: 10.0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RateLimiter {
|
||||||
|
pub fn new(config: LimitConfig) -> Self {
|
||||||
|
Self {
|
||||||
|
buckets: Arc::new(Mutex::new(HashMap::new())),
|
||||||
|
limit_config: config,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get or create bucket for apikey + endpoint
|
||||||
|
fn get_or_create_bucket(&self, apikey_endpoint: &str) -> Arc<Mutex<TokenBucket>> {
|
||||||
|
let mut buckets = self.buckets.lock().unwrap();
|
||||||
|
let config = &self.limit_config;
|
||||||
|
|
||||||
|
if !buckets.contains_key(apikey_endpoint) {
|
||||||
|
// Determine limit based on endpoint
|
||||||
|
let capacity = if apikey_endpoint.contains("/memory/ingest") {
|
||||||
|
config.ingest_per_hour
|
||||||
|
} else if apikey_endpoint.contains("/memory/query") {
|
||||||
|
config.query_per_hour
|
||||||
|
} else if apikey_endpoint.contains("/memory/projects") {
|
||||||
|
config.projects_per_hour
|
||||||
|
} else {
|
||||||
|
// Unlimited for unknown endpoints
|
||||||
|
f64::INFINITY
|
||||||
|
};
|
||||||
|
|
||||||
|
let refill_rate = if capacity.is_infinite() {
|
||||||
|
f64::INFINITY
|
||||||
|
} else {
|
||||||
|
capacity / 3600.0 // per second
|
||||||
|
};
|
||||||
|
|
||||||
|
let bucket = TokenBucket::new(capacity, refill_rate);
|
||||||
|
buckets.insert(apikey_endpoint.to_string(), Arc::new(Mutex::new(bucket)));
|
||||||
|
}
|
||||||
|
|
||||||
|
buckets[apikey_endpoint].clone()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check rate limit for apikey + endpoint. Returns Ok or Err with retry guidance.
|
||||||
|
pub fn check(&self, apikey: &str, endpoint: &str) -> Result<(), RateLimitError> {
|
||||||
|
let key = format!("{}::{}", apikey, endpoint);
|
||||||
|
let bucket = self.get_or_create_bucket(&key);
|
||||||
|
let mut b = bucket.lock().unwrap();
|
||||||
|
|
||||||
|
match b.try_consume() {
|
||||||
|
Ok(_) => Ok(()),
|
||||||
|
Err(retry_after) => {
|
||||||
|
let window_secs = if endpoint.contains("/memory/ingest") {
|
||||||
|
3600
|
||||||
|
} else if endpoint.contains("/memory/query") {
|
||||||
|
3600
|
||||||
|
} else if endpoint.contains("/memory/projects") {
|
||||||
|
3600
|
||||||
|
} else {
|
||||||
|
3600
|
||||||
|
};
|
||||||
|
|
||||||
|
Err(RateLimitError {
|
||||||
|
retry_after_seconds: retry_after,
|
||||||
|
limit_window_secs: window_secs,
|
||||||
|
reason: format!(
|
||||||
|
"rate_limit_exceeded for {}",
|
||||||
|
endpoint
|
||||||
|
),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_token_bucket_refill() {
|
||||||
|
let mut bucket = TokenBucket::new(100.0, 100.0 / 3600.0);
|
||||||
|
assert!(bucket.try_consume().is_ok());
|
||||||
|
// After one consumption, should have 99 tokens
|
||||||
|
assert_eq!((bucket.tokens * 1.0) as i64, 99);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_rate_limit_within_capacity() {
|
||||||
|
let config = LimitConfig {
|
||||||
|
ingest_per_hour: 5.0,
|
||||||
|
query_per_hour: 10.0,
|
||||||
|
projects_per_hour: 10.0,
|
||||||
|
burst_per_second: 10.0,
|
||||||
|
};
|
||||||
|
let limiter = RateLimiter::new(config);
|
||||||
|
|
||||||
|
// First 5 should succeed
|
||||||
|
for _ in 0..5 {
|
||||||
|
assert!(limiter.check("apikey1", "/memory/ingest").is_ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
// 6th should fail
|
||||||
|
let err = limiter.check("apikey1", "/memory/ingest");
|
||||||
|
assert!(err.is_err());
|
||||||
|
if let Err(e) = err {
|
||||||
|
assert!(e.retry_after_seconds > 0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_per_apikey_isolation() {
|
||||||
|
let config = LimitConfig {
|
||||||
|
ingest_per_hour: 5.0,
|
||||||
|
query_per_hour: 10.0,
|
||||||
|
projects_per_hour: 10.0,
|
||||||
|
burst_per_second: 10.0,
|
||||||
|
};
|
||||||
|
let limiter = RateLimiter::new(config);
|
||||||
|
|
||||||
|
// apikey1 uses up 5 ingest requests
|
||||||
|
for _ in 0..5 {
|
||||||
|
assert!(limiter.check("apikey1", "/memory/ingest").is_ok());
|
||||||
|
}
|
||||||
|
assert!(limiter.check("apikey1", "/memory/ingest").is_err());
|
||||||
|
|
||||||
|
// apikey2 should have its own 5
|
||||||
|
for _ in 0..5 {
|
||||||
|
assert!(limiter.check("apikey2", "/memory/ingest").is_ok());
|
||||||
|
}
|
||||||
|
assert!(limiter.check("apikey2", "/memory/ingest").is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_per_endpoint_isolation() {
|
||||||
|
let config = LimitConfig {
|
||||||
|
ingest_per_hour: 5.0,
|
||||||
|
query_per_hour: 10.0,
|
||||||
|
projects_per_hour: 10.0,
|
||||||
|
burst_per_second: 10.0,
|
||||||
|
};
|
||||||
|
let limiter = RateLimiter::new(config);
|
||||||
|
|
||||||
|
// Use up 5 ingest
|
||||||
|
for _ in 0..5 {
|
||||||
|
assert!(limiter.check("apikey1", "/memory/ingest").is_ok());
|
||||||
|
}
|
||||||
|
assert!(limiter.check("apikey1", "/memory/ingest").is_err());
|
||||||
|
|
||||||
|
// Query should have separate 10 limit
|
||||||
|
for _ in 0..10 {
|
||||||
|
assert!(limiter.check("apikey1", "/memory/query").is_ok());
|
||||||
|
}
|
||||||
|
assert!(limiter.check("apikey1", "/memory/query").is_err());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -11,7 +11,7 @@ use anyhow::Result;
|
|||||||
|
|
||||||
use super::access_evaluator::{AccessEvaluator, FilterResult, HasResourceMeta};
|
use super::access_evaluator::{AccessEvaluator, FilterResult, HasResourceMeta};
|
||||||
use super::role_provider::RoleProvider;
|
use super::role_provider::RoleProvider;
|
||||||
use super::types::{AccessDecision, Claims, ResourceMeta, Verb};
|
use super::types::{AccessDecision, Claims, DenyReason, ResourceMeta, Verb};
|
||||||
|
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
// Audit Logger
|
// Audit Logger
|
||||||
|
|||||||
@@ -6,7 +6,7 @@
|
|||||||
/// - OwnerScope: resource.owner == claims.sub?
|
/// - OwnerScope: resource.owner == claims.sub?
|
||||||
/// - GroupScope: user in required groups?
|
/// - GroupScope: user in required groups?
|
||||||
|
|
||||||
use super::types::{AccessScope, Claims, DenyReason, OwnerConstraint, ResourceMeta};
|
use super::types::{AccessScope, Claims, DenyReason, OwnerConstraint, ResourceMeta, Visibility};
|
||||||
|
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
// Trait
|
// Trait
|
||||||
@@ -247,7 +247,7 @@ impl Default for CompositeScopeChecker {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use crate::rbac::types::{ResourceType, Visibility};
|
use crate::rbac::types::ResourceType;
|
||||||
|
|
||||||
fn test_claims() -> Claims {
|
fn test_claims() -> Claims {
|
||||||
Claims::new("alice")
|
Claims::new("alice")
|
||||||
|
|||||||
@@ -7,6 +7,7 @@
|
|||||||
/// - ResourceMeta: metadata attached to each document/wiki entry
|
/// - ResourceMeta: metadata attached to each document/wiki entry
|
||||||
|
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
// Verbs
|
// Verbs
|
||||||
|
|||||||
@@ -4,8 +4,9 @@
|
|||||||
//! Uses LLM (Qwen-7B or similar) to judge if retrieved results are relevant.
|
//! Uses LLM (Qwen-7B or similar) to judge if retrieved results are relevant.
|
||||||
//! Tracks precision, recall, F1 via Prometheus metrics.
|
//! Tracks precision, recall, F1 via Prometheus metrics.
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use tracing::debug;
|
use tracing::{debug, error};
|
||||||
|
|
||||||
use crate::metrics;
|
use crate::metrics;
|
||||||
|
|
||||||
|
|||||||
@@ -1,3 +1,13 @@
|
|||||||
|
/// Result Compressor: Optimize response size without losing essential information
|
||||||
|
///
|
||||||
|
/// Strategies:
|
||||||
|
/// - Truncate long texts to summary
|
||||||
|
/// - Extract key sentences
|
||||||
|
/// - Remove redundant metadata
|
||||||
|
/// - Compress to multiple formats (JSON, msgpack, CBOR)
|
||||||
|
/// - Progressive disclosure (compact by default, expand on demand)
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
/// Compression strategy
|
/// Compression strategy
|
||||||
|
|||||||
@@ -0,0 +1,137 @@
|
|||||||
|
//! M8.6 — Simple Hybrid Search (Semantic + Lexical Fusion)
|
||||||
|
//!
|
||||||
|
//! Combines pgvector semantic search with OpenSearch lexical search using RRF.
|
||||||
|
//! Simpler than HybridQueryWorker - uses only existing VectorStore/OpenSearchClient APIs.
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
|
use mem_store::VectorStore;
|
||||||
|
use pgvector::Vector;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use crate::opensearch_client::OpenSearchClient;
|
||||||
|
use crate::query_optimizer::{RRFFusion, RRFConfig};
|
||||||
|
|
||||||
|
/// Hybrid search result with score breakdown
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct SimpleHybridResult {
|
||||||
|
pub id: String,
|
||||||
|
pub content: String,
|
||||||
|
pub project: String,
|
||||||
|
pub semantic_score: Option<f32>,
|
||||||
|
pub lexical_score: Option<f32>,
|
||||||
|
pub final_score: f32,
|
||||||
|
pub rank: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Simple hybrid search orchestrator
|
||||||
|
pub struct SimpleHybridSearch {
|
||||||
|
vector_store: Arc<VectorStore>,
|
||||||
|
opensearch: Option<Arc<OpenSearchClient>>,
|
||||||
|
rrf: RRFFusion,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SimpleHybridSearch {
|
||||||
|
pub fn new(
|
||||||
|
vector_store: Arc<VectorStore>,
|
||||||
|
opensearch: Option<Arc<OpenSearchClient>>,
|
||||||
|
) -> Self {
|
||||||
|
// Create RRF with default config (k=60 per academic standards)
|
||||||
|
let rrf_config = RRFConfig {
|
||||||
|
k: 60.0,
|
||||||
|
retrieve_k: 50,
|
||||||
|
final_k: 10,
|
||||||
|
};
|
||||||
|
let rrf = RRFFusion::new(rrf_config);
|
||||||
|
|
||||||
|
Self {
|
||||||
|
vector_store,
|
||||||
|
opensearch,
|
||||||
|
rrf,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute hybrid search: semantic + lexical with RRF fusion
|
||||||
|
pub async fn search(
|
||||||
|
&self,
|
||||||
|
project: &str,
|
||||||
|
query: &str,
|
||||||
|
embedding: &Vector,
|
||||||
|
jwt_token: &str,
|
||||||
|
limit: usize,
|
||||||
|
) -> Result<Vec<SimpleHybridResult>> {
|
||||||
|
// 1. Semantic search (pgvector)
|
||||||
|
let semantic_results = self
|
||||||
|
.vector_store
|
||||||
|
.search_l1(project, embedding, limit as i64)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let semantic_scores: Vec<(String, f32)> = semantic_results
|
||||||
|
.into_iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, result)| {
|
||||||
|
// Rank to score conversion
|
||||||
|
let rank_score = 1.0 / (i as f32 + 1.0);
|
||||||
|
(result.item.id.to_string(), rank_score)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// 2. Lexical search (OpenSearch) - optional if available
|
||||||
|
// TODO: Implement OpenSearchClient.search() method
|
||||||
|
let lexical_scores: Vec<(String, f32)> = vec![];
|
||||||
|
|
||||||
|
// 3. Fuse with RRF
|
||||||
|
let fused = self.rrf.fuse(semantic_scores.clone(), lexical_scores.clone());
|
||||||
|
|
||||||
|
// 4. Convert to response format
|
||||||
|
let results = fused
|
||||||
|
.into_iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(rank, (id, score))| {
|
||||||
|
let semantic_score = semantic_scores
|
||||||
|
.iter()
|
||||||
|
.find(|(sid, _)| sid == &id)
|
||||||
|
.map(|(_, s)| *s);
|
||||||
|
|
||||||
|
let lexical_score = lexical_scores
|
||||||
|
.iter()
|
||||||
|
.find(|(sid, _)| sid == &id)
|
||||||
|
.map(|(_, s)| *s);
|
||||||
|
|
||||||
|
SimpleHybridResult {
|
||||||
|
id: id.clone(),
|
||||||
|
content: String::new(), // Would fetch from store
|
||||||
|
project: project.to_string(),
|
||||||
|
semantic_score,
|
||||||
|
lexical_score,
|
||||||
|
final_score: score,
|
||||||
|
rank: rank + 1,
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(results)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_simple_hybrid_result_creation() {
|
||||||
|
let result = SimpleHybridResult {
|
||||||
|
id: "doc1".to_string(),
|
||||||
|
content: "test".to_string(),
|
||||||
|
project: "test".to_string(),
|
||||||
|
semantic_score: Some(0.95),
|
||||||
|
lexical_score: Some(8.5),
|
||||||
|
final_score: 0.067,
|
||||||
|
rank: 1,
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(result.id, "doc1");
|
||||||
|
assert_eq!(result.rank, 1);
|
||||||
|
assert!(result.semantic_score.is_some());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,10 +1,10 @@
|
|||||||
use anyhow::Result;
|
use anyhow::{anyhow, Result};
|
||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use serde::{Serialize, Deserialize};
|
use serde::{Serialize, Deserialize};
|
||||||
|
|
||||||
use mem_store::PgRepo;
|
use mem_store::{PgRepo, Level};
|
||||||
|
|
||||||
/// Memory record from log
|
/// Memory record from log
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
@@ -39,7 +39,7 @@ pub struct VerifyOpts {
|
|||||||
pub check_db: bool,
|
pub check_db: bool,
|
||||||
pub check_log: bool,
|
pub check_log: bool,
|
||||||
pub log_dir: Option<PathBuf>,
|
pub log_dir: Option<PathBuf>,
|
||||||
pub _format: OutputFormat,
|
pub format: OutputFormat,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy)]
|
#[derive(Debug, Clone, Copy)]
|
||||||
@@ -69,13 +69,13 @@ pub struct VerificationResult {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub struct Verifier {
|
pub struct Verifier {
|
||||||
_repo: PgRepo,
|
repo: PgRepo,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Verifier {
|
impl Verifier {
|
||||||
pub async fn new(db_url: &str) -> Result<Self> {
|
pub async fn new(db_url: &str) -> Result<Self> {
|
||||||
let repo = PgRepo::connect(db_url).await?;
|
let repo = PgRepo::connect(db_url).await?;
|
||||||
Ok(Self { _repo: repo })
|
Ok(Self { repo })
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Run all verifications
|
/// Run all verifications
|
||||||
@@ -136,7 +136,7 @@ impl Verifier {
|
|||||||
let mut evidence_gate_count = 0;
|
let mut evidence_gate_count = 0;
|
||||||
let mut evidence_records = 0;
|
let mut evidence_records = 0;
|
||||||
|
|
||||||
for (_line_num, memory) in memories.iter().enumerate() {
|
for (line_num, memory) in memories.iter().enumerate() {
|
||||||
let sha = Self::memory_sha(&memory.text);
|
let sha = Self::memory_sha(&memory.text);
|
||||||
memory_map.insert(sha.clone(), memory);
|
memory_map.insert(sha.clone(), memory);
|
||||||
level_map.insert(sha.clone(), memory.level.clone());
|
level_map.insert(sha.clone(), memory.level.clone());
|
||||||
@@ -188,7 +188,7 @@ impl Verifier {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Invariant 2: Every parent sha resolves to a memory that exists
|
// Invariant 2: Every parent sha resolves to a memory that exists
|
||||||
for (_sha, parents) in &memory_parents {
|
for (sha, parents) in &memory_parents {
|
||||||
for parent_sha in parents {
|
for parent_sha in parents {
|
||||||
if !memory_map.contains_key(parent_sha) {
|
if !memory_map.contains_key(parent_sha) {
|
||||||
violations.push(Violation {
|
violations.push(Violation {
|
||||||
@@ -206,7 +206,7 @@ impl Verifier {
|
|||||||
// Invariant 3: Every evidence sha appears as a parent of at least one memory
|
// Invariant 3: Every evidence sha appears as a parent of at least one memory
|
||||||
for evidence_sha in &evidence_shas {
|
for evidence_sha in &evidence_shas {
|
||||||
let mut is_cited = false;
|
let mut is_cited = false;
|
||||||
for (_sha, parents) in &memory_parents {
|
for (sha, parents) in &memory_parents {
|
||||||
if parents.contains(evidence_sha) {
|
if parents.contains(evidence_sha) {
|
||||||
is_cited = true;
|
is_cited = true;
|
||||||
break;
|
break;
|
||||||
|
|||||||
@@ -44,7 +44,7 @@ impl ExtractedEntity {
|
|||||||
#[async_trait]
|
#[async_trait]
|
||||||
pub trait EntityExtractor: Send + Sync {
|
pub trait EntityExtractor: Send + Sync {
|
||||||
async fn extract(&self, text: &str) -> Result<Vec<ExtractedEntity>>;
|
async fn extract(&self, text: &str) -> Result<Vec<ExtractedEntity>>;
|
||||||
async fn extract_with_auth(&self, text: &str, _x_forward_user: Option<&str>) -> Result<Vec<ExtractedEntity>> {
|
async fn extract_with_auth(&self, text: &str, x_forward_user: Option<&str>) -> Result<Vec<ExtractedEntity>> {
|
||||||
// Default: ignore auth header, use regular extract
|
// Default: ignore auth header, use regular extract
|
||||||
self.extract(text).await
|
self.extract(text).await
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -243,7 +243,7 @@ pub struct FilterStatistics {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use mem_core::entity::{Entity, EntityType};
|
use mem_core::entity::Entity;
|
||||||
|
|
||||||
fn create_test_entity(name: &str) -> Entity {
|
fn create_test_entity(name: &str) -> Entity {
|
||||||
Entity::new("poimen", name, EntityType::Person)
|
Entity::new("poimen", name, EntityType::Person)
|
||||||
|
|||||||
@@ -64,15 +64,9 @@ impl EmbeddingsClient {
|
|||||||
/// - `LLM_API_BASE`: Gateway endpoint (default: https://api.riotpiao.com)
|
/// - `LLM_API_BASE`: Gateway endpoint (default: https://api.riotpiao.com)
|
||||||
/// - `LLM_API_KEY`: API key (optional)
|
/// - `LLM_API_KEY`: API key (optional)
|
||||||
pub fn from_env() -> Result<Self> {
|
pub fn from_env() -> Result<Self> {
|
||||||
let mut base_url = env::var("LLM_API_BASE")
|
let base_url = env::var("LLM_API_BASE")
|
||||||
.unwrap_or_else(|_| "https://api.riotpiao.com".to_string());
|
.unwrap_or_else(|_| "https://api.riotpiao.com".to_string());
|
||||||
|
|
||||||
// Strip trailing /v1 to avoid double /v1/v1/embeddings
|
|
||||||
base_url = base_url.trim_end_matches('/').to_string();
|
|
||||||
if base_url.ends_with("/v1") {
|
|
||||||
base_url = base_url[..base_url.len() - 3].to_string();
|
|
||||||
}
|
|
||||||
|
|
||||||
let model = env::var("EMBEDDINGS_MODEL")
|
let model = env::var("EMBEDDINGS_MODEL")
|
||||||
.unwrap_or_else(|_| "nomic-ai/nomic-embed-text-v2-moe".to_string());
|
.unwrap_or_else(|_| "nomic-ai/nomic-embed-text-v2-moe".to_string());
|
||||||
|
|
||||||
@@ -167,9 +161,10 @@ impl EmbeddingsClient {
|
|||||||
let url = format!("{}/v1/embeddings", self.base_url);
|
let url = format!("{}/v1/embeddings", self.base_url);
|
||||||
let mut builder = self.http.post(&url);
|
let mut builder = self.http.post(&url);
|
||||||
|
|
||||||
// Send as Bearer token (gateway expects Authorization: Bearer <key>)
|
// Send apikey header even though route currently doesn't require auth
|
||||||
|
// This future-proofs for when the route's auth plugin gets enabled
|
||||||
if !self.api_key.is_empty() {
|
if !self.api_key.is_empty() {
|
||||||
builder = builder.header("Authorization", format!("Bearer {}", &self.api_key));
|
builder = builder.header("apikey", &self.api_key);
|
||||||
}
|
}
|
||||||
|
|
||||||
let resp = builder.json(&req).send().await?;
|
let resp = builder.json(&req).send().await?;
|
||||||
@@ -220,28 +215,6 @@ mod tests {
|
|||||||
assert_eq!(EMBEDDINGS_DIM, 768);
|
assert_eq!(EMBEDDINGS_DIM, 768);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_strip_trailing_v1() {
|
|
||||||
// Simulates LLM_API_BASE=https://api.riotpiao.com/v1
|
|
||||||
let mut base = "https://api.riotpiao.com/v1".to_string();
|
|
||||||
base = base.trim_end_matches('/').to_string();
|
|
||||||
if base.ends_with("/v1") {
|
|
||||||
base = base[..base.len() - 3].to_string();
|
|
||||||
}
|
|
||||||
assert_eq!(base, "https://api.riotpiao.com");
|
|
||||||
assert_eq!(format!("{}/v1/embeddings", base), "https://api.riotpiao.com/v1/embeddings");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_no_strip_when_no_v1() {
|
|
||||||
let mut base = "https://api.riotpiao.com".to_string();
|
|
||||||
base = base.trim_end_matches('/').to_string();
|
|
||||||
if base.ends_with("/v1") {
|
|
||||||
base = base[..base.len() - 3].to_string();
|
|
||||||
}
|
|
||||||
assert_eq!(base, "https://api.riotpiao.com");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_parse_real_embedding_response() {
|
fn test_parse_real_embedding_response() {
|
||||||
// Exact format returned by embeddings-predictor service
|
// Exact format returned by embeddings-predictor service
|
||||||
|
|||||||
Generated
+52
@@ -0,0 +1,52 @@
|
|||||||
|
{
|
||||||
|
"db_name": "PostgreSQL",
|
||||||
|
"query": "\n SELECT \n version_num,\n operation,\n snapshot,\n changed_at,\n changed_by,\n COALESCE(fields_changed, '{}') as \"fields_changed!\"\n FROM memory_entity_version\n WHERE entity_id = $1\n ORDER BY version_num DESC\n ",
|
||||||
|
"describe": {
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"ordinal": 0,
|
||||||
|
"name": "version_num",
|
||||||
|
"type_info": "Int4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 1,
|
||||||
|
"name": "operation",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 2,
|
||||||
|
"name": "snapshot",
|
||||||
|
"type_info": "Jsonb"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 3,
|
||||||
|
"name": "changed_at",
|
||||||
|
"type_info": "Timestamptz"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 4,
|
||||||
|
"name": "changed_by",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 5,
|
||||||
|
"name": "fields_changed!",
|
||||||
|
"type_info": "TextArray"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"parameters": {
|
||||||
|
"Left": [
|
||||||
|
"Text"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"nullable": [
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
null
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"hash": "1e81bb729531ca33e4cef21623bcfe4fafb0c1bd435353b205f582bfda8873bc"
|
||||||
|
}
|
||||||
Generated
+52
@@ -0,0 +1,52 @@
|
|||||||
|
{
|
||||||
|
"db_name": "PostgreSQL",
|
||||||
|
"query": "\n SELECT \n version_num,\n operation,\n snapshot,\n changed_at,\n changed_by,\n COALESCE(fields_changed, '{}') as \"fields_changed!\"\n FROM memory_edge_version\n WHERE edge_id = $1\n ORDER BY version_num DESC\n ",
|
||||||
|
"describe": {
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"ordinal": 0,
|
||||||
|
"name": "version_num",
|
||||||
|
"type_info": "Int4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 1,
|
||||||
|
"name": "operation",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 2,
|
||||||
|
"name": "snapshot",
|
||||||
|
"type_info": "Jsonb"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 3,
|
||||||
|
"name": "changed_at",
|
||||||
|
"type_info": "Timestamptz"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 4,
|
||||||
|
"name": "changed_by",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 5,
|
||||||
|
"name": "fields_changed!",
|
||||||
|
"type_info": "TextArray"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"parameters": {
|
||||||
|
"Left": [
|
||||||
|
"Uuid"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"nullable": [
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
null
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"hash": "62d65d4afc4d292b37de8e5cb59fbd51c602bdc1b437988f54e6c7fe268b9816"
|
||||||
|
}
|
||||||
Generated
+53
@@ -0,0 +1,53 @@
|
|||||||
|
{
|
||||||
|
"db_name": "PostgreSQL",
|
||||||
|
"query": "\n SELECT \n version_num,\n operation,\n snapshot,\n changed_at,\n changed_by,\n COALESCE(fields_changed, '{}') as \"fields_changed!\"\n FROM memory_entity_version\n WHERE entity_id = $1 AND changed_at <= $2\n ORDER BY version_num DESC\n LIMIT 1\n ",
|
||||||
|
"describe": {
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"ordinal": 0,
|
||||||
|
"name": "version_num",
|
||||||
|
"type_info": "Int4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 1,
|
||||||
|
"name": "operation",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 2,
|
||||||
|
"name": "snapshot",
|
||||||
|
"type_info": "Jsonb"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 3,
|
||||||
|
"name": "changed_at",
|
||||||
|
"type_info": "Timestamptz"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 4,
|
||||||
|
"name": "changed_by",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 5,
|
||||||
|
"name": "fields_changed!",
|
||||||
|
"type_info": "TextArray"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"parameters": {
|
||||||
|
"Left": [
|
||||||
|
"Text",
|
||||||
|
"Timestamptz"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"nullable": [
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
null
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"hash": "aee5900f5e3d7cbba23729bbf2dd033dcc4cb41f6c851bf447a9238810684d18"
|
||||||
|
}
|
||||||
Generated
+53
@@ -0,0 +1,53 @@
|
|||||||
|
{
|
||||||
|
"db_name": "PostgreSQL",
|
||||||
|
"query": "\n SELECT \n version_num,\n operation,\n snapshot,\n changed_at,\n changed_by,\n COALESCE(fields_changed, '{}') as \"fields_changed!\"\n FROM memory_entity_version\n WHERE entity_id = $1 AND version_num = $2\n ",
|
||||||
|
"describe": {
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"ordinal": 0,
|
||||||
|
"name": "version_num",
|
||||||
|
"type_info": "Int4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 1,
|
||||||
|
"name": "operation",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 2,
|
||||||
|
"name": "snapshot",
|
||||||
|
"type_info": "Jsonb"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 3,
|
||||||
|
"name": "changed_at",
|
||||||
|
"type_info": "Timestamptz"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 4,
|
||||||
|
"name": "changed_by",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 5,
|
||||||
|
"name": "fields_changed!",
|
||||||
|
"type_info": "TextArray"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"parameters": {
|
||||||
|
"Left": [
|
||||||
|
"Text",
|
||||||
|
"Int4"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"nullable": [
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
null
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"hash": "c045466e1fe037dbdafea1008f262f4e48f104ea77732aa1d32ecb797f70e71d"
|
||||||
|
}
|
||||||
Generated
+53
@@ -0,0 +1,53 @@
|
|||||||
|
{
|
||||||
|
"db_name": "PostgreSQL",
|
||||||
|
"query": "\n SELECT \n version_num,\n operation,\n snapshot,\n changed_at,\n changed_by,\n COALESCE(fields_changed, '{}') as \"fields_changed!\"\n FROM memory_edge_version\n WHERE edge_id = $1 AND version_num = $2\n ",
|
||||||
|
"describe": {
|
||||||
|
"columns": [
|
||||||
|
{
|
||||||
|
"ordinal": 0,
|
||||||
|
"name": "version_num",
|
||||||
|
"type_info": "Int4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 1,
|
||||||
|
"name": "operation",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 2,
|
||||||
|
"name": "snapshot",
|
||||||
|
"type_info": "Jsonb"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 3,
|
||||||
|
"name": "changed_at",
|
||||||
|
"type_info": "Timestamptz"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 4,
|
||||||
|
"name": "changed_by",
|
||||||
|
"type_info": "Varchar"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"ordinal": 5,
|
||||||
|
"name": "fields_changed!",
|
||||||
|
"type_info": "TextArray"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"parameters": {
|
||||||
|
"Left": [
|
||||||
|
"Uuid",
|
||||||
|
"Int4"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"nullable": [
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
null
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"hash": "ca6872495bc04c6a65531279af8c758637c902dda2cc10366662988c6973ca48"
|
||||||
|
}
|
||||||
@@ -1,20 +1,20 @@
|
|||||||
-- Migration 009: Temporal knowledge graph edge schema (Zep paper §2.2.2)
|
-- Migration 009: Temporal edge schema (Zep paper §2.2.2)
|
||||||
-- Replaces old memory_edge (child_sha/parent_sha provenance DAG)
|
-- Replaces old memory_edge (child_sha/parent_sha node graph)
|
||||||
-- with temporal edge schema for the knowledge graph.
|
-- with temporal edge schema supporting relation types, facts, and validity periods.
|
||||||
-- Idempotent: safe to run multiple times.
|
-- Idempotent: safe to run multiple times.
|
||||||
|
|
||||||
-- Rename old provenance DAG table if it still has child_sha columns
|
-- Rename old table if it still exists (skip if already migrated)
|
||||||
DO $$
|
DO $$
|
||||||
BEGIN
|
BEGIN
|
||||||
IF EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'memory_edge'
|
IF EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'memory_edge'
|
||||||
AND EXISTS (SELECT 1 FROM information_schema.columns
|
AND EXISTS (SELECT 1 FROM information_schema.columns
|
||||||
WHERE table_name = 'memory_edge' AND column_name = 'child_sha'))
|
WHERE table_name = 'memory_edge' AND column_name = 'child_sha'))
|
||||||
THEN
|
THEN
|
||||||
ALTER TABLE memory_edge RENAME TO memory_edge_provenance;
|
ALTER TABLE memory_edge RENAME TO memory_edge_legacy;
|
||||||
END IF;
|
END IF;
|
||||||
END $$;
|
END $$;
|
||||||
|
|
||||||
-- Create temporal knowledge graph edge table
|
-- Create temporal edge table
|
||||||
CREATE TABLE IF NOT EXISTS memory_edge (
|
CREATE TABLE IF NOT EXISTS memory_edge (
|
||||||
id TEXT PRIMARY KEY,
|
id TEXT PRIMARY KEY,
|
||||||
project_id TEXT NOT NULL DEFAULT 'default',
|
project_id TEXT NOT NULL DEFAULT 'default',
|
||||||
@@ -64,5 +64,4 @@ CREATE UNIQUE INDEX IF NOT EXISTS idx_memory_entity_project_name ON memory_entit
|
|||||||
|
|
||||||
-- ROLLBACK instructions:
|
-- ROLLBACK instructions:
|
||||||
-- DROP TABLE IF EXISTS memory_edge;
|
-- DROP TABLE IF EXISTS memory_edge;
|
||||||
-- ALTER TABLE IF EXISTS memory_edge_provenance RENAME TO memory_edge;
|
-- ALTER TABLE IF EXISTS memory_edge_legacy RENAME TO memory_edge;
|
||||||
-- DROP INDEX IF EXISTS idx_memory_entity_project_name;
|
|
||||||
|
|||||||
@@ -177,7 +177,6 @@ pub struct AuditedEntityRepo {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use serde_json::json;
|
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_diff_fields_modified() {
|
fn test_diff_fields_modified() {
|
||||||
|
|||||||
@@ -0,0 +1,543 @@
|
|||||||
|
/// PostgreSQL repository implementation for Phase 2.6 DB Integration.
|
||||||
|
///
|
||||||
|
/// Connects ingest pipeline to persistent storage.
|
||||||
|
/// Handles transactions, error recovery, and audit logging.
|
||||||
|
|
||||||
|
use sqlx::{Pool, Postgres, Row, Transaction, Error as SqlxError};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use mem_core::entity::Entity;
|
||||||
|
use mem_core::edge::Edge;
|
||||||
|
use crate::entity_repo::EntityRepoOps;
|
||||||
|
use crate::edge_repo::EdgeRepoOps;
|
||||||
|
|
||||||
|
/// Database connection error types
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub enum DbError {
|
||||||
|
ConnectionFailed(String),
|
||||||
|
QueryFailed(String),
|
||||||
|
TransactionFailed(String),
|
||||||
|
DuplicateKey(String),
|
||||||
|
NotFound(String),
|
||||||
|
InvalidData(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Display for DbError {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
||||||
|
match self {
|
||||||
|
DbError::ConnectionFailed(msg) => write!(f, "Connection failed: {}", msg),
|
||||||
|
DbError::QueryFailed(msg) => write!(f, "Query failed: {}", msg),
|
||||||
|
DbError::TransactionFailed(msg) => write!(f, "Transaction failed: {}", msg),
|
||||||
|
DbError::DuplicateKey(msg) => write!(f, "Duplicate key: {}", msg),
|
||||||
|
DbError::NotFound(msg) => write!(f, "Not found: {}", msg),
|
||||||
|
DbError::InvalidData(msg) => write!(f, "Invalid data: {}", msg),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::error::Error for DbError {}
|
||||||
|
|
||||||
|
/// PostgreSQL repository pool
|
||||||
|
pub struct DbPool {
|
||||||
|
pool: Pool<Postgres>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DbPool {
|
||||||
|
/// Create new DB pool from connection string
|
||||||
|
pub async fn new(database_url: &str) -> Result<Self, DbError> {
|
||||||
|
let pool = Pool::<Postgres>::connect(database_url)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::ConnectionFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(DbPool { pool })
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get pool for queries
|
||||||
|
pub fn pool(&self) -> &Pool<Postgres> {
|
||||||
|
&self.pool
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Test connection
|
||||||
|
pub async fn health_check(&self) -> Result<(), DbError> {
|
||||||
|
sqlx::query("SELECT 1")
|
||||||
|
.fetch_one(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::ConnectionFailed(e.to_string()))?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Persistent entity repository
|
||||||
|
pub struct PersistentEntityRepo {
|
||||||
|
pool: Pool<Postgres>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PersistentEntityRepo {
|
||||||
|
pub fn new(pool: Pool<Postgres>) -> Self {
|
||||||
|
Self { pool }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Save entity to database (idempotent)
|
||||||
|
pub async fn save(&self, entity: &Entity) -> Result<String, DbError> {
|
||||||
|
let query = r#"
|
||||||
|
INSERT INTO memory_entity (id, entity_type, name, description, embedding, created_at, updated_at)
|
||||||
|
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||||
|
ON CONFLICT(id) DO UPDATE SET
|
||||||
|
name = EXCLUDED.name,
|
||||||
|
description = EXCLUDED.description,
|
||||||
|
updated_at = EXCLUDED.updated_at
|
||||||
|
RETURNING id;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
let id = sqlx::query_scalar::<_, String>(query)
|
||||||
|
.bind(&entity.id)
|
||||||
|
.bind(&entity.entity_type)
|
||||||
|
.bind(&entity.name)
|
||||||
|
.bind(&entity.description)
|
||||||
|
.bind(&entity.embedding)
|
||||||
|
.bind(Utc::now())
|
||||||
|
.bind(Utc::now())
|
||||||
|
.fetch_one(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| {
|
||||||
|
if e.to_string().contains("duplicate") {
|
||||||
|
DbError::DuplicateKey(format!("Entity {} already exists", entity.id))
|
||||||
|
} else {
|
||||||
|
DbError::QueryFailed(e.to_string())
|
||||||
|
}
|
||||||
|
})?;
|
||||||
|
|
||||||
|
Ok(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get entity by ID
|
||||||
|
pub async fn get(&self, id: &str) -> Result<Option<Entity>, DbError> {
|
||||||
|
let query = r#"
|
||||||
|
SELECT id, entity_type, name, description, embedding, created_at, updated_at
|
||||||
|
FROM memory_entity
|
||||||
|
WHERE id = $1 AND deleted_at IS NULL;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
let row = sqlx::query(query)
|
||||||
|
.bind(id)
|
||||||
|
.fetch_optional(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(row.map(|r| Entity {
|
||||||
|
id: r.get("id"),
|
||||||
|
entity_type: r.get("entity_type"),
|
||||||
|
name: r.get("name"),
|
||||||
|
description: r.get("description"),
|
||||||
|
embedding: r.get("embedding"),
|
||||||
|
created_at: r.get("created_at"),
|
||||||
|
updated_at: r.get("updated_at"),
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List entities with pagination
|
||||||
|
pub async fn list(&self, limit: i64, offset: i64) -> Result<Vec<Entity>, DbError> {
|
||||||
|
let query = r#"
|
||||||
|
SELECT id, entity_type, name, description, embedding, created_at, updated_at
|
||||||
|
FROM memory_entity
|
||||||
|
WHERE deleted_at IS NULL
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
LIMIT $1 OFFSET $2;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
let rows = sqlx::query(query)
|
||||||
|
.bind(limit)
|
||||||
|
.bind(offset)
|
||||||
|
.fetch_all(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(rows.iter().map(|r| Entity {
|
||||||
|
id: r.get("id"),
|
||||||
|
entity_type: r.get("entity_type"),
|
||||||
|
name: r.get("name"),
|
||||||
|
description: r.get("description"),
|
||||||
|
embedding: r.get("embedding"),
|
||||||
|
created_at: r.get("created_at"),
|
||||||
|
updated_at: r.get("updated_at"),
|
||||||
|
}).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Soft delete entity
|
||||||
|
pub async fn delete(&self, id: &str) -> Result<(), DbError> {
|
||||||
|
let query = r#"
|
||||||
|
UPDATE memory_entity
|
||||||
|
SET deleted_at = $1
|
||||||
|
WHERE id = $2;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
sqlx::query(query)
|
||||||
|
.bind(Utc::now())
|
||||||
|
.bind(id)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Persistent edge repository
|
||||||
|
pub struct PersistentEdgeRepo {
|
||||||
|
pool: Pool<Postgres>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PersistentEdgeRepo {
|
||||||
|
pub fn new(pool: Pool<Postgres>) -> Self {
|
||||||
|
Self { pool }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Save edge to database (idempotent)
|
||||||
|
pub async fn save(&self, edge: &Edge) -> Result<String, DbError> {
|
||||||
|
let query = r#"
|
||||||
|
INSERT INTO memory_edge (id, source_id, target_id, relation_type, fact, strength, t_valid, t_invalid, t_created, t_expired)
|
||||||
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
||||||
|
ON CONFLICT(id) DO UPDATE SET
|
||||||
|
strength = EXCLUDED.strength,
|
||||||
|
t_invalid = EXCLUDED.t_invalid,
|
||||||
|
t_expired = EXCLUDED.t_expired
|
||||||
|
RETURNING id;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
let id = sqlx::query_scalar::<_, String>(query)
|
||||||
|
.bind(&edge.id)
|
||||||
|
.bind(&edge.source_id)
|
||||||
|
.bind(&edge.target_id)
|
||||||
|
.bind(&edge.relation_type)
|
||||||
|
.bind(&edge.fact)
|
||||||
|
.bind(edge.strength)
|
||||||
|
.bind(edge.t_valid)
|
||||||
|
.bind(edge.t_invalid)
|
||||||
|
.bind(edge.t_created)
|
||||||
|
.bind(edge.t_expired)
|
||||||
|
.fetch_one(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| {
|
||||||
|
if e.to_string().contains("duplicate") {
|
||||||
|
DbError::DuplicateKey(format!("Edge {} already exists", edge.id))
|
||||||
|
} else {
|
||||||
|
DbError::QueryFailed(e.to_string())
|
||||||
|
}
|
||||||
|
})?;
|
||||||
|
|
||||||
|
Ok(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get edge by ID
|
||||||
|
pub async fn get(&self, id: &str) -> Result<Option<Edge>, DbError> {
|
||||||
|
let query = r#"
|
||||||
|
SELECT id, source_id, target_id, relation_type, fact, strength, t_valid, t_invalid, t_created, t_expired
|
||||||
|
FROM memory_edge
|
||||||
|
WHERE id = $1 AND t_expired IS NULL;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
let row = sqlx::query(query)
|
||||||
|
.bind(id)
|
||||||
|
.fetch_optional(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(row.map(|r| Edge {
|
||||||
|
id: r.get("id"),
|
||||||
|
source_id: r.get("source_id"),
|
||||||
|
target_id: r.get("target_id"),
|
||||||
|
relation_type: r.get("relation_type"),
|
||||||
|
fact: r.get("fact"),
|
||||||
|
strength: r.get("strength"),
|
||||||
|
t_valid: r.get("t_valid"),
|
||||||
|
t_invalid: r.get("t_invalid"),
|
||||||
|
t_created: r.get("t_created"),
|
||||||
|
t_expired: r.get("t_expired"),
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List edges for a source entity
|
||||||
|
pub async fn list_from(&self, source_id: &str, limit: i64) -> Result<Vec<Edge>, DbError> {
|
||||||
|
let query = r#"
|
||||||
|
SELECT id, source_id, target_id, relation_type, fact, strength, t_valid, t_invalid, t_created, t_expired
|
||||||
|
FROM memory_edge
|
||||||
|
WHERE source_id = $1 AND t_expired IS NULL AND t_invalid IS NULL
|
||||||
|
ORDER BY t_created DESC
|
||||||
|
LIMIT $2;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
let rows = sqlx::query(query)
|
||||||
|
.bind(source_id)
|
||||||
|
.bind(limit)
|
||||||
|
.fetch_all(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(rows.iter().map(|r| Edge {
|
||||||
|
id: r.get("id"),
|
||||||
|
source_id: r.get("source_id"),
|
||||||
|
target_id: r.get("target_id"),
|
||||||
|
relation_type: r.get("relation_type"),
|
||||||
|
fact: r.get("fact"),
|
||||||
|
strength: r.get("strength"),
|
||||||
|
t_valid: r.get("t_valid"),
|
||||||
|
t_invalid: r.get("t_invalid"),
|
||||||
|
t_created: r.get("t_created"),
|
||||||
|
t_expired: r.get("t_expired"),
|
||||||
|
}).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mark edge as contradicted (soft delete)
|
||||||
|
pub async fn invalidate(&self, id: &str) -> Result<(), DbError> {
|
||||||
|
let query = r#"
|
||||||
|
UPDATE memory_edge
|
||||||
|
SET t_invalid = $1
|
||||||
|
WHERE id = $2;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
sqlx::query(query)
|
||||||
|
.bind(Utc::now())
|
||||||
|
.bind(id)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Review queue entry for human verification
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ReviewQueueEntry {
|
||||||
|
pub id: String,
|
||||||
|
pub extraction_type: String, // "entity" | "edge" | "contradiction"
|
||||||
|
pub content: serde_json::Value, // Full extracted data
|
||||||
|
pub status: String, // "pending" | "approved" | "rejected"
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
pub reviewed_at: Option<DateTime<Utc>>,
|
||||||
|
pub reviewed_by: Option<String>, // User ID who reviewed
|
||||||
|
pub rejection_reason: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Review queue repository
|
||||||
|
pub struct ReviewQueueRepo {
|
||||||
|
pool: Pool<Postgres>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ReviewQueueRepo {
|
||||||
|
pub fn new(pool: Pool<Postgres>) -> Self {
|
||||||
|
Self { pool }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Add item to review queue
|
||||||
|
pub async fn enqueue(&self, entry: &ReviewQueueEntry) -> Result<String, DbError> {
|
||||||
|
let query = r#"
|
||||||
|
INSERT INTO review_queue (id, extraction_type, content, status, created_at)
|
||||||
|
VALUES ($1, $2, $3, $4, $5)
|
||||||
|
RETURNING id;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
let id = sqlx::query_scalar::<_, String>(query)
|
||||||
|
.bind(&entry.id)
|
||||||
|
.bind(&entry.extraction_type)
|
||||||
|
.bind(&entry.content)
|
||||||
|
.bind(&entry.status)
|
||||||
|
.bind(Utc::now())
|
||||||
|
.fetch_one(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get pending items for review
|
||||||
|
pub async fn list_pending(&self, limit: i64) -> Result<Vec<ReviewQueueEntry>, DbError> {
|
||||||
|
let query = r#"
|
||||||
|
SELECT id, extraction_type, content, status, created_at, reviewed_at, reviewed_by, rejection_reason
|
||||||
|
FROM review_queue
|
||||||
|
WHERE status = 'pending'
|
||||||
|
ORDER BY created_at ASC
|
||||||
|
LIMIT $1;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
let rows = sqlx::query(query)
|
||||||
|
.bind(limit)
|
||||||
|
.fetch_all(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(rows.iter().map(|r| ReviewQueueEntry {
|
||||||
|
id: r.get("id"),
|
||||||
|
extraction_type: r.get("extraction_type"),
|
||||||
|
content: r.get("content"),
|
||||||
|
status: r.get("status"),
|
||||||
|
created_at: r.get("created_at"),
|
||||||
|
reviewed_at: r.get("reviewed_at"),
|
||||||
|
reviewed_by: r.get("reviewed_by"),
|
||||||
|
rejection_reason: r.get("rejection_reason"),
|
||||||
|
}).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Approve review queue entry
|
||||||
|
pub async fn approve(&self, id: &str, reviewed_by: &str) -> Result<(), DbError> {
|
||||||
|
let query = r#"
|
||||||
|
UPDATE review_queue
|
||||||
|
SET status = 'approved', reviewed_at = $1, reviewed_by = $2
|
||||||
|
WHERE id = $3;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
sqlx::query(query)
|
||||||
|
.bind(Utc::now())
|
||||||
|
.bind(reviewed_by)
|
||||||
|
.bind(id)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Reject review queue entry
|
||||||
|
pub async fn reject(&self, id: &str, reviewed_by: &str, reason: &str) -> Result<(), DbError> {
|
||||||
|
let query = r#"
|
||||||
|
UPDATE review_queue
|
||||||
|
SET status = 'rejected', reviewed_at = $1, reviewed_by = $2, rejection_reason = $3
|
||||||
|
WHERE id = $4;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
sqlx::query(query)
|
||||||
|
.bind(Utc::now())
|
||||||
|
.bind(reviewed_by)
|
||||||
|
.bind(reason)
|
||||||
|
.bind(id)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extraction Audit Repository (Immutable log for audit trail)
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ExtractionAuditEntry {
|
||||||
|
pub id: String,
|
||||||
|
pub extraction_type: String, // "entity" | "edge"
|
||||||
|
pub extraction_id: String, // ID of extracted entity/edge
|
||||||
|
pub source_content: String, // Original text
|
||||||
|
pub extracted_data: serde_json::Value,
|
||||||
|
pub llm_confidence: Option<f32>,
|
||||||
|
pub contradiction_score: Option<f32>,
|
||||||
|
pub status: String, // "extracted" | "approved" | "rejected"
|
||||||
|
pub extracted_at: DateTime<Utc>,
|
||||||
|
pub extracted_by: String, // User or "system"
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ExtractionAuditRepo {
|
||||||
|
pool: Pool<Postgres>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ExtractionAuditRepo {
|
||||||
|
pub fn new(pool: Pool<Postgres>) -> Self {
|
||||||
|
Self { pool }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Log an extraction attempt (immutable append)
|
||||||
|
pub async fn log_extraction(&self, entry: &ExtractionAuditEntry) -> Result<String, DbError> {
|
||||||
|
let query = r#"
|
||||||
|
INSERT INTO extraction_audit (id, extraction_type, extraction_id, source_content, extracted_data, llm_confidence, contradiction_score, status, extracted_at, extracted_by)
|
||||||
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
||||||
|
RETURNING id;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
let id = sqlx::query_scalar::<_, String>(query)
|
||||||
|
.bind(&entry.id)
|
||||||
|
.bind(&entry.extraction_type)
|
||||||
|
.bind(&entry.extraction_id)
|
||||||
|
.bind(&entry.source_content)
|
||||||
|
.bind(&entry.extracted_data)
|
||||||
|
.bind(entry.llm_confidence)
|
||||||
|
.bind(entry.contradiction_score)
|
||||||
|
.bind(&entry.status)
|
||||||
|
.bind(entry.extracted_at)
|
||||||
|
.bind(&entry.extracted_by)
|
||||||
|
.fetch_one(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get audit trail for an extracted item
|
||||||
|
pub async fn get_history(&self, extraction_id: &str) -> Result<Vec<ExtractionAuditEntry>, DbError> {
|
||||||
|
let query = r#"
|
||||||
|
SELECT id, extraction_type, extraction_id, source_content, extracted_data, llm_confidence, contradiction_score, status, extracted_at, extracted_by
|
||||||
|
FROM extraction_audit
|
||||||
|
WHERE extraction_id = $1
|
||||||
|
ORDER BY extracted_at DESC;
|
||||||
|
"#;
|
||||||
|
|
||||||
|
let rows = sqlx::query(query)
|
||||||
|
.bind(extraction_id)
|
||||||
|
.fetch_all(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| DbError::QueryFailed(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(rows.iter().map(|r| ExtractionAuditEntry {
|
||||||
|
id: r.get("id"),
|
||||||
|
extraction_type: r.get("extraction_type"),
|
||||||
|
extraction_id: r.get("extraction_id"),
|
||||||
|
source_content: r.get("source_content"),
|
||||||
|
extracted_data: r.get("extracted_data"),
|
||||||
|
llm_confidence: r.get("llm_confidence"),
|
||||||
|
contradiction_score: r.get("contradiction_score"),
|
||||||
|
status: r.get("status"),
|
||||||
|
extracted_at: r.get("extracted_at"),
|
||||||
|
extracted_by: r.get("extracted_by"),
|
||||||
|
}).collect())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_db_error_display() {
|
||||||
|
let err = DbError::ConnectionFailed("test".to_string());
|
||||||
|
assert!(err.to_string().contains("Connection failed"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_review_queue_entry_creation() {
|
||||||
|
let entry = ReviewQueueEntry {
|
||||||
|
id: "test-1".to_string(),
|
||||||
|
extraction_type: "entity".to_string(),
|
||||||
|
content: serde_json::json!({"name": "test"}),
|
||||||
|
status: "pending".to_string(),
|
||||||
|
created_at: Utc::now(),
|
||||||
|
reviewed_at: None,
|
||||||
|
reviewed_by: None,
|
||||||
|
rejection_reason: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(entry.extraction_type, "entity");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_dead_letter_entry_creation() {
|
||||||
|
let entry = DeadLetterEntry {
|
||||||
|
id: "dlq-1".to_string(),
|
||||||
|
original_content: "test content".to_string(),
|
||||||
|
error_message: "extraction failed".to_string(),
|
||||||
|
error_type: "extraction_failed".to_string(),
|
||||||
|
retry_count: 0,
|
||||||
|
max_retries: 3,
|
||||||
|
created_at: Utc::now(),
|
||||||
|
last_retry_at: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(entry.retry_count, 0);
|
||||||
|
assert!(entry.retry_count < entry.max_retries);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -9,6 +9,7 @@ pub mod community_repo;
|
|||||||
pub mod versioning;
|
pub mod versioning;
|
||||||
pub mod audit_logger;
|
pub mod audit_logger;
|
||||||
pub mod agent_repo;
|
pub mod agent_repo;
|
||||||
|
// pub mod db_repo; // TODO: Fix Entity schema integration
|
||||||
|
|
||||||
pub use event_log::{EventRecord, LogWriter};
|
pub use event_log::{EventRecord, LogWriter};
|
||||||
pub use pgvector::{VectorRecord, VectorStore, ChunkL0, MemoryL1, MemoryL2};
|
pub use pgvector::{VectorRecord, VectorStore, ChunkL0, MemoryL1, MemoryL2};
|
||||||
|
|||||||
@@ -12,7 +12,7 @@ use mem_ingest::OptimizationMetrics;
|
|||||||
|
|
||||||
/// Memory record from log (local copy for rebuild purposes)
|
/// Memory record from log (local copy for rebuild purposes)
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct MemoryRecord {
|
struct MemoryRecord {
|
||||||
pub level: String,
|
pub level: String,
|
||||||
pub project: String,
|
pub project: String,
|
||||||
pub query_id: Option<String>,
|
pub query_id: Option<String>,
|
||||||
@@ -25,7 +25,7 @@ pub struct MemoryRecord {
|
|||||||
|
|
||||||
/// Parent reference for provenance
|
/// Parent reference for provenance
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct MemoryParent {
|
struct MemoryParent {
|
||||||
pub source: String,
|
pub source: String,
|
||||||
pub t: i32,
|
pub t: i32,
|
||||||
pub description: Option<String>,
|
pub description: Option<String>,
|
||||||
|
|||||||
@@ -218,97 +218,6 @@ pub async fn init_schema(pool: &PgPool) -> Result<()> {
|
|||||||
.execute(pool)
|
.execute(pool)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
// Memory entity table (temporal knowledge graph)
|
tracing::info!("Database schema initialized");
|
||||||
sqlx::query(
|
|
||||||
r#"
|
|
||||||
CREATE TABLE IF NOT EXISTS memory_entity (
|
|
||||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
||||||
project_id VARCHAR(255) NOT NULL,
|
|
||||||
name VARCHAR(500) NOT NULL,
|
|
||||||
name_embedding VECTOR(768),
|
|
||||||
summary TEXT,
|
|
||||||
description TEXT,
|
|
||||||
summary_embedding VECTOR(768),
|
|
||||||
entity_type VARCHAR(50),
|
|
||||||
t_created TIMESTAMPTZ DEFAULT NOW(),
|
|
||||||
t_updated TIMESTAMPTZ DEFAULT NOW(),
|
|
||||||
t_expired TIMESTAMPTZ,
|
|
||||||
confidence FLOAT DEFAULT 1.0,
|
|
||||||
source_count INT DEFAULT 1,
|
|
||||||
source_episodes UUID[] DEFAULT '{}',
|
|
||||||
access_count BIGINT DEFAULT 0,
|
|
||||||
last_accessed TIMESTAMPTZ,
|
|
||||||
UNIQUE(project_id, name)
|
|
||||||
)
|
|
||||||
"#,
|
|
||||||
)
|
|
||||||
.execute(pool)
|
|
||||||
.await?;
|
|
||||||
|
|
||||||
sqlx::query("CREATE INDEX IF NOT EXISTS idx_entity_project ON memory_entity(project_id)")
|
|
||||||
.execute(pool)
|
|
||||||
.await?;
|
|
||||||
sqlx::query("CREATE INDEX IF NOT EXISTS idx_entity_type ON memory_entity(project_id, entity_type)")
|
|
||||||
.execute(pool)
|
|
||||||
.await?;
|
|
||||||
|
|
||||||
// Memory edge table (temporal knowledge graph)
|
|
||||||
sqlx::query(
|
|
||||||
r#"
|
|
||||||
CREATE TABLE IF NOT EXISTS memory_edge (
|
|
||||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
||||||
project_id VARCHAR(255) NOT NULL,
|
|
||||||
source_id UUID NOT NULL,
|
|
||||||
target_id UUID NOT NULL,
|
|
||||||
relation_type VARCHAR(100) NOT NULL,
|
|
||||||
fact TEXT NOT NULL,
|
|
||||||
fact_embedding VECTOR(768),
|
|
||||||
t_valid TIMESTAMPTZ,
|
|
||||||
t_invalid TIMESTAMPTZ,
|
|
||||||
t_created TIMESTAMPTZ DEFAULT NOW(),
|
|
||||||
t_expired TIMESTAMPTZ,
|
|
||||||
confidence FLOAT DEFAULT 1.0,
|
|
||||||
contradiction_status VARCHAR(20) DEFAULT 'active',
|
|
||||||
contradiction_confidence FLOAT
|
|
||||||
)
|
|
||||||
"#,
|
|
||||||
)
|
|
||||||
.execute(pool)
|
|
||||||
.await?;
|
|
||||||
|
|
||||||
sqlx::query("CREATE INDEX IF NOT EXISTS idx_edge_project ON memory_edge(project_id)")
|
|
||||||
.execute(pool)
|
|
||||||
.await?;
|
|
||||||
sqlx::query("CREATE INDEX IF NOT EXISTS idx_edge_source ON memory_edge(source_id)")
|
|
||||||
.execute(pool)
|
|
||||||
.await?;
|
|
||||||
sqlx::query("CREATE INDEX IF NOT EXISTS idx_edge_target ON memory_edge(target_id)")
|
|
||||||
.execute(pool)
|
|
||||||
.await?;
|
|
||||||
|
|
||||||
// HNSW vector indexes for semantic search (RAG-001)
|
|
||||||
// name_embedding: primary entity search vector
|
|
||||||
sqlx::query(
|
|
||||||
"CREATE INDEX IF NOT EXISTS idx_entity_name_emb ON memory_entity \
|
|
||||||
USING hnsw (name_embedding vector_cosine_ops) WITH (m = 16, ef_construction = 128)"
|
|
||||||
)
|
|
||||||
.execute(pool)
|
|
||||||
.await?;
|
|
||||||
// summary_embedding: secondary entity search vector
|
|
||||||
sqlx::query(
|
|
||||||
"CREATE INDEX IF NOT EXISTS idx_entity_summary_emb ON memory_entity \
|
|
||||||
USING hnsw (summary_embedding vector_cosine_ops) WITH (m = 16, ef_construction = 128)"
|
|
||||||
)
|
|
||||||
.execute(pool)
|
|
||||||
.await?;
|
|
||||||
// fact_embedding: edge/relationship search vector
|
|
||||||
sqlx::query(
|
|
||||||
"CREATE INDEX IF NOT EXISTS idx_edge_fact_emb ON memory_edge \
|
|
||||||
USING hnsw (fact_embedding vector_cosine_ops) WITH (m = 16, ef_construction = 128)"
|
|
||||||
)
|
|
||||||
.execute(pool)
|
|
||||||
.await?;
|
|
||||||
|
|
||||||
tracing::info!("Database schema initialized (including memory_entity + memory_edge + HNSW indexes)");
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,16 +1,15 @@
|
|||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use sqlx::{PgPool, FromRow};
|
use sqlx::PgPool;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
use chrono::{DateTime, Utc};
|
use chrono::{DateTime, Utc};
|
||||||
|
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct VersionSnapshot {
|
pub struct VersionSnapshot {
|
||||||
pub version_num: i32,
|
pub version_num: i32,
|
||||||
pub operation: String, // 'create' | 'update' | 'delete'
|
pub operation: String, // 'create' | 'update' | 'delete'
|
||||||
pub snapshot: serde_json::Value,
|
pub snapshot: serde_json::Value,
|
||||||
pub changed_at: DateTime<Utc>,
|
pub changed_at: DateTime<Utc>,
|
||||||
pub changed_by: String,
|
pub changed_by: String,
|
||||||
#[sqlx(default)]
|
|
||||||
pub fields_changed: Vec<String>,
|
pub fields_changed: Vec<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -41,7 +40,8 @@ impl EntityVersioningService {
|
|||||||
|
|
||||||
/// Get all versions of an entity in descending order
|
/// Get all versions of an entity in descending order
|
||||||
pub async fn get_versions(&self, entity_id: &str) -> Result<Vec<VersionSnapshot>, sqlx::Error> {
|
pub async fn get_versions(&self, entity_id: &str) -> Result<Vec<VersionSnapshot>, sqlx::Error> {
|
||||||
sqlx::query_as::<_, VersionSnapshot>(
|
sqlx::query_as!(
|
||||||
|
VersionSnapshot,
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
version_num,
|
version_num,
|
||||||
@@ -49,13 +49,13 @@ impl EntityVersioningService {
|
|||||||
snapshot,
|
snapshot,
|
||||||
changed_at,
|
changed_at,
|
||||||
changed_by,
|
changed_by,
|
||||||
COALESCE(fields_changed, '{}') as fields_changed
|
COALESCE(fields_changed, '{}') as "fields_changed!"
|
||||||
FROM memory_entity_version
|
FROM memory_entity_version
|
||||||
WHERE entity_id = $1
|
WHERE entity_id = $1
|
||||||
ORDER BY version_num DESC
|
ORDER BY version_num DESC
|
||||||
"#,
|
"#,
|
||||||
|
entity_id
|
||||||
)
|
)
|
||||||
.bind(entity_id)
|
|
||||||
.fetch_all(&self.pool)
|
.fetch_all(&self.pool)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
@@ -66,7 +66,8 @@ impl EntityVersioningService {
|
|||||||
entity_id: &str,
|
entity_id: &str,
|
||||||
version_num: i32,
|
version_num: i32,
|
||||||
) -> Result<Option<VersionSnapshot>, sqlx::Error> {
|
) -> Result<Option<VersionSnapshot>, sqlx::Error> {
|
||||||
sqlx::query_as::<_, VersionSnapshot>(
|
sqlx::query_as!(
|
||||||
|
VersionSnapshot,
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
version_num,
|
version_num,
|
||||||
@@ -74,13 +75,13 @@ impl EntityVersioningService {
|
|||||||
snapshot,
|
snapshot,
|
||||||
changed_at,
|
changed_at,
|
||||||
changed_by,
|
changed_by,
|
||||||
COALESCE(fields_changed, '{}') as fields_changed
|
COALESCE(fields_changed, '{}') as "fields_changed!"
|
||||||
FROM memory_entity_version
|
FROM memory_entity_version
|
||||||
WHERE entity_id = $1 AND version_num = $2
|
WHERE entity_id = $1 AND version_num = $2
|
||||||
"#,
|
"#,
|
||||||
|
entity_id,
|
||||||
|
version_num
|
||||||
)
|
)
|
||||||
.bind(entity_id)
|
|
||||||
.bind(version_num)
|
|
||||||
.fetch_optional(&self.pool)
|
.fetch_optional(&self.pool)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
@@ -94,7 +95,78 @@ impl EntityVersioningService {
|
|||||||
) -> Result<DiffResult, sqlx::Error> {
|
) -> Result<DiffResult, sqlx::Error> {
|
||||||
let from_snap = self.get_version(entity_id, from_v).await?;
|
let from_snap = self.get_version(entity_id, from_v).await?;
|
||||||
let to_snap = self.get_version(entity_id, to_v).await?;
|
let to_snap = self.get_version(entity_id, to_v).await?;
|
||||||
compute_diff(from_snap, to_snap, from_v, to_v)
|
|
||||||
|
let from_obj = from_snap
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|s| s.snapshot.as_object())
|
||||||
|
.map(|o| o.clone());
|
||||||
|
|
||||||
|
let to_obj = to_snap
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|s| s.snapshot.as_object())
|
||||||
|
.map(|o| o.clone());
|
||||||
|
|
||||||
|
let mut added = Vec::new();
|
||||||
|
let mut removed = Vec::new();
|
||||||
|
let mut modified = Vec::new();
|
||||||
|
|
||||||
|
// Check removed and modified
|
||||||
|
if let Some(ref from) = from_obj {
|
||||||
|
for (key, from_val) in from {
|
||||||
|
if let Some(to) = &to_obj {
|
||||||
|
if let Some(to_val) = to.get(key) {
|
||||||
|
if from_val != to_val {
|
||||||
|
modified.push(DiffField {
|
||||||
|
name: key.clone(),
|
||||||
|
from_value: Some(from_val.clone()),
|
||||||
|
to_value: Some(to_val.clone()),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
removed.push(DiffField {
|
||||||
|
name: key.clone(),
|
||||||
|
from_value: Some(from_val.clone()),
|
||||||
|
to_value: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
removed.push(DiffField {
|
||||||
|
name: key.clone(),
|
||||||
|
from_value: Some(from_val.clone()),
|
||||||
|
to_value: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check added
|
||||||
|
if let Some(to) = to_obj {
|
||||||
|
for (key, to_val) in to {
|
||||||
|
if let Some(from) = &from_obj {
|
||||||
|
if !from.contains_key(&key) {
|
||||||
|
added.push(DiffField {
|
||||||
|
name: key,
|
||||||
|
from_value: None,
|
||||||
|
to_value: Some(to_val),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
added.push(DiffField {
|
||||||
|
name: key,
|
||||||
|
from_value: None,
|
||||||
|
to_value: Some(to_val),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(DiffResult {
|
||||||
|
from_version: from_v,
|
||||||
|
to_version: to_v,
|
||||||
|
added_fields: added,
|
||||||
|
removed_fields: removed,
|
||||||
|
modified_fields: modified,
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Get entity state at a point in time
|
/// Get entity state at a point in time
|
||||||
@@ -103,7 +175,8 @@ impl EntityVersioningService {
|
|||||||
entity_id: &str,
|
entity_id: &str,
|
||||||
as_of: DateTime<Utc>,
|
as_of: DateTime<Utc>,
|
||||||
) -> Result<Option<VersionSnapshot>, sqlx::Error> {
|
) -> Result<Option<VersionSnapshot>, sqlx::Error> {
|
||||||
sqlx::query_as::<_, VersionSnapshot>(
|
sqlx::query_as!(
|
||||||
|
VersionSnapshot,
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
version_num,
|
version_num,
|
||||||
@@ -111,15 +184,15 @@ impl EntityVersioningService {
|
|||||||
snapshot,
|
snapshot,
|
||||||
changed_at,
|
changed_at,
|
||||||
changed_by,
|
changed_by,
|
||||||
COALESCE(fields_changed, '{}') as fields_changed
|
COALESCE(fields_changed, '{}') as "fields_changed!"
|
||||||
FROM memory_entity_version
|
FROM memory_entity_version
|
||||||
WHERE entity_id = $1 AND changed_at <= $2
|
WHERE entity_id = $1 AND changed_at <= $2
|
||||||
ORDER BY version_num DESC
|
ORDER BY version_num DESC
|
||||||
LIMIT 1
|
LIMIT 1
|
||||||
"#,
|
"#,
|
||||||
|
entity_id,
|
||||||
|
as_of
|
||||||
)
|
)
|
||||||
.bind(entity_id)
|
|
||||||
.bind(as_of)
|
|
||||||
.fetch_optional(&self.pool)
|
.fetch_optional(&self.pool)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
@@ -137,7 +210,8 @@ impl EdgeVersioningService {
|
|||||||
|
|
||||||
/// Get all versions of an edge
|
/// Get all versions of an edge
|
||||||
pub async fn get_versions(&self, edge_id: Uuid) -> Result<Vec<VersionSnapshot>, sqlx::Error> {
|
pub async fn get_versions(&self, edge_id: Uuid) -> Result<Vec<VersionSnapshot>, sqlx::Error> {
|
||||||
sqlx::query_as::<_, VersionSnapshot>(
|
sqlx::query_as!(
|
||||||
|
VersionSnapshot,
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
version_num,
|
version_num,
|
||||||
@@ -145,13 +219,13 @@ impl EdgeVersioningService {
|
|||||||
snapshot,
|
snapshot,
|
||||||
changed_at,
|
changed_at,
|
||||||
changed_by,
|
changed_by,
|
||||||
COALESCE(fields_changed, '{}') as fields_changed
|
COALESCE(fields_changed, '{}') as "fields_changed!"
|
||||||
FROM memory_edge_version
|
FROM memory_edge_version
|
||||||
WHERE edge_id = $1
|
WHERE edge_id = $1
|
||||||
ORDER BY version_num DESC
|
ORDER BY version_num DESC
|
||||||
"#,
|
"#,
|
||||||
|
edge_id
|
||||||
)
|
)
|
||||||
.bind(edge_id)
|
|
||||||
.fetch_all(&self.pool)
|
.fetch_all(&self.pool)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
@@ -163,7 +237,8 @@ impl EdgeVersioningService {
|
|||||||
from_v: i32,
|
from_v: i32,
|
||||||
to_v: i32,
|
to_v: i32,
|
||||||
) -> Result<DiffResult, sqlx::Error> {
|
) -> Result<DiffResult, sqlx::Error> {
|
||||||
let from_snap = sqlx::query_as::<_, VersionSnapshot>(
|
let from_snap = sqlx::query_as!(
|
||||||
|
VersionSnapshot,
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
version_num,
|
version_num,
|
||||||
@@ -171,17 +246,18 @@ impl EdgeVersioningService {
|
|||||||
snapshot,
|
snapshot,
|
||||||
changed_at,
|
changed_at,
|
||||||
changed_by,
|
changed_by,
|
||||||
COALESCE(fields_changed, '{}') as fields_changed
|
COALESCE(fields_changed, '{}') as "fields_changed!"
|
||||||
FROM memory_edge_version
|
FROM memory_edge_version
|
||||||
WHERE edge_id = $1 AND version_num = $2
|
WHERE edge_id = $1 AND version_num = $2
|
||||||
"#,
|
"#,
|
||||||
|
edge_id,
|
||||||
|
from_v
|
||||||
)
|
)
|
||||||
.bind(edge_id)
|
|
||||||
.bind(from_v)
|
|
||||||
.fetch_optional(&self.pool)
|
.fetch_optional(&self.pool)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let to_snap = sqlx::query_as::<_, VersionSnapshot>(
|
let to_snap = sqlx::query_as!(
|
||||||
|
VersionSnapshot,
|
||||||
r#"
|
r#"
|
||||||
SELECT
|
SELECT
|
||||||
version_num,
|
version_num,
|
||||||
@@ -189,16 +265,17 @@ impl EdgeVersioningService {
|
|||||||
snapshot,
|
snapshot,
|
||||||
changed_at,
|
changed_at,
|
||||||
changed_by,
|
changed_by,
|
||||||
COALESCE(fields_changed, '{}') as fields_changed
|
COALESCE(fields_changed, '{}') as "fields_changed!"
|
||||||
FROM memory_edge_version
|
FROM memory_edge_version
|
||||||
WHERE edge_id = $1 AND version_num = $2
|
WHERE edge_id = $1 AND version_num = $2
|
||||||
"#,
|
"#,
|
||||||
|
edge_id,
|
||||||
|
to_v
|
||||||
)
|
)
|
||||||
.bind(edge_id)
|
|
||||||
.bind(to_v)
|
|
||||||
.fetch_optional(&self.pool)
|
.fetch_optional(&self.pool)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
|
// Same diff logic as entities
|
||||||
compute_diff(from_snap, to_snap, from_v, to_v)
|
compute_diff(from_snap, to_snap, from_v, to_v)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,21 +0,0 @@
|
|||||||
version: '3.8'
|
|
||||||
|
|
||||||
services:
|
|
||||||
postgres-test:
|
|
||||||
image: postgres:16-alpine
|
|
||||||
environment:
|
|
||||||
POSTGRES_USER: app
|
|
||||||
POSTGRES_PASSWORD: testpass
|
|
||||||
POSTGRES_DB: memory
|
|
||||||
ports:
|
|
||||||
- "5433:5432"
|
|
||||||
volumes:
|
|
||||||
- postgres-test-data:/var/lib/postgresql/data
|
|
||||||
healthcheck:
|
|
||||||
test: ["CMD-SHELL", "pg_isready -U app -d memory"]
|
|
||||||
interval: 2s
|
|
||||||
timeout: 5s
|
|
||||||
retries: 10
|
|
||||||
|
|
||||||
volumes:
|
|
||||||
postgres-test-data:
|
|
||||||
@@ -6,7 +6,6 @@ resources:
|
|||||||
- deployment.yaml
|
- deployment.yaml
|
||||||
- service.yaml
|
- service.yaml
|
||||||
- config.yaml # Production config (SOPS-encrypted)
|
- config.yaml # Production config (SOPS-encrypted)
|
||||||
- tekton-pipeline.yaml
|
|
||||||
|
|
||||||
generators:
|
generators:
|
||||||
- secret-generator.yaml
|
- secret-generator.yaml
|
||||||
|
|||||||
@@ -1,15 +0,0 @@
|
|||||||
apiVersion: monitoring.coreos.com/v1
|
|
||||||
kind: ServiceMonitor
|
|
||||||
metadata:
|
|
||||||
name: poimen-memory
|
|
||||||
namespace: poimen
|
|
||||||
labels:
|
|
||||||
app.kubernetes.io/name: poimen-memory
|
|
||||||
spec:
|
|
||||||
selector:
|
|
||||||
matchLabels:
|
|
||||||
app.kubernetes.io/name: poimen-memory
|
|
||||||
endpoints:
|
|
||||||
- port: http
|
|
||||||
path: /metrics
|
|
||||||
interval: 30s
|
|
||||||
@@ -1,121 +0,0 @@
|
|||||||
---
|
|
||||||
# Tekton Pipeline: poimen-ci
|
|
||||||
# Runs smoke tests against the live service after image push.
|
|
||||||
# Referenced by .gitea/workflows/build.yaml CI.
|
|
||||||
apiVersion: tekton.dev/v1
|
|
||||||
kind: Pipeline
|
|
||||||
metadata:
|
|
||||||
name: poimen-ci
|
|
||||||
namespace: poimen
|
|
||||||
spec:
|
|
||||||
params:
|
|
||||||
- name: image
|
|
||||||
type: string
|
|
||||||
- name: registry-user
|
|
||||||
type: string
|
|
||||||
default: ""
|
|
||||||
- name: registry-token
|
|
||||||
type: string
|
|
||||||
default: ""
|
|
||||||
tasks:
|
|
||||||
- name: integration-tests
|
|
||||||
taskRef:
|
|
||||||
name: poimen-integration-test
|
|
||||||
params:
|
|
||||||
- name: image
|
|
||||||
value: $(params.image)
|
|
||||||
- name: gate-on-tests
|
|
||||||
runAfter:
|
|
||||||
- integration-tests
|
|
||||||
taskSpec:
|
|
||||||
steps:
|
|
||||||
- name: check
|
|
||||||
image: alpine:latest
|
|
||||||
script: |
|
|
||||||
echo "Integration tests passed"
|
|
||||||
---
|
|
||||||
# Tekton Task: smoke test against live poimen-memory service
|
|
||||||
apiVersion: tekton.dev/v1
|
|
||||||
kind: Task
|
|
||||||
metadata:
|
|
||||||
name: poimen-integration-test
|
|
||||||
namespace: poimen
|
|
||||||
spec:
|
|
||||||
params:
|
|
||||||
- name: image
|
|
||||||
type: string
|
|
||||||
results:
|
|
||||||
- name: summary
|
|
||||||
type: string
|
|
||||||
steps:
|
|
||||||
- name: test
|
|
||||||
image: curlimages/curl:latest
|
|
||||||
env:
|
|
||||||
- name: DB_PASSWORD
|
|
||||||
valueFrom:
|
|
||||||
secretKeyRef:
|
|
||||||
key: password
|
|
||||||
name: memory-db-app
|
|
||||||
script: |
|
|
||||||
#!/bin/sh
|
|
||||||
set -e
|
|
||||||
SVC="http://poimen-memory.poimen.svc.cluster.local:8080"
|
|
||||||
PASS=0; FAIL=0
|
|
||||||
|
|
||||||
echo "=== Smoke Test: poimen-memory ==="
|
|
||||||
|
|
||||||
# 1. Health
|
|
||||||
echo "1. GET /health"
|
|
||||||
if curl -sf "$SVC/health" | grep -q '"status":"ok"'; then
|
|
||||||
echo " PASS"; PASS=$((PASS+1))
|
|
||||||
else
|
|
||||||
echo " FAIL"; FAIL=$((FAIL+1))
|
|
||||||
fi
|
|
||||||
|
|
||||||
# 2. Projects
|
|
||||||
echo "2. GET /memory/projects"
|
|
||||||
if curl -sf "$SVC/memory/projects" | grep -q '"projects"'; then
|
|
||||||
echo " PASS"; PASS=$((PASS+1))
|
|
||||||
else
|
|
||||||
echo " FAIL"; FAIL=$((FAIL+1))
|
|
||||||
fi
|
|
||||||
|
|
||||||
# 3. Ingest
|
|
||||||
echo "3. POST /memory/ingest"
|
|
||||||
ID="smoke-$(date +%s)"
|
|
||||||
RESP=$(curl -sf -X POST "$SVC/memory/ingest" \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d "{\"project\":\"ci-smoke\",\"source\":\"tekton\",\"ingest_id\":\"$ID\",\"records\":[{\"role\":\"user\",\"text\":\"[[Kubernetes]] uses [[Docker]]\",\"timestamp\":\"2026-01-01T00:00:00Z\",\"source_position\":0}]}")
|
|
||||||
if echo "$RESP" | grep -q '"status"'; then
|
|
||||||
echo " PASS"; PASS=$((PASS+1))
|
|
||||||
else
|
|
||||||
echo " FAIL"; FAIL=$((FAIL+1))
|
|
||||||
fi
|
|
||||||
|
|
||||||
sleep 3
|
|
||||||
|
|
||||||
# 4. Poll status
|
|
||||||
echo "4. GET /memory/ingest/$ID"
|
|
||||||
if curl -sf "$SVC/memory/ingest/$ID" | grep -q '"status"'; then
|
|
||||||
echo " PASS"; PASS=$((PASS+1))
|
|
||||||
else
|
|
||||||
echo " FAIL"; FAIL=$((FAIL+1))
|
|
||||||
fi
|
|
||||||
|
|
||||||
# 5. Query
|
|
||||||
echo "5. GET /memory/query"
|
|
||||||
if curl -sf "$SVC/memory/query?project=ci-smoke&question=Docker" | grep -q '"project"'; then
|
|
||||||
echo " PASS"; PASS=$((PASS+1))
|
|
||||||
else
|
|
||||||
echo " FAIL"; FAIL=$((FAIL+1))
|
|
||||||
fi
|
|
||||||
|
|
||||||
echo ""
|
|
||||||
echo "Result: $PASS passed, $FAIL failed"
|
|
||||||
|
|
||||||
if [ $FAIL -eq 0 ]; then
|
|
||||||
echo "PASS: $PASS/$((PASS+FAIL))" > /tekton/results/summary
|
|
||||||
else
|
|
||||||
echo "FAIL: $FAIL/$((PASS+FAIL))" > /tekton/results/summary
|
|
||||||
exit 1
|
|
||||||
fi
|
|
||||||
@@ -1,136 +0,0 @@
|
|||||||
---
|
|
||||||
# ClusterRole for CI/Tekton triggers to deploy and manage resources across cluster
|
|
||||||
apiVersion: rbac.authorization.k8s.io/v1
|
|
||||||
kind: ClusterRole
|
|
||||||
metadata:
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
rules:
|
|
||||||
# Tekton resources
|
|
||||||
- apiGroups: ["tekton.dev"]
|
|
||||||
resources: ["pipelineruns", "taskruns", "pipelines", "tasks"]
|
|
||||||
verbs: ["create", "get", "list", "watch", "patch", "update", "delete"]
|
|
||||||
|
|
||||||
# Deployments and pods
|
|
||||||
- apiGroups: ["apps"]
|
|
||||||
resources: ["deployments", "statefulsets", "daemonsets"]
|
|
||||||
verbs: ["get", "list", "watch", "create", "patch", "update"]
|
|
||||||
- apiGroups: [""]
|
|
||||||
resources: ["pods", "pods/log", "pods/status"]
|
|
||||||
verbs: ["get", "list", "watch"]
|
|
||||||
|
|
||||||
# Services and networking
|
|
||||||
- apiGroups: [""]
|
|
||||||
resources: ["services", "endpoints"]
|
|
||||||
verbs: ["get", "list", "watch"]
|
|
||||||
|
|
||||||
# ConfigMaps and Secrets
|
|
||||||
- apiGroups: [""]
|
|
||||||
resources: ["configmaps", "secrets"]
|
|
||||||
verbs: ["get", "list", "watch"]
|
|
||||||
|
|
||||||
# Events
|
|
||||||
- apiGroups: [""]
|
|
||||||
resources: ["events"]
|
|
||||||
verbs: ["create", "patch"]
|
|
||||||
|
|
||||||
# Namespaces
|
|
||||||
- apiGroups: [""]
|
|
||||||
resources: ["namespaces"]
|
|
||||||
verbs: ["get", "list", "watch"]
|
|
||||||
|
|
||||||
# Persistent volumes
|
|
||||||
- apiGroups: [""]
|
|
||||||
resources: ["persistentvolumeclaims", "persistentvolumes"]
|
|
||||||
verbs: ["get", "list", "watch"]
|
|
||||||
|
|
||||||
---
|
|
||||||
# ClusterRoleBinding for ci-tekton-trigger service account
|
|
||||||
apiVersion: rbac.authorization.k8s.io/v1
|
|
||||||
kind: ClusterRoleBinding
|
|
||||||
metadata:
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
roleRef:
|
|
||||||
apiGroup: rbac.authorization.k8s.io
|
|
||||||
kind: ClusterRole
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
subjects:
|
|
||||||
- kind: ServiceAccount
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
namespace: api
|
|
||||||
|
|
||||||
---
|
|
||||||
# Additional ClusterRole for poimen namespace operations
|
|
||||||
apiVersion: rbac.authorization.k8s.io/v1
|
|
||||||
kind: Role
|
|
||||||
metadata:
|
|
||||||
name: ci-tekton-trigger-poimen
|
|
||||||
namespace: poimen
|
|
||||||
rules:
|
|
||||||
# Allow full access in poimen namespace for CI
|
|
||||||
- apiGroups: ["*"]
|
|
||||||
resources: ["*"]
|
|
||||||
verbs: ["*"]
|
|
||||||
|
|
||||||
---
|
|
||||||
# RoleBinding in poimen namespace
|
|
||||||
apiVersion: rbac.authorization.k8s.io/v1
|
|
||||||
kind: RoleBinding
|
|
||||||
metadata:
|
|
||||||
name: ci-tekton-trigger-poimen
|
|
||||||
namespace: poimen
|
|
||||||
roleRef:
|
|
||||||
apiGroup: rbac.authorization.k8s.io
|
|
||||||
kind: Role
|
|
||||||
name: ci-tekton-trigger-poimen
|
|
||||||
subjects:
|
|
||||||
- kind: ServiceAccount
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
namespace: api
|
|
||||||
|
|
||||||
---
|
|
||||||
# RoleBinding in tekton-pipelines namespace
|
|
||||||
apiVersion: rbac.authorization.k8s.io/v1
|
|
||||||
kind: RoleBinding
|
|
||||||
metadata:
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
namespace: tekton-pipelines
|
|
||||||
roleRef:
|
|
||||||
apiGroup: rbac.authorization.k8s.io
|
|
||||||
kind: ClusterRole
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
subjects:
|
|
||||||
- kind: ServiceAccount
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
namespace: api
|
|
||||||
|
|
||||||
---
|
|
||||||
# RoleBinding in llm-serving namespace (for LLM checks)
|
|
||||||
apiVersion: rbac.authorization.k8s.io/v1
|
|
||||||
kind: RoleBinding
|
|
||||||
metadata:
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
namespace: llm-serving
|
|
||||||
roleRef:
|
|
||||||
apiGroup: rbac.authorization.k8s.io
|
|
||||||
kind: ClusterRole
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
subjects:
|
|
||||||
- kind: ServiceAccount
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
namespace: api
|
|
||||||
|
|
||||||
---
|
|
||||||
# RoleBinding in kube-system namespace (for cluster info)
|
|
||||||
apiVersion: rbac.authorization.k8s.io/v1
|
|
||||||
kind: RoleBinding
|
|
||||||
metadata:
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
namespace: kube-system
|
|
||||||
roleRef:
|
|
||||||
apiGroup: rbac.authorization.k8s.io
|
|
||||||
kind: ClusterRole
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
subjects:
|
|
||||||
- kind: ServiceAccount
|
|
||||||
name: ci-tekton-trigger
|
|
||||||
namespace: api
|
|
||||||
@@ -142,3 +142,9 @@ spec:
|
|||||||
fi
|
fi
|
||||||
|
|
||||||
resources:
|
resources:
|
||||||
|
requests:
|
||||||
|
memory: "1Gi"
|
||||||
|
cpu: "500m"
|
||||||
|
limits:
|
||||||
|
memory: "2Gi"
|
||||||
|
cpu: "2000m"
|
||||||
|
|||||||
@@ -98,7 +98,8 @@ CREATE TABLE IF NOT EXISTS role_prompt_mapping (
|
|||||||
active BOOLEAN DEFAULT true,
|
active BOOLEAN DEFAULT true,
|
||||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||||||
UNIQUE(project_id, role_name, prompt_id)
|
UNIQUE(project_id, role_name, prompt_id),
|
||||||
|
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE
|
||||||
);
|
);
|
||||||
|
|
||||||
CREATE INDEX idx_role_prompt_mapping_role ON role_prompt_mapping(project_id, role_name, active);
|
CREATE INDEX idx_role_prompt_mapping_role ON role_prompt_mapping(project_id, role_name, active);
|
||||||
@@ -116,12 +117,10 @@ CREATE TABLE IF NOT EXISTS agent_metrics (
|
|||||||
p95_latency_ms FLOAT DEFAULT 0.0,
|
p95_latency_ms FLOAT DEFAULT 0.0,
|
||||||
p99_latency_ms FLOAT DEFAULT 0.0,
|
p99_latency_ms FLOAT DEFAULT 0.0,
|
||||||
error_rate FLOAT DEFAULT 0.0,
|
error_rate FLOAT DEFAULT 0.0,
|
||||||
recorded_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
|
recorded_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||||||
|
UNIQUE(project_id, agent_id, DATE(recorded_at))
|
||||||
);
|
);
|
||||||
|
|
||||||
-- Note: Use daily rollup job or materialized view for DATE(recorded_at) unique constraint
|
|
||||||
-- PostgreSQL doesn't allow functions in UNIQUE constraints, so we use trigger-based rollup instead
|
|
||||||
|
|
||||||
CREATE INDEX idx_agent_metrics_agent ON agent_metrics(project_id, agent_id, recorded_at DESC);
|
CREATE INDEX idx_agent_metrics_agent ON agent_metrics(project_id, agent_id, recorded_at DESC);
|
||||||
|
|
||||||
-- Prompt Usage Log: detailed invocation tracking
|
-- Prompt Usage Log: detailed invocation tracking
|
||||||
@@ -141,4 +140,3 @@ CREATE TABLE IF NOT EXISTS prompt_usage_log (
|
|||||||
|
|
||||||
CREATE INDEX idx_prompt_usage_log_prompt ON prompt_usage_log(prompt_id, created_at DESC);
|
CREATE INDEX idx_prompt_usage_log_prompt ON prompt_usage_log(prompt_id, created_at DESC);
|
||||||
CREATE INDEX idx_prompt_usage_log_agent ON prompt_usage_log(agent_id, created_at DESC);
|
CREATE INDEX idx_prompt_usage_log_agent ON prompt_usage_log(agent_id, created_at DESC);
|
||||||
CREATE INDEX idx_prompt_usage_log_project ON prompt_usage_log(project_id, created_at DESC);
|
|
||||||
|
|||||||
@@ -1,333 +0,0 @@
|
|||||||
//! Integration test: Full ingest + embedding flow with api-gw
|
|
||||||
//!
|
|
||||||
//! Tests:
|
|
||||||
//! 1. POST /memory/ingest with sample records
|
|
||||||
//! 2. Poll /memory/ingest/{id} until done
|
|
||||||
//! 3. Log root causes of errors
|
|
||||||
//!
|
|
||||||
//! Features:
|
|
||||||
//! - RAII guard for port-forward cleanup (fix for resource leak)
|
|
||||||
//! - Enhanced error handling with resource cleanup
|
|
||||||
//!
|
|
||||||
//! Requires:
|
|
||||||
//! - DATABASE_URL set (postgres)
|
|
||||||
//! - LLM_ENDPOINT set (for embeddings)
|
|
||||||
//! - Server running locally or started by test
|
|
||||||
//!
|
|
||||||
//! Usage:
|
|
||||||
//! ```
|
|
||||||
//! RUST_LOG=debug cargo test --test integration_ingest_with_gw -- --nocapture
|
|
||||||
//! ```
|
|
||||||
|
|
||||||
use std::env;
|
|
||||||
use std::time::Duration;
|
|
||||||
use tokio::time::sleep;
|
|
||||||
use serde_json::json;
|
|
||||||
use std::process::Child;
|
|
||||||
|
|
||||||
/// RAII guard for port-forward cleanup — ensures process is killed even if test panics
|
|
||||||
struct PortForwardGuard {
|
|
||||||
process: Option<Child>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl PortForwardGuard {
|
|
||||||
fn spawn(namespace: &str, service: &str, local_port: u16, remote_port: u16) -> std::io::Result<Self> {
|
|
||||||
let process = std::process::Command::new("kubectl")
|
|
||||||
.args(&["-n", namespace, "port-forward", &format!("svc/{}", service), &format!("{}:{}", local_port, remote_port)])
|
|
||||||
.spawn()?;
|
|
||||||
|
|
||||||
Ok(PortForwardGuard {
|
|
||||||
process: Some(process),
|
|
||||||
})
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Drop for PortForwardGuard {
|
|
||||||
fn drop(&mut self) {
|
|
||||||
if let Some(mut process) = self.process.take() {
|
|
||||||
let _ = process.kill();
|
|
||||||
let _ = process.wait();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
#[ignore] // Run manually: cargo test --test integration_ingest_with_gw -- --ignored --nocapture
|
|
||||||
async fn test_ingest_with_embeddings_and_logging() {
|
|
||||||
// Initialize tracing with DEBUG level to see all logs
|
|
||||||
let _ = tracing_subscriber::fmt()
|
|
||||||
.with_max_level(tracing::Level::DEBUG)
|
|
||||||
.with_writer(std::io::stderr)
|
|
||||||
.try_init();
|
|
||||||
|
|
||||||
let base_url = env::var("MEM_API_URL").unwrap_or_else(|_| "http://localhost:8080".to_string());
|
|
||||||
let api_key = env::var("MEM_API_KEY").unwrap_or_else(|_| "test-key".to_string());
|
|
||||||
let local_port: u16 = 9990;
|
|
||||||
const NAMESPACE: &str = "poimen";
|
|
||||||
const SERVICE: &str = "poimen-memory";
|
|
||||||
|
|
||||||
// Spawn port-forward with RAII guard — guaranteed cleanup
|
|
||||||
let _pf_guard = match PortForwardGuard::spawn(NAMESPACE, SERVICE, local_port, 8080) {
|
|
||||||
Ok(guard) => {
|
|
||||||
sleep(Duration::from_secs(2)).await;
|
|
||||||
println!("[TEST] ✓ Port-forward started");
|
|
||||||
guard
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
eprintln!("[ERROR] Failed to spawn port-forward: {}", e);
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let client = reqwest::Client::new();
|
|
||||||
|
|
||||||
// Sample ingest payload
|
|
||||||
let payload = json!({
|
|
||||||
"project": "test-project",
|
|
||||||
"records": [
|
|
||||||
{
|
|
||||||
"content": "Kubernetes is an open-source container orchestration platform. [[Docker]] [[Go]]",
|
|
||||||
"source": "wiki/kubernetes"
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"content": "Docker is a containerization platform that makes it easier to build, ship, and run applications. [[Linux]] [[Container]]",
|
|
||||||
"source": "wiki/docker"
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"content": "Go is a programming language designed at Google. [[Concurrency]] [[Static Typing]]",
|
|
||||||
"source": "wiki/go"
|
|
||||||
}
|
|
||||||
]
|
|
||||||
});
|
|
||||||
|
|
||||||
println!("[TEST] Sending ingest request...");
|
|
||||||
tracing::info!(
|
|
||||||
target: "integration_test",
|
|
||||||
"Ingest payload: {}",
|
|
||||||
serde_json::to_string_pretty(&payload).unwrap()
|
|
||||||
);
|
|
||||||
|
|
||||||
// POST /memory/ingest
|
|
||||||
let response = match client
|
|
||||||
.post(&format!("{}/memory/ingest", base_url))
|
|
||||||
.header("Authorization", format!("Bearer {}", api_key))
|
|
||||||
.json(&payload)
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(resp) => resp,
|
|
||||||
Err(e) => {
|
|
||||||
eprintln!("[ERROR] Failed to send ingest request: {}", e);
|
|
||||||
tracing::error!(
|
|
||||||
target: "integration_test",
|
|
||||||
error = %e,
|
|
||||||
"Failed to POST /memory/ingest"
|
|
||||||
);
|
|
||||||
panic!("Request failed: {}", e);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let status = response.status();
|
|
||||||
println!("[TEST] Ingest response status: {}", status);
|
|
||||||
|
|
||||||
let body_text = match response.text().await {
|
|
||||||
Ok(text) => text,
|
|
||||||
Err(e) => {
|
|
||||||
tracing::error!(target: "integration_test", error = %e, "Failed to read response body");
|
|
||||||
panic!("Failed to read response body: {}", e);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
println!("[TEST] Response body:\n{}", body_text);
|
|
||||||
|
|
||||||
// Parse response
|
|
||||||
let resp_json: serde_json::Value = match serde_json::from_str(&body_text) {
|
|
||||||
Ok(j) => j,
|
|
||||||
Err(e) => {
|
|
||||||
tracing::error!(
|
|
||||||
target: "integration_test",
|
|
||||||
error = %e,
|
|
||||||
body = %body_text,
|
|
||||||
"Failed to parse JSON response"
|
|
||||||
);
|
|
||||||
panic!("Failed to parse JSON: {}", e);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let ingest_id = match resp_json["id"].as_str() {
|
|
||||||
Some(id) => id.to_string(),
|
|
||||||
None => {
|
|
||||||
tracing::error!(
|
|
||||||
target: "integration_test",
|
|
||||||
response = %serde_json::to_string_pretty(&resp_json).unwrap(),
|
|
||||||
"Missing 'id' in response"
|
|
||||||
);
|
|
||||||
panic!("Missing 'id' in response: {}", resp_json);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
println!("[TEST] Ingest ID: {}", ingest_id);
|
|
||||||
tracing::info!(target: "integration_test", ingest_id = %ingest_id, "Ingest queued");
|
|
||||||
|
|
||||||
// Poll until complete or timeout
|
|
||||||
let max_polls = 60; // 10 minutes with 10s intervals
|
|
||||||
for poll_num in 1..=max_polls {
|
|
||||||
sleep(Duration::from_secs(10)).await;
|
|
||||||
|
|
||||||
println!(
|
|
||||||
"[TEST] Poll #{}/{}: Checking status of ingest {}",
|
|
||||||
poll_num, max_polls, ingest_id
|
|
||||||
);
|
|
||||||
|
|
||||||
let status_response = match client
|
|
||||||
.get(&format!("{}/memory/ingest/{}", base_url, ingest_id))
|
|
||||||
.header("Authorization", format!("Bearer {}", api_key))
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(resp) => resp,
|
|
||||||
Err(e) => {
|
|
||||||
tracing::error!(
|
|
||||||
target: "integration_test",
|
|
||||||
error = %e,
|
|
||||||
ingest_id = %ingest_id,
|
|
||||||
poll = poll_num,
|
|
||||||
"Failed to fetch status"
|
|
||||||
);
|
|
||||||
eprintln!("[ERROR] Failed to fetch status: {}", e);
|
|
||||||
sleep(Duration::from_secs(5)).await;
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let status_text = match status_response.text().await {
|
|
||||||
Ok(text) => text,
|
|
||||||
Err(e) => {
|
|
||||||
tracing::error!(
|
|
||||||
target: "integration_test",
|
|
||||||
error = %e,
|
|
||||||
ingest_id = %ingest_id,
|
|
||||||
"Failed to read status response"
|
|
||||||
);
|
|
||||||
eprintln!("[ERROR] Failed to read status: {}", e);
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let status_json: serde_json::Value = match serde_json::from_str(&status_text) {
|
|
||||||
Ok(j) => j,
|
|
||||||
Err(e) => {
|
|
||||||
tracing::error!(
|
|
||||||
target: "integration_test",
|
|
||||||
error = %e,
|
|
||||||
body = %status_text,
|
|
||||||
"Failed to parse status JSON"
|
|
||||||
);
|
|
||||||
eprintln!("[ERROR] Failed to parse status JSON: {}", e);
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let status = status_json["status"].as_str().unwrap_or("unknown");
|
|
||||||
println!(
|
|
||||||
"[TEST] Poll #{}: status = {}",
|
|
||||||
poll_num, status
|
|
||||||
);
|
|
||||||
|
|
||||||
tracing::info!(
|
|
||||||
target: "integration_test",
|
|
||||||
ingest_id = %ingest_id,
|
|
||||||
poll = poll_num,
|
|
||||||
status = %status,
|
|
||||||
full_response = %serde_json::to_string_pretty(&status_json).unwrap(),
|
|
||||||
"Status check"
|
|
||||||
);
|
|
||||||
|
|
||||||
match status {
|
|
||||||
"done" => {
|
|
||||||
println!("[TEST] ✓ Ingest completed successfully!");
|
|
||||||
tracing::info!(target: "integration_test", "Ingest completed");
|
|
||||||
|
|
||||||
// Extract and log results
|
|
||||||
if let Some(results) = status_json.get("results") {
|
|
||||||
println!("[TEST] Results:\n{}", serde_json::to_string_pretty(results).unwrap());
|
|
||||||
tracing::info!(
|
|
||||||
target: "integration_test",
|
|
||||||
results = %serde_json::to_string_pretty(results).unwrap(),
|
|
||||||
"Ingest results"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
"failed" | "error" => {
|
|
||||||
let error_msg = status_json["error"].as_str().unwrap_or("unknown error");
|
|
||||||
println!("[TEST] ✗ Ingest FAILED: {}", error_msg);
|
|
||||||
tracing::error!(
|
|
||||||
target: "integration_test",
|
|
||||||
ingest_id = %ingest_id,
|
|
||||||
error = %error_msg,
|
|
||||||
full_response = %serde_json::to_string_pretty(&status_json).unwrap(),
|
|
||||||
"Ingest failed"
|
|
||||||
);
|
|
||||||
panic!("Ingest failed: {}", error_msg);
|
|
||||||
}
|
|
||||||
"processing" | "queued" => {
|
|
||||||
// Continue polling
|
|
||||||
println!("[TEST] Still processing, poll again...");
|
|
||||||
}
|
|
||||||
_ => {
|
|
||||||
println!("[TEST] Unknown status: {}", status);
|
|
||||||
tracing::warn!(target: "integration_test", status = %status, "Unknown status");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Timeout — _pf_guard will be dropped here, cleaning up port-forward
|
|
||||||
let msg = format!("Ingest did not complete after {} polls (timeout)", max_polls);
|
|
||||||
println!("[TEST] ✗ {}", msg);
|
|
||||||
tracing::error!(target: "integration_test", ingest_id = %ingest_id, "Ingest timeout");
|
|
||||||
panic!("{}", msg);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
#[ignore]
|
|
||||||
async fn test_ingest_endpoint_only() {
|
|
||||||
let _ = tracing_subscriber::fmt()
|
|
||||||
.with_max_level(tracing::Level::DEBUG)
|
|
||||||
.try_init();
|
|
||||||
|
|
||||||
let base_url = env::var("MEM_API_URL").unwrap_or_else(|_| "http://localhost:8080".to_string());
|
|
||||||
let api_key = env::var("MEM_API_KEY").unwrap_or_else(|_| "test-key".to_string());
|
|
||||||
|
|
||||||
let client = reqwest::Client::new();
|
|
||||||
|
|
||||||
let payload = json!({
|
|
||||||
"project": "test-project",
|
|
||||||
"records": [
|
|
||||||
{
|
|
||||||
"content": "Simple test record",
|
|
||||||
"source": "test"
|
|
||||||
}
|
|
||||||
]
|
|
||||||
});
|
|
||||||
|
|
||||||
println!("[TEST] Testing /memory/ingest endpoint only");
|
|
||||||
|
|
||||||
let response = client
|
|
||||||
.post(&format!("{}/memory/ingest", base_url))
|
|
||||||
.header("Authorization", format!("Bearer {}", api_key))
|
|
||||||
.json(&payload)
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
.expect("Failed to send request");
|
|
||||||
|
|
||||||
println!("[TEST] Status: {}", response.status());
|
|
||||||
|
|
||||||
let body = response.text().await.expect("Failed to read body");
|
|
||||||
println!("[TEST] Response: {}", body);
|
|
||||||
|
|
||||||
let json: serde_json::Value = serde_json::from_str(&body).expect("Invalid JSON");
|
|
||||||
println!("[TEST] Parsed: {}", serde_json::to_string_pretty(&json).unwrap());
|
|
||||||
|
|
||||||
assert!(json.get("id").is_some(), "Response should contain 'id'");
|
|
||||||
}
|
|
||||||
@@ -1,341 +0,0 @@
|
|||||||
//! Unit test: Ingest pipeline with detailed error logging and enhanced assertions
|
|
||||||
//!
|
|
||||||
//! Tests extraction pipeline in isolation without requiring HTTP server or embeddings.
|
|
||||||
//! Useful for debugging extraction errors.
|
|
||||||
//!
|
|
||||||
//! Features:
|
|
||||||
//! - Verify extracted entity names (not just count)
|
|
||||||
//! - Verify edge connections between entities
|
|
||||||
//! - Detailed error logging
|
|
||||||
//!
|
|
||||||
//! Usage:
|
|
||||||
//! ```
|
|
||||||
//! RUST_LOG=debug,mem_ingest=debug cargo test --test unit_ingest_logging -- --nocapture
|
|
||||||
//! ```
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use mem_ingest::ingest_pipeline::{IngestPipeline, Episode};
|
|
||||||
use mem_ingest::entity_extractor::WikiLinkFallbackExtractor;
|
|
||||||
use mem_ingest::fact_extractor::SimpleFactExtractor;
|
|
||||||
use mem_ingest::contradiction_detector::ContradictionHandler;
|
|
||||||
use std::sync::Arc;
|
|
||||||
|
|
||||||
fn init_logging() {
|
|
||||||
let _ = tracing_subscriber::fmt()
|
|
||||||
.with_max_level(tracing::Level::DEBUG)
|
|
||||||
.with_writer(std::io::stderr)
|
|
||||||
.try_init();
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_wiki_link_extraction_with_entity_verification() {
|
|
||||||
init_logging();
|
|
||||||
|
|
||||||
println!("\n[TEST] Wiki link extraction with entity name verification\n");
|
|
||||||
|
|
||||||
let entity_extractor = Arc::new(WikiLinkFallbackExtractor);
|
|
||||||
let fact_extractor = Arc::new(SimpleFactExtractor);
|
|
||||||
let contradiction_detector = Arc::new(ContradictionHandler::default());
|
|
||||||
|
|
||||||
let pipeline = IngestPipeline::new(
|
|
||||||
entity_extractor,
|
|
||||||
fact_extractor,
|
|
||||||
contradiction_detector,
|
|
||||||
);
|
|
||||||
|
|
||||||
let episode = Episode {
|
|
||||||
id: "test-1".to_string(),
|
|
||||||
project_id: "test-project".to_string(),
|
|
||||||
text: "Kubernetes [[Docker]] is a [[Container]] orchestration platform. It works with [[Go]] programs."
|
|
||||||
.to_string(),
|
|
||||||
wiki_links: vec!["Docker".to_string(), "Container".to_string(), "Go".to_string()],
|
|
||||||
};
|
|
||||||
|
|
||||||
tracing::info!(
|
|
||||||
target: "test",
|
|
||||||
episode_id = %episode.id,
|
|
||||||
wiki_links = ?episode.wiki_links,
|
|
||||||
"Starting pipeline ingest"
|
|
||||||
);
|
|
||||||
|
|
||||||
match pipeline.ingest(&episode).await {
|
|
||||||
Ok(result) => {
|
|
||||||
tracing::info!(
|
|
||||||
target: "test",
|
|
||||||
entities = result.entities.len(),
|
|
||||||
edges = result.edges.len(),
|
|
||||||
reviews = result.reviews.len(),
|
|
||||||
"Pipeline succeeded"
|
|
||||||
);
|
|
||||||
|
|
||||||
println!("✓ Extracted {} entities", result.entities.len());
|
|
||||||
for entity in &result.entities {
|
|
||||||
println!(" - {} ({}): {}", entity.name, entity.entity_type.as_str(), entity.summary.as_deref().unwrap_or(""));
|
|
||||||
}
|
|
||||||
|
|
||||||
println!("✓ Extracted {} edges", result.edges.len());
|
|
||||||
for edge in &result.edges {
|
|
||||||
println!(" - {} --[{}]--> {}", edge.source_entity_id, edge.relation_type, edge.target_entity_id);
|
|
||||||
}
|
|
||||||
|
|
||||||
// ENHANCED: Verify extracted entity names (not just count)
|
|
||||||
assert!(!result.entities.is_empty(), "Should extract at least one entity");
|
|
||||||
|
|
||||||
let entity_names: Vec<&str> = result.entities.iter().map(|e| e.name.as_str()).collect();
|
|
||||||
println!("\nEntity names extracted: {:?}", entity_names);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
entity_names.iter().any(|&name| name.contains("Docker") || name.contains("docker")),
|
|
||||||
"Should extract Docker entity"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
entity_names.iter().any(|&name| name.contains("Container") || name.contains("container")),
|
|
||||||
"Should extract Container entity"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
entity_names.iter().any(|&name| name.contains("Go") || name.contains("go")),
|
|
||||||
"Should extract Go entity"
|
|
||||||
);
|
|
||||||
|
|
||||||
// ENHANCED: Verify edges connect correct entity pairs
|
|
||||||
if !result.edges.is_empty() {
|
|
||||||
println!("\nEdge connections:");
|
|
||||||
for edge in &result.edges {
|
|
||||||
println!(" {} → {}", edge.source_entity_id, edge.target_entity_id);
|
|
||||||
|
|
||||||
// Verify both endpoints exist in entities
|
|
||||||
let source_exists = result.entities.iter().any(|e| e.id == edge.source_entity_id);
|
|
||||||
let target_exists = result.entities.iter().any(|e| e.id == edge.target_entity_id);
|
|
||||||
|
|
||||||
assert!(source_exists, "Edge source entity {} must exist in extracted entities", edge.source_entity_id);
|
|
||||||
assert!(target_exists, "Edge target entity {} must exist in extracted entities", edge.target_entity_id);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
tracing::error!(
|
|
||||||
target: "test",
|
|
||||||
error = %e,
|
|
||||||
"Pipeline failed"
|
|
||||||
);
|
|
||||||
panic!("Pipeline failed: {}", e);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_extraction_error_logging() {
|
|
||||||
init_logging();
|
|
||||||
|
|
||||||
println!("\n[TEST] Pipeline error handling with logging\n");
|
|
||||||
|
|
||||||
let entity_extractor = Arc::new(WikiLinkFallbackExtractor);
|
|
||||||
let fact_extractor = Arc::new(SimpleFactExtractor);
|
|
||||||
let contradiction_detector = Arc::new(ContradictionHandler::default());
|
|
||||||
|
|
||||||
let pipeline = IngestPipeline::new(
|
|
||||||
entity_extractor,
|
|
||||||
fact_extractor,
|
|
||||||
contradiction_detector,
|
|
||||||
);
|
|
||||||
|
|
||||||
// Episode with problematic content (empty, or only whitespace)
|
|
||||||
let episode = Episode {
|
|
||||||
id: "test-empty".to_string(),
|
|
||||||
project_id: "test-project".to_string(),
|
|
||||||
text: "".to_string(),
|
|
||||||
wiki_links: vec![],
|
|
||||||
};
|
|
||||||
|
|
||||||
tracing::info!(
|
|
||||||
target: "test",
|
|
||||||
episode_id = %episode.id,
|
|
||||||
text_len = episode.text.len(),
|
|
||||||
"Processing empty episode"
|
|
||||||
);
|
|
||||||
|
|
||||||
match pipeline.ingest(&episode).await {
|
|
||||||
Ok(result) => {
|
|
||||||
tracing::info!(
|
|
||||||
target: "test",
|
|
||||||
entities = result.entities.len(),
|
|
||||||
edges = result.edges.len(),
|
|
||||||
"Empty episode processed (no error expected)"
|
|
||||||
);
|
|
||||||
println!("✓ Empty episode handled gracefully");
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
tracing::error!(
|
|
||||||
target: "test",
|
|
||||||
error = %e,
|
|
||||||
"Empty episode caused error"
|
|
||||||
);
|
|
||||||
// Empty is OK for some extractors
|
|
||||||
println!("⚠ Empty episode error (may be expected): {}", e);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_multiple_records_with_entity_verification() {
|
|
||||||
init_logging();
|
|
||||||
|
|
||||||
println!("\n[TEST] Processing multiple records with entity name verification\n");
|
|
||||||
|
|
||||||
let entity_extractor = Arc::new(WikiLinkFallbackExtractor);
|
|
||||||
let fact_extractor = Arc::new(SimpleFactExtractor);
|
|
||||||
let contradiction_detector = Arc::new(ContradictionHandler::default());
|
|
||||||
|
|
||||||
let pipeline = IngestPipeline::new(
|
|
||||||
entity_extractor,
|
|
||||||
fact_extractor,
|
|
||||||
contradiction_detector,
|
|
||||||
);
|
|
||||||
|
|
||||||
let records = vec![
|
|
||||||
("Kubernetes [[Docker]] is a container orchestrator", "wiki/k8s"),
|
|
||||||
("Docker [[Linux]] containers enable microservices", "wiki/docker"),
|
|
||||||
("", "wiki/empty"),
|
|
||||||
("Go [[Concurrency]] is powerful for backend services", "wiki/go"),
|
|
||||||
];
|
|
||||||
|
|
||||||
let mut success_count = 0;
|
|
||||||
let mut error_count = 0;
|
|
||||||
let mut all_extracted_entities = Vec::new();
|
|
||||||
|
|
||||||
for (idx, (text, source)) in records.iter().enumerate() {
|
|
||||||
let episode = Episode {
|
|
||||||
id: format!("record-{}", idx),
|
|
||||||
project_id: "test-project".to_string(),
|
|
||||||
text: text.to_string(),
|
|
||||||
wiki_links: vec![],
|
|
||||||
};
|
|
||||||
|
|
||||||
tracing::info!(
|
|
||||||
target: "test",
|
|
||||||
record_idx = idx,
|
|
||||||
source = source,
|
|
||||||
text_len = text.len(),
|
|
||||||
"Processing record"
|
|
||||||
);
|
|
||||||
|
|
||||||
match pipeline.ingest(&episode).await {
|
|
||||||
Ok(result) => {
|
|
||||||
tracing::debug!(
|
|
||||||
target: "test",
|
|
||||||
record_idx = idx,
|
|
||||||
entities = result.entities.len(),
|
|
||||||
edges = result.edges.len(),
|
|
||||||
"Record succeeded"
|
|
||||||
);
|
|
||||||
println!(" ✓ Record {}: {} entities, {} edges", idx, result.entities.len(), result.edges.len());
|
|
||||||
|
|
||||||
// Collect entity names for batch verification
|
|
||||||
for entity in &result.entities {
|
|
||||||
all_extracted_entities.push(entity.name.clone());
|
|
||||||
}
|
|
||||||
|
|
||||||
success_count += 1;
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
tracing::warn!(
|
|
||||||
target: "test",
|
|
||||||
record_idx = idx,
|
|
||||||
error = %e,
|
|
||||||
source = source,
|
|
||||||
"Record failed"
|
|
||||||
);
|
|
||||||
println!(" ✗ Record {}: {}", idx, e);
|
|
||||||
error_count += 1;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
println!("\nSummary: {} success, {} errors", success_count, error_count);
|
|
||||||
println!("All extracted entities: {:?}", all_extracted_entities);
|
|
||||||
|
|
||||||
tracing::info!(
|
|
||||||
target: "test",
|
|
||||||
total_records = records.len(),
|
|
||||||
success = success_count,
|
|
||||||
errors = error_count,
|
|
||||||
"Batch processing complete"
|
|
||||||
);
|
|
||||||
|
|
||||||
// ENHANCED: Verify that expected entities were extracted across records
|
|
||||||
assert!(success_count > 0, "At least some records should succeed");
|
|
||||||
assert!(
|
|
||||||
all_extracted_entities.iter().any(|name| name.contains("Docker") || name.contains("docker")),
|
|
||||||
"Docker entity should be extracted from at least one record"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
all_extracted_entities.iter().any(|name| name.contains("Linux") || name.contains("linux")),
|
|
||||||
"Linux entity should be extracted from at least one record"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
all_extracted_entities.iter().any(|name| name.contains("Concurrency") || name.contains("concurrency")),
|
|
||||||
"Concurrency entity should be extracted from at least one record (via [[Concurrency]] wiki link)"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_entity_deduplication() {
|
|
||||||
init_logging();
|
|
||||||
|
|
||||||
println!("\n[TEST] Entity deduplication (same entity from multiple records)\n");
|
|
||||||
|
|
||||||
let entity_extractor = Arc::new(WikiLinkFallbackExtractor);
|
|
||||||
let fact_extractor = Arc::new(SimpleFactExtractor);
|
|
||||||
let contradiction_detector = Arc::new(ContradictionHandler::default());
|
|
||||||
|
|
||||||
let pipeline = IngestPipeline::new(
|
|
||||||
entity_extractor,
|
|
||||||
fact_extractor,
|
|
||||||
contradiction_detector,
|
|
||||||
);
|
|
||||||
|
|
||||||
// Two records with overlapping entity references
|
|
||||||
let episode1 = Episode {
|
|
||||||
id: "record-1".to_string(),
|
|
||||||
project_id: "test-project".to_string(),
|
|
||||||
text: "Kubernetes uses [[Docker]] containers".to_string(),
|
|
||||||
wiki_links: vec!["Docker".to_string()],
|
|
||||||
};
|
|
||||||
|
|
||||||
let episode2 = Episode {
|
|
||||||
id: "record-2".to_string(),
|
|
||||||
project_id: "test-project".to_string(),
|
|
||||||
text: "Docker is used by [[Kubernetes]]".to_string(),
|
|
||||||
wiki_links: vec!["Kubernetes".to_string()],
|
|
||||||
};
|
|
||||||
|
|
||||||
let mut all_entities = Vec::new();
|
|
||||||
|
|
||||||
for episode in &[episode1, episode2] {
|
|
||||||
match pipeline.ingest(episode).await {
|
|
||||||
Ok(result) => {
|
|
||||||
all_entities.extend(result.entities);
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
tracing::error!(target: "test", error = %e, "Failed to ingest");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
println!("Total entities extracted: {}", all_entities.len());
|
|
||||||
for entity in &all_entities {
|
|
||||||
println!(" - {}", entity.name);
|
|
||||||
}
|
|
||||||
|
|
||||||
// Verify both Docker and Kubernetes were extracted
|
|
||||||
assert!(
|
|
||||||
all_entities.iter().any(|e| e.name.contains("Docker") || e.name.contains("docker")),
|
|
||||||
"Docker should be extracted"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
all_entities.iter().any(|e| e.name.contains("Kubernetes") || e.name.contains("kubernetes")),
|
|
||||||
"Kubernetes should be extracted"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Reference in New Issue
Block a user