Compare commits

..
Author SHA1 Message Date
poimenandrock fb61de6b47 feat: LLM entity + fact extraction pipeline (Zep paper alignment) (#48)
CI / CI (push) Successful in 12m9s
Deploy / Tag & Push Latest (push) Failing after 41s
DB Migration / Run Migrations (push) Failing after 18s
## Changes

### Entity Extraction
- Switch from WikiLinkFallbackExtractor to LlmEntityExtractor when LLM_ENDPOINT set
- `clean_llm_response()`: strips `<think>` tags, markdown fences, extracts JSON
- Handle array responses (Ollama returns `[...]` not `{entities: [...]}`)
- EntityType custom Deserialize: unknown variants → Unknown (no crash)
- Increase timeout 30s→90s, max_tokens 500→1500 for reasoning models
- Graceful reflection fallback: keep entities if verification fails

### Fact Extraction (NEW)
- LlmFactExtractor: LLM-based relationship extraction between entity pairs
- Validates source/target against known entity list (drops hallucinated edges)
- Same robust JSON cleaning for reasoning models + Ollama
- IngestWorker auto-selects LLM vs Simple based on LLM_ENDPOINT env

### K8s Deployment
- Add `command: ["/app/mem"]` (fix args replacing CMD)
- Add LLM_ENDPOINT, LLM_MODEL env vars for in-cluster LLM

## E2E Tested (local Ollama qwen2.5:3b)
- 12 entities extracted (person, tool, concept, organization)
- 5 edges with relationships and facts
- 781 tests pass

## Zep Paper Alignment (§2.2)
- Entity extraction + resolution (§2.2.1)
- Fact extraction between entity pairs (§2.2.2)
- Temporal edge invalidation ready (t_valid/t_invalid schema)
- Reflection verification (§2.2.1, graceful fallback)

---------

Co-authored-by: rock <[email protected]>
Reviewed-on: #48
Co-authored-by: poimen <[email protected]>
2026-09-11 01:11:15 +00:00
3 changed files with 48 additions and 77 deletions
+2 -4
View File
@@ -15,10 +15,8 @@ jobs:
name: Tag & Push Latest name: Tag & Push Latest
runs-on: rust runs-on: rust
steps: steps:
- name: Install Node.js and Docker - name: Install Docker
run: | run: apt-get update && apt-get install -y docker.io
apt-get update
apt-get install -y nodejs docker.io
- name: Checkout code - name: Checkout code
uses: actions/checkout@v4 uses: actions/checkout@v4
+43 -56
View File
@@ -11,79 +11,66 @@ env:
DB_HOST: memory-db-rw.poimen.svc.cluster.local DB_HOST: memory-db-rw.poimen.svc.cluster.local
DB_PORT: "5432" DB_PORT: "5432"
DB_NAME: memory DB_NAME: memory
MIGRATIONS_DIR: crates/mem-store/migrations
DOCKER_HOST: tcp://localhost:2375
jobs: jobs:
migrate: migrate:
name: Run Migrations name: Run Migrations
runs-on: rust runs-on: rust
steps: steps:
- name: Install Node.js, Docker, and psql - name: Install psql
run: | run: apt-get update && apt-get install -y postgresql-client
apt-get update
apt-get install -y nodejs docker.io postgresql-client
- name: Checkout code - name: Checkout code
uses: actions/checkout@v4 uses: actions/checkout@v4
with:
fetch-depth: 2
- name: Detect changed migrations - name: Fetch previous migrations state
id: detect
run: | run: |
CHANGED=$(git diff --name-only HEAD~1 HEAD -- "$MIGRATIONS_DIR"/*.sql 2>/dev/null || echo "") git fetch origin main --depth=2
if [ -n "$CHANGED" ]; then # List changed migration files
echo "files=$CHANGED" >> $GITHUB_OUTPUT CHANGED=$(git diff --name-only HEAD~1 HEAD -- crates/mem-store/migrations/ || echo "")
echo "found=true" >> $GITHUB_OUTPUT echo "Changed migrations: $CHANGED"
echo "Changed: $CHANGED" echo "CHANGED_MIGRATIONS=$CHANGED" >> $GITHUB_ENV
else
echo "found=false" >> $GITHUB_OUTPUT
echo "No migration changes detected"
fi
- name: Apply changed migrations (push) - name: Run migrations
if: github.event_name == 'push' && steps.detect.outputs.found == 'true' if: env.CHANGED_MIGRATIONS != ''
env:
PGHOST: ${{ env.DB_HOST }}
PGPORT: ${{ env.DB_PORT }}
PGDATABASE: ${{ env.DB_NAME }}
PGUSER: ${{ secrets.DB_USER }}
PGPASSWORD: ${{ secrets.DB_PASSWORD }}
run: | run: |
for f in ${{ steps.detect.outputs.files }}; do export PGPASSWORD="${DB_PASSWORD}"
[ -f "$f" ] || continue
echo "=== Applying: $f ===" echo "=== Running changed migrations ==="
psql -v ON_ERROR_STOP=1 -f "$f" for f in $CHANGED_MIGRATIONS; do
echo "=== OK ===" 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 done
- name: Apply all migrations (dispatch) 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 (manual trigger)
if: github.event_name == 'workflow_dispatch' if: github.event_name == 'workflow_dispatch'
env:
PGHOST: ${{ env.DB_HOST }}
PGPORT: ${{ env.DB_PORT }}
PGDATABASE: ${{ env.DB_NAME }}
PGUSER: ${{ secrets.DB_USER }}
PGPASSWORD: ${{ secrets.DB_PASSWORD }}
run: | run: |
for f in $(ls "$MIGRATIONS_DIR"/*.sql | sort); do export PGPASSWORD="${DB_PASSWORD}"
echo "=== Applying: $f ==="
psql -v ON_ERROR_STOP=1 -f "$f" || true echo "=== Running all migrations in order ==="
echo "=== Done ===" for f in $(ls crates/mem-store/migrations/*.sql | sort); do
echo "--- Applying: $f ---"
psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -f "$f" 2>&1 || true
echo "--- Done: $f ---"
done done
- name: Verify schema 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: env:
PGHOST: ${{ env.DB_HOST }} DB_USER: ${{ secrets.DB_USER }}
PGPORT: ${{ env.DB_PORT }} DB_PASSWORD: ${{ secrets.DB_PASSWORD }}
PGDATABASE: ${{ env.DB_NAME }}
PGUSER: ${{ secrets.DB_USER }}
PGPASSWORD: ${{ secrets.DB_PASSWORD }}
run: |
echo "=== Tables ==="
psql -c "\dt memory*"
echo "=== Entity Schema ==="
psql -c "\d memory_entity"
echo "=== Edge Schema ==="
psql -c "\d memory_edge"
+3 -17
View File
@@ -1342,29 +1342,15 @@ 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 question (fuzzy name/description search) // Step 1: Find entities (order by name for deterministic results)
let entities_rows: Vec<(String, String, String)> = sqlx::query_as( let entities_rows: Vec<(String, String, String)> = sqlx::query_as(
"SELECT id, 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 (name ILIKE '%' || $2 || '%' OR description ILIKE '%' || $2 || '%')
ORDER BY confidence DESC
LIMIT $3"
) )
.bind(&params.project) .bind(&params.project)
.bind(&params.question)
.bind(params.limit as i32) .bind(params.limit as i32)
.fetch_all(&state.pool) .fetch_all(&state.pool)
.await .await
.unwrap_or_default(); .unwrap_or_default();
tracing::info!(
target: "observability",
event = "query_entity_search",
project = %params.project,
question = %params.question,
matched = entities_rows.len(),
"Entity search complete"
);
// Step 2: Traverse edges from found entities // Step 2: Traverse edges from found entities
// NOTE: Edges will be empty until temporal schema is migrated // NOTE: Edges will be empty until temporal schema is migrated
@@ -1374,7 +1360,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, target_id, relation_type, fact, confidence, t_valid, t_invalid FROM memory_edge WHERE project_id = $1 AND source_id = $2" "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(&params.project) .bind(&params.project)
.bind(entity_id) .bind(entity_id)